diff --git a/.clang-format b/.clang-format index 66aeb8e..81253e8 100644 --- a/.clang-format +++ b/.clang-format @@ -10,6 +10,7 @@ MaxEmptyLinesToKeep: 2 SortIncludes: true CompactNamespaces: false ReflowComments: true +FixNamespaceComments: false DerivePointerAlignment: false PointerAlignment: Left diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index bd9bd1c..09f24e1 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -19,16 +19,16 @@ repos: hooks: - id: check-added-large-files args: ['--maxkb=512'] - exclude: '^(test_files/|src/tts/tokenizer/mandarin_data/)' + exclude: '^(test_files/|src/tts/tokenizer/mandarin_data/|native/tts-normalizer/testdata/en-US/(tokenize_and_classify|verbalize)\.far$)' - id: check-yaml - id: check-merge-conflict - id: end-of-file-fixer - exclude: '^(patches/|proto/riva-common/).*' + exclude: '^(patches/|proto/riva-common/|deploy/nemotron_native/kernels/|native/tts-normalizer/|python/nemotron_native_backend/|deploy/nemotron_native/(convert_model|serve)\.py$|scripts/nemotron_native/(build_kernels|build_package|runtime_dependencies)\.py$)' - id: trailing-whitespace - exclude: '^(patches/|proto/riva-common/).*' + exclude: '^(patches/|proto/riva-common/|deploy/nemotron_native/kernels/|native/tts-normalizer/|python/nemotron_native_backend/|deploy/nemotron_native/(convert_model|serve)\.py$|scripts/nemotron_native/(build_kernels|build_package|runtime_dependencies)\.py$)' - id: mixed-line-ending args: ['--fix=lf'] - exclude: '^(patches/|proto/riva-common/).*' + exclude: '^(patches/|proto/riva-common/|deploy/nemotron_native/kernels/|native/tts-normalizer/|python/nemotron_native_backend/|deploy/nemotron_native/(convert_model|serve)\.py$|scripts/nemotron_native/(build_kernels|build_package|runtime_dependencies)\.py$)' - repo: https://github.com/pre-commit/mirrors-clang-format rev: v18.1.8 @@ -36,20 +36,21 @@ repos: - id: clang-format types_or: [c++, c, cuda] # llama.cpp patches + vendored riva protos are upstream code — don't reformat. - exclude: '^(patches/|proto/riva-common/|build/|build-.*/).*' + exclude: '^(patches/|proto/riva-common/|build/|build-.*/|deploy/nemotron_native/kernels/|native/tts-normalizer/)' - repo: https://github.com/psf/black rev: 24.10.0 hooks: - id: black args: ['--skip-string-normalization', '--line-length=100'] - exclude: '^(proto/riva-common/|build/).*' + exclude: '^(proto/riva-common/|build/|python/nemotron_native_backend/|native/tts-normalizer/|deploy/nemotron_native/kernels/|deploy/nemotron_native/(convert_model|serve)\.py$|scripts/nemotron_native/(build_kernels|build_package|runtime_dependencies)\.py$)' - repo: https://github.com/pycqa/isort rev: 5.13.2 hooks: - id: isort name: isort (python) + exclude: '^(python/nemotron_native_backend/|native/tts-normalizer/|deploy/nemotron_native/kernels/|deploy/nemotron_native/(convert_model|serve)\.py$|scripts/nemotron_native/(build_kernels|build_package|runtime_dependencies)\.py$)' args: - '--profile=black' - '--line-length=100' diff --git a/README.md b/README.md index d343214..ba887e0 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,7 @@ | Speaker diarization | [Streaming Sortformer 4-speaker v2](https://huggingface.co/nvidia/diar_streaming_sortformer_4spk-v2) and [Nemotron 3 Diarization](https://huggingface.co/nvidia/Nemotron-3-Diarization), standalone or combined with ASR | | Text and speech translation | [Riva Translate 4B Instruct v2](https://huggingface.co/nvidia/Riva-Translate-4B-Instruct-v2), with composed ASR-to-NMT-to-TTS speech translation | | Speech synthesis | [MagpieTTS Multilingual 357M](https://huggingface.co/nvidia/magpie_tts_multilingual_357m) with [NeMo NanoCodec](https://huggingface.co/nvidia/nemo-nano-codec-22khz-1.89kbps-21.5fps) | +| Nemotron-TTS megakernels | Optional H100 backend with fixed cohorts or in-flight admission; [integration and benchmarks](docs/nemotron-megakernel.md) | | Full-duplex voicechat | [Nemotron Labs VoiceChat](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B), including realtime audio, transcripts, and tool calling | | Speech processing | [Silero VAD](https://github.com/snakers4/silero-vad), punctuation and capitalization, endpointing, text normalization, and subtitles | diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 6d71c60..a950f5d 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -119,7 +119,7 @@ include each selected dependency's license. ### Container system libraries -The Linux container copies the shared libraries required by its executables +The default Linux container copies the shared libraries required by its executables from the CUDA/Ubuntu builder into a `FROM scratch` runtime. This includes the GNU C Library, GCC runtime libraries, c-ares, zlib, bzip2, and XZ Utils/liblzma, subject to the exact feature set and package revisions in the built image. @@ -142,6 +142,51 @@ applicable agreement is installed with the archive under ## Other incorporated third-party code and data +### Imported Nemotron native deployment closure + +The retained deployment code in `deploy/nemotron_native/`, the Python controller +in `python/nemotron_native_backend/`, the native normalizer in +`native/tts-normalizer/`, and their imported build utilities retain their original +license terms and file-level notices. Importing these components does not +relicense them under this repository's Apache-2.0 license. Apache-2.0 headers on +new integration files do not change the terms of the imported code or its +dependencies. + +The legacy NVIDIA notice is preserved in +[`deploy/nemotron_native/LICENSE`](deploy/nemotron_native/LICENSE) and +[`native/tts-normalizer/LICENSE.legacy`](native/tts-normalizer/LICENSE.legacy). +The normalizer also preserves the Riva source notice in +[`native/tts-normalizer/LICENSE.riva`](native/tts-normalizer/LICENSE.riva). +The processor ABI and media-schema sources used by the native preparation +library retain the same legacy terms in +[`native/text_preparation/LICENSE.imported`](native/text_preparation/LICENSE.imported), +with their source revision recorded in +[`native/text_preparation/NOTICE`](native/text_preparation/NOTICE). The library +pins Hugging Face tokenizers 0.21.4 and its transitive dependencies through +`native/text_preparation/Cargo.lock`; their upstream licenses also apply. +These notices reserve NVIDIA's proprietary rights and require an express +license agreement for use, reproduction, disclosure, or distribution. Licensing +and redistribution review of the imported closure is pending; this notice and +a successful build do not grant publication or redistribution permission. + +The imported kernel dependencies include two ThunderKittens snapshots, +Copyright (c) 2024-2026 HazyResearch, under the MIT License. Their complete +notices remain in +[`thunderkittens/LICENSE`](deploy/nemotron_native/kernels/third_party/thunderkittens/LICENSE) +and +[`thunderkittens-v4-prefill/LICENSE`](deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/LICENSE). +The normalizer's pinned OpenFst, Sparrowhawk, Thrax, Abseil, protobuf, RE2, +and other dependencies are listed separately in +[`native/tts-normalizer/THIRD_PARTY_NOTICES.md`](native/tts-normalizer/THIRD_PARTY_NOTICES.md). +Their individual license terms remain applicable. + +The optional `docker/Dockerfile.nemotron` runtime uses the pinned NVIDIA PyTorch +container rather than the default scratch image described above. Its CUDA, +PyTorch, Python packages, and operating-system dependencies retain their own +licenses and notices. Model weights, tokenizer files, voice context, and +production normalization grammars are supplied separately; this repository does +not grant rights to those assets. + ### miniaudio - Source: [`mackron/miniaudio`](https://github.com/mackron/miniaudio), version diff --git a/app/bench_tts.cpp b/app/bench_tts.cpp index 8f85e12..e343a94 100644 --- a/app/bench_tts.cpp +++ b/app/bench_tts.cpp @@ -59,6 +59,10 @@ class TtsWorkload : public Workload { texts_.push_back(value()); } else if (arg == "--text-file") { text_files_.push_back(value()); + } else if (arg == "--tts-backend") { + parsed_.backend = value(); + } else if (arg == "--native-socket") { + parsed_.native_socket = value(); } else if (arg == "--magpie-model" || arg == "--model" || arg == "-m") { parsed_.runtime.magpie_model = value(); } else if (arg == "--codec-model") { diff --git a/app/serve.cpp b/app/serve.cpp index b066602..0ced0c9 100644 --- a/app/serve.cpp +++ b/app/serve.cpp @@ -102,7 +102,8 @@ validate_compiled_feature_options(int argc, char** argv) { add(option, "NEMO_SPEECH_BUILD_NMT=ON"); #endif #if !defined(NEMO_SPEECH_CLI_TTS) - if (option == "--tts-model" || option == "--codec-model" || option == "--tokenizer-dir" || + if (option == "--tts-backend" || option == "--native-socket" || option == "--tts-model" || + option == "--codec-model" || option == "--tokenizer-dir" || option == "--tn-model-dir" || option.rfind("--tts.", 0) == 0) add(option, "NEMO_SPEECH_BUILD_TTS=ON"); #endif @@ -304,7 +305,7 @@ run_server(int argc, char** argv) { server_config.address = value(i, arg); else if (arg == "--port") server_config.port = parse_int(value(i, arg), arg, 1, 65535); - else if (arg == "--threads") { + else if (arg == "--threads" || arg == "--http-threads") { server_config.threads = parse_int(value(i, arg), arg, 1, 1024); http_threads_set = true; } else if (arg == "--max-upload-mb") @@ -365,6 +366,10 @@ run_server(int argc, char** argv) { nmt_model = value(i, arg); #endif #if defined(NEMO_SPEECH_CLI_TTS) + else if (arg == "--tts-backend") + tts_config.backend = value(i, arg); + else if (arg == "--native-socket") + tts_config.native_socket = value(i, arg); else if (arg == "--tts-model") tts_model = value(i, arg); else if (arg == "--codec-model") @@ -482,15 +487,19 @@ run_server(int argc, char** argv) { } #endif #if defined(NEMO_SPEECH_CLI_TTS) + const bool native_tts = tts_enabled != 0 && tts_config.backend == "nemotron"; + if (tts_config.backend != "magpie" && tts_config.backend != "nemotron") + throw std::invalid_argument("tts.backend must be magpie or nemotron"); const bool tts_requested = tts_enabled > 0 || !tts_model.empty() || !tts_config.runtime.magpie_model.empty() || !codec_model.empty() || !tokenizer_model.empty() || !tn_model.empty(); - const auto magpie_path = tts_enabled == 0 ? std::string() : resolve("TTS model", [&] { - return optional_model( - tts_model.empty() ? tts_config.runtime.magpie_model : tts_model, "tts", - "MagpieTTS model", tts_requested); - }); - if (tts_enabled != 0 && (tts_requested || !magpie_path.empty())) { + const auto magpie_path = + (tts_enabled == 0 || native_tts) ? std::string() : resolve("TTS model", [&] { + return optional_model( + tts_model.empty() ? tts_config.runtime.magpie_model : tts_model, "tts", + "MagpieTTS model", tts_requested); + }); + if (!native_tts && tts_enabled != 0 && (tts_requested || !magpie_path.empty())) { tts_config.runtime.codec_model = resolve("TTS codec model", [&] { return optional_model( codec_model.empty() ? tts_config.runtime.codec_model : codec_model, "codec", @@ -579,6 +588,25 @@ run_server(int argc, char** argv) { } #endif #if defined(NEMO_SPEECH_CLI_TTS) + if (native_tts) { + nemo_speech::tts::SynthesizerConfig config; + config.backend = "nemotron"; + config.runtime = tts_config.runtime; + config.native_socket = tts_config.native_socket; + config.native_timeout_seconds = tts_config.native_timeout_seconds; + config.default_language_code = tts_config.default_language_code; + config.default_voice_name = tts_config.default_voice_name; + engines.load_tts(std::move(config)); + const int needed = engines.tts()->backend_batch_size() + 8; + if (!http_threads_set) + server_config.threads = std::max(server_config.threads, needed); + if (engines.tts()->backend_mode() == "fixed_cohort" && + server_config.threads < engines.tts()->backend_batch_size()) + throw std::invalid_argument( + "fixed cohorts require at least physical-width HTTP workers"); + if (server_config.preempt_tts) + throw std::invalid_argument("Nemotron batching does not support tts.preempt"); + } if (!magpie_path.empty()) { bool tts_cuda = false; #if defined(NEMO_SPEECH_CLI_CUDA) @@ -735,6 +763,8 @@ print_serve_help(const char* program) { " --diar-model MODEL Optional Sortformer path or indexed model\n" #endif #if defined(NEMO_SPEECH_CLI_TTS) + " --tts-backend NAME magpie (default) or nemotron\n" + " --native-socket PATH Nemotron controller bridge Unix socket\n" " --tts-model MODEL Optional MagpieTTS path or indexed model\n" " --codec-model MODEL NanoCodec path or indexed model\n" " --tokenizer-dir MODEL TTS tokenizer directory or indexed model\n" diff --git a/app/synthesize.cpp b/app/synthesize.cpp index b904cdd..9928f10 100644 --- a/app/synthesize.cpp +++ b/app/synthesize.cpp @@ -62,6 +62,17 @@ nemo_speech::tts::SynthesizerConfig make_synthesizer_config( nemo_speech::tts::MagpieTtsServerConfig parsed, const std::string& device_name, bool device_set) { + if (parsed.backend == "nemotron") { + nemo_speech::tts::SynthesizerConfig config; + config.backend = parsed.backend; + config.native_socket = parsed.native_socket; + config.native_timeout_seconds = parsed.native_timeout_seconds; + config.default_language_code = parsed.default_language_code; + config.default_voice_name = parsed.default_voice_name; + return config; + } + if (parsed.backend != "magpie") + throw std::invalid_argument("tts.backend must be magpie or nemotron"); parsed.runtime.magpie_model = resolve_model_file(parsed.runtime.magpie_model, "tts", "MagpieTTS model").string(); parsed.runtime.codec_model = @@ -110,6 +121,8 @@ print_synthesize_help(const char* program) { std::printf( "Usage: %s synthesize TEXT [options]\n\n" "Options:\n" + " --tts-backend NAME magpie (default) or nemotron\n" + " --native-socket PATH Nemotron controller bridge Unix socket\n" " --magpie-model MODEL MagpieTTS GGUF path or indexed HF repo\n" " (default: nvidia/magpie_tts_multilingual_357m)\n" " --codec-model MODEL NanoCodec GGUF path or indexed HF repo\n" @@ -173,6 +186,10 @@ command_synthesize(int argc, char** argv) { const std::string arg = argv[i]; if (arg == "--config") ++i; + else if (arg == "--tts-backend") + parsed.backend = value(i, arg); + else if (arg == "--native-socket") + parsed.native_socket = value(i, arg); else if (arg == "--magpie-model") parsed.runtime.magpie_model = value(i, arg); else if (arg == "--codec-model") diff --git a/deploy/nemotron_native/LICENSE b/deploy/nemotron_native/LICENSE new file mode 100644 index 0000000..009736b --- /dev/null +++ b/deploy/nemotron_native/LICENSE @@ -0,0 +1,7 @@ +Copyright (c) 2020-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. diff --git a/deploy/nemotron_native/convert_model.py b/deploy/nemotron_native/convert_model.py new file mode 100644 index 0000000..48156f9 --- /dev/null +++ b/deploy/nemotron_native/convert_model.py @@ -0,0 +1,129 @@ +import argparse +import hashlib +import json +import os +import shlex +import subprocess +import sys +from pathlib import Path + +LOCK = Path(__file__).resolve().parent / "model.lock.json" + + +def sha256(path): + digest = hashlib.sha256() + with open(path, "rb") as stream: + for block in iter(lambda: stream.read(16 * 1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def verify(root, rows, label): + for row in rows: + path = root / row["path"] + if not path.is_file(): + raise SystemExit(f"missing {label}: {path}") + if "bytes" in row and path.stat().st_size != row["bytes"]: + raise SystemExit(f"{label} size mismatch: {path}") + if "sha256" in row and sha256(path) != row["sha256"]: + raise SystemExit(f"{label} sha256 mismatch: {path}") + + +def download(lock, inputs): + source = lock["inputs"] + if not os.environ.get(source["token_environment_variable"]): + raise SystemExit( + f"set {source['token_environment_variable']} to download {source['repository']}" + ) + argv = ["hf", "download", source["repository"], "--repo-type", "model"] + argv += ["--revision", source["revision"], "--local-dir", str(inputs)] + argv += [row["path"] for row in source["files"]] + subprocess.run(argv, check=True) + + +def checkout(lock, converter): + spec = lock["converter"] + if not (converter / ".git").exists(): + subprocess.run(["git", "init", "--quiet", str(converter)], check=True) + subprocess.run( + ["git", "-C", str(converter), "remote", "add", "origin", spec["repository"]], check=True + ) + subprocess.run( + ["git", "-C", str(converter), "fetch", "--depth", "1", "origin", spec["commit"]], check=True + ) + subprocess.run( + ["git", "-C", str(converter), "checkout", "--quiet", "--detach", "FETCH_HEAD"], check=True + ) + head = subprocess.check_output( + ["git", "-C", str(converter), "rev-parse", "HEAD"], text=True + ).strip() + if head != spec["commit"]: + raise SystemExit("converter commit mismatch") + if sha256(converter / spec["script"]) != spec["script_sha256"]: + raise SystemExit("converter script sha256 mismatch") + + +def convert(lock, inputs, converter, output, gpu): + spec = lock["converter"] + output.mkdir(parents=True, exist_ok=True) + install = " && ".join( + f"python3 -m pip install --no-cache-dir {shlex.quote(r)}" + for r in spec["extra_pip_requirements"] + ) + command = " ".join(shlex.quote(a) for a in lock["argv"]) + script = f"{install} && {command}" if install else command + argv = ["docker", "run", "--rm", "--gpus", f"device={gpu}", "--network", "bridge"] + argv += ["--user", f"{os.getuid()}:{os.getgid()}", "--env", "HOME=/tmp"] + argv += [ + "--env", + "PYTHONPATH=/repo", + "--env", + "PYTHONDONTWRITEBYTECODE=1", + "--workdir", + "/repo", + ] + argv += ["--volume", f"{converter}:/repo:ro", "--volume", f"{inputs}:/ckpt:ro"] + argv += ["--volume", f"{output}:/out", "--entrypoint", "bash", spec["image"], "-c", script] + subprocess.run(argv, check=True) + + +def seal(lock, output, path): + verify(output, lock["outputs"], "model output") + files = {row["path"]: dict(sha256=row["sha256"], bytes=row["bytes"]) for row in lock["outputs"]} + value = dict(model_identity=lock["model_identity"], files=files) + path.write_text(json.dumps(value, indent=2) + "\n") + return path + + +def main(argv=None): + parser = argparse.ArgumentParser() + parser.add_argument("--inputs", type=Path) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--converter", type=Path) + parser.add_argument("--seal", type=Path) + parser.add_argument("--gpu", default="0") + parser.add_argument("--download", action="store_true") + parser.add_argument("--seal-only", action="store_true") + args = parser.parse_args(argv) + lock = json.loads(LOCK.read_text()) + output = args.output.resolve() + if not args.seal_only: + if args.inputs is None: + parser.error("--inputs is required unless --seal-only") + inputs = args.inputs.resolve() + if args.converter is None: + parser.error("--converter is required unless --seal-only") + if args.download: + download(lock, inputs) + verify(inputs, lock["inputs"]["files"], "input") + converter = args.converter.resolve() + checkout(lock, converter) + convert(lock, inputs, converter, output, args.gpu) + written = seal(lock, output, (args.seal or output / "model-seal.json").resolve()) + print( + json.dumps(dict(status="SEALED", model_identity=lock["model_identity"], seal=str(written))) + ) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/deploy/nemotron_native/kernels/b1024/prefill.cu b/deploy/nemotron_native/kernels/b1024/prefill.cu new file mode 100644 index 0000000..6d21d7e --- /dev/null +++ b/deploy/nemotron_native/kernels/b1024/prefill.cu @@ -0,0 +1 @@ +#include "regions/prefill/prefill.cu" diff --git a/deploy/nemotron_native/kernels/b1024/regions/prefill/native/prefill/dedup_native_prefill_impl.inc b/deploy/nemotron_native/kernels/b1024/regions/prefill/native/prefill/dedup_native_prefill_impl.inc new file mode 100644 index 0000000..fe4c125 --- /dev/null +++ b/deploy/nemotron_native/kernels/b1024/regions/prefill/native/prefill/dedup_native_prefill_impl.inc @@ -0,0 +1,267 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 16 || BATCH == 32 || BATCH == 64 || BATCH == 128 || BATCH == 256 || BATCH == 512 || BATCH == 1024 || BATCH == 2048)) + throw std::runtime_error("native prefill requires a matching supported physical batch"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((std::uint64_t(c.prompt_positions) + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (rows > (INT_MAX - 132 * THREADS) / D_MODEL) + throw std::runtime_error("native prefill full-width prompt exceeds safe device loop bounds"); + if (padded > (INT_MAX - 132 * THREADS) / (BB_EXPERTS * BB_MOE_INTER)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((size_t(config->prompt_positions) + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + float *r9_workspace = nullptr; + mk2_check(owned.alloc(&r9_workspace,size_t(BATCH)*ssm_elements*sizeof(float))); + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + r9_state_t *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(r9_state_t))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(r9_state_t))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = r9_workspace; + st->mb_host[i].r9_private = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + mk2_check(owned.alloc(&st->mb_host[i].r8_canonical_ssm, ssm_elements * sizeof(float))); + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_dtda, 2 * MB_HEADS); + { + const size_t estride = padded; + const size_t grows = size_t(BB_EXPERTS) * estride; + mk2_check(owned.alloc(&hb.pf_moe_gath, grows * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_gath, 0, grows * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_hc, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_hc, 0, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_gidx, grows * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_gidx, 0, grows * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ecnt, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ecnt, 0, BB_EXPERTS * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ntiles, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ntiles, 0, BB_EXPERTS * sizeof(int))); + hb.pf_moe_estride = int(estride); + hb.pf_moe_gath_v = act_gl{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_v = act_gl{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + hb.pf_moe_gath_vn = act_gl_n{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_vn = act_gl_n{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + scratch_bytes += grows * (D_MODEL + BB_MOE_INTER) * sizeof(bf16); + } + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(owned.alloc(&hb.native_pf_cohort_check, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(hb.native_pf_cohort_check, 0, 2 * sizeof(unsigned int))); +#if PF_ACTIVE_TOKEN_EXPERTS + mk2_check(owned.alloc(&hb.native_pf_expert_masks, (padded/BATCH_PAD)*sizeof(unsigned int))); +#endif + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} +extern "C" int mk2_native_prefill_validate(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null compact prefill handle"); + auto *st = reinterpret_cast(p->base); + mk2_check(cudaStreamSynchronize(static_cast(stream))); + unsigned int status[2] = {0, 0}; + mk2_check(cudaMemcpy(status, st->host_bb.native_pf_cohort_check, + sizeof(status), cudaMemcpyDeviceToHost)); + if (status[0] != 0 || status[1] != 1) + throw std::runtime_error("compact prefill requires bit-identical prompts and identity slots"); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown compact prefill validation failure"; return 1; } +} +extern "C" int mk2_native_prefill_workspace_sequences() { return 1; } +extern "C" int mk2_native_prefill_effective_sequences(MK2NativePrefill *p, int *out) { + try { + if (!p || !out) throw std::runtime_error("null effective-sequence query"); + auto *st = reinterpret_cast(p->base); + mk2_check(cudaMemcpy(out, st->host_bb.native_pf_cohort_check + 1, + sizeof(int), cudaMemcpyDeviceToHost)); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" int mk2_r8_initial_state(MK2NativePrefill *p, void **initial, void **ready) { + if (!p || !initial || !ready) return 1; + auto *st = reinterpret_cast(p->base); + for (int i = 0; i < 14; ++i) initial[i] = st->mb_host[i].r8_canonical_ssm; + *ready = st->host_bb.native_pf_cohort_check + 1; + return 0; +} +extern "C" int mk2_r9_prefill_bind(MK2NativePrefill*p,const void*const*scales,const void*const*maps){if(!p)return 1;auto*st=(MK2Stage0State*)p->base;for(int i=0;i<14;++i){st->mb_host[i].r9_scales=(const float*)scales[i];st->mb_host[i].r9_map=(const int*)maps[i];}return (int)cudaMemcpy((void*)st->host_bb.mamba,st->mb_host.data(),14*sizeof(MambaLayer),cudaMemcpyHostToDevice);} diff --git a/deploy/nemotron_native/kernels/b1024/regions/prefill/native/prefill/dedup_predictor_mk2.cuh b/deploy/nemotron_native/kernels/b1024/regions/prefill/native/prefill/dedup_predictor_mk2.cuh new file mode 100644 index 0000000..f98dd53 --- /dev/null +++ b/deploy/nemotron_native/kernels/b1024/regions/prefill/native/prefill/dedup_predictor_mk2.cuh @@ -0,0 +1,4600 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef PF_ACTIVE_TOKEN_EXPERTS +#define PF_ACTIVE_TOKEN_EXPERTS 0 +#endif +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +#ifndef MK2_PREFILL_ONLY +#define MK2_PREFILL_ONLY 0 +#endif +#ifndef MK2_BATCH_TILE +#define MK2_BATCH_TILE (((MK2_BATCH + 15) / 16) * 16) +#endif +constexpr int BATCH_PAD = MK2_BATCH_TILE; +static_assert(BATCH >= 1 && (MK2_PREFILL_ONLY || BATCH <= BATCH_PAD)); +static_assert(BATCH_PAD == 16 || BATCH_PAD == 32 || BATCH_PAD == 64 || BATCH_PAD == 128, + "native prefill tile must be 16, 32, 64 or 128"); +static_assert(MK2_PREFILL_ONLY || BATCH_PAD == ((BATCH + 15) / 16) * 16, + "a wider token tile requires a prefill-only library"); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_PW_MULTI +#define PF_SSM_PW_MULTI 4 +#endif +#ifndef PF_SSM_PW_SINGLE +#define PF_SSM_PW_SINGLE 2 +#endif +static_assert(MB_HEAD_DIM % PF_SSM_PW_MULTI == 0 && + MB_HEAD_DIM % PF_SSM_PW_SINGLE == 0, + "head_dim must split evenly into scan channel groups"); +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef MK2_SMEM_SLACK +#define MK2_SMEM_SLACK 1024 +#endif +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +#ifndef MK2_THREADS +#define MK2_THREADS 256 +#endif +constexpr int THREADS = MK2_THREADS; +constexpr int NORM_THREADS = 256; +constexpr int NORM_WARPS = NORM_THREADS / 32; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +#ifndef MK2_GEMM_CONSUMERS +#define MK2_GEMM_CONSUMERS 1 +#endif +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 +#define MK2_HEAD_STAGE_BYTES 0 +#else +#define MK2_HEAD_STAGE_BYTES sizeof(out_tile_f) +#endif +constexpr int NCONS = MK2_GEMM_CONSUMERS; +static_assert(NCONS >= 1 && NCONS <= 3, "consumers occupy warpgroups 0, 2 and 3"); +static_assert(THREADS >= (NCONS + 1) * 128, "one producer plus NCONS consumers"); +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +constexpr int PAD_N = 32; +using act_tile_n = st_bf; +using out_tile_n = st_bf; +using act_gl_n = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_MOE_ACC_C, EPI_ACC_F32 }; +#include +constexpr int R9_STATE_MODE=1; +using r9_state_t=bf16; +__device__ __forceinline__ int r9_scale_index(int head,int channel,int n) {return (((head/8)*4+(head%8)/2)*4+channel/6)*128+n;} +__device__ __forceinline__ float r9_load(r9_state_t v,float scale) { +return __bfloat162float(v); +} +__device__ __forceinline__ r9_state_t r9_store(float x,float scale) { +return __float2bfloat16(x); +} +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; + float *r8_canonical_ssm = nullptr; + r9_state_t *r9_private = nullptr; + const float *r9_scales = nullptr; + const int *r9_map = nullptr; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + unsigned int *native_pf_expert_masks = nullptr; + float *pf_dtda = nullptr; + bf16 *pf_moe_gath = nullptr; + bf16 *pf_moe_hc = nullptr; + int *pf_moe_gidx = nullptr; + int *pf_moe_ecnt = nullptr; + int *pf_moe_ntiles = nullptr; + int pf_moe_estride = 0; + act_gl pf_moe_gath_v; + act_gl pf_moe_hc_v; + act_gl_n pf_moe_gath_vn; + act_gl_n pf_moe_hc_vn; + unsigned int *native_pf_cohort_check = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + volatile unsigned int *generation = state + 1; + const unsigned int seen = *generation; + if (atomicAdd(state, 1u) == gridDim.x - 1) { + atomicExch(state, 0u); + atomicAdd(state + 1, 1u); + } else { + while (*generation == seen) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const bf16 *in = src + (size_t)row * D_MODEL; + float v[PER_THREAD]; + float sum = 0.f, sq = 0.f; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (threadIdx.x < NORM_THREADS) { + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + v[j] = __bfloat162float(in[threadIdx.x + j * NORM_THREADS]); + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + if (lane == 0) { scratch[warp] = sum; scratch[NORM_WARPS + warp] = sq; } + } + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f, q = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) { s += scratch[w]; q += scratch[NORM_WARPS + w]; } + scratch[2 * NORM_WARPS] = s / D_MODEL; + scratch[2 * NORM_WARPS + 1] = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + } + __syncthreads(); + const float mean = scratch[2 * NORM_WARPS], inv = scratch[2 * NORM_WARPS + 1]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + if (threadIdx.x >= NORM_THREADS) break; + const int i = threadIdx.x + j * NORM_THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * weight[i]); + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, + int idx, unsigned mask); +__device__ __forceinline__ void norm_sub_sync(int sub) { + if (THREADS == NORM_THREADS) { __syncthreads(); return; } + asm volatile("bar.sync %0, %1;" :: "r"(6 + sub), "n"(NORM_THREADS) : "memory"); +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, const float *fold_src = nullptr, + const float *fold_routes = nullptr) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + constexpr int NSUB = THREADS / NORM_THREADS; + if (threadIdx.x >= NSUB * NORM_THREADS) return; + const int sub = threadIdx.x / NORM_THREADS; + const int tid = threadIdx.x - sub * NORM_THREADS; + float *const sc = scratch + sub * (NORM_WARPS + 1); + for (int row = blockIdx.x * NSUB + sub; row < n_rows; + row += gridDim.x * NSUB) { + const size_t base = (size_t)row * D_MODEL; + unsigned fold_mask = (1u << BB_EXPERTS) - 1u; + if (fold_src && fold_routes) { + fold_mask = 0u; + const float *rw = fold_routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) + if (rw[e] != 0.f) fold_mask |= 1u << e; + } + float v[PER_THREAD]; + float sq = 0.f; + const int warp = tid >> 5, lane = tid & 31; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + bf16 hv; + if (fold_src) { + hv = bb_moe_fold_deterministic(fold_src, (int)(base + i), fold_mask); + h[base + i] = hv; + } else { + hv = h[base + i]; + } + v[j] = __bfloat162float(hv) + __bfloat162float(r[base + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + if (lane == 0) sc[warp] = sq; + norm_sub_sync(sub); + if (tid == 0) { + float s = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) s += sc[w]; + sc[NORM_WARPS] = rsqrtf(s / D_MODEL + RMS_EPS); + } + norm_sub_sync(sub); + const float inv = sc[NORM_WARPS]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + out[base + i] = __float2bfloat16( + __bfloat162float(weight[i]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + } + norm_sub_sync(sub); + } +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + const int c = idx - row * MB_CONV_DIM; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + const float raw = __bfloat162float( + zxbcdt[(size_t)row * MB_IN_PROJ + MB_INNER + c]); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + xbc_act[(size_t)row * MB_CONV_DIM + c] = __float2bfloat16(acc); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + const int pairs = BATCH * MB_HEADS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_HEADS * split); + const int head = (pair / split) % MB_HEADS; + const int first_channel = (pair % split) * channels; + const int group = head / MB_HEADS_PER_GROUP; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + L.dt_bias[head]); + const float dA = __expf(L.A[head] * dt); + const float Dh = L.D[head]; + float *st = L.ssm_state + (size_t)slot_map[row] * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + for (int p = first_channel; p < first_channel + channels; ++p) { + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + const size_t at = (size_t)p * MB_STATE + n; + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + static_assert(MB_CONV_K == 4, "packed conv taps assume a 4-tap kernel"); + constexpr int CPT = 4; + static_assert(MB_CONV_DIM % CPT == 0, "channel vector width"); + constexpr int CGROUPS = MB_CONV_DIM / CPT; + const int total = n_tokens * CGROUPS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / CGROUPS; + const int c = (idx - t * CGROUPS) * CPT; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + float wv[CPT][MB_CONV_K]; + { + const float4 *wsrc = reinterpret_cast( + L.conv_w + (size_t)c * MB_CONV_K); + #pragma unroll + for (int j = 0; j < CPT / 2; ++j) { + const float4 w4 = wsrc[j]; + const unsigned uu0 = __float_as_uint(w4.x); + const unsigned uu1 = __float_as_uint(w4.y); + const unsigned uu2 = __float_as_uint(w4.z); + const unsigned uu3 = __float_as_uint(w4.w); + wv[2 * j][0] = __uint_as_float(uu0 << 16); + wv[2 * j][1] = __uint_as_float(uu0 & 0xffff0000u); + wv[2 * j][2] = __uint_as_float(uu1 << 16); + wv[2 * j][3] = __uint_as_float(uu1 & 0xffff0000u); + wv[2 * j + 1][0] = __uint_as_float(uu2 << 16); + wv[2 * j + 1][1] = __uint_as_float(uu2 & 0xffff0000u); + wv[2 * j + 1][2] = __uint_as_float(uu3 << 16); + wv[2 * j + 1][3] = __uint_as_float(uu3 & 0xffff0000u); + } + } + float acc[CPT]; + { + const uint2 b2 = *reinterpret_cast(L.conv_b + c); + acc[0] = __uint_as_float(b2.x << 16); + acc[1] = __uint_as_float(b2.x & 0xffff0000u); + acc[2] = __uint_as_float(b2.y << 16); + acc[3] = __uint_as_float(b2.y & 0xffff0000u); + } + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v[CPT]; + if (src < 0) { + #pragma unroll + for (int q = 0; q < CPT; ++q) + v[q] = __bfloat162float( + state[(size_t)q * (MB_CONV_K - 1) + pos + k]); + } else { + const uint2 p = *reinterpret_cast( + zxbcdt + (size_t)src * MB_IN_PROJ + MB_INNER + c); + v[0] = __uint_as_float(p.x << 16); + v[1] = __uint_as_float(p.x & 0xffff0000u); + v[2] = __uint_as_float(p.y << 16); + v[3] = __uint_as_float(p.y & 0xffff0000u); + } + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] += wv[q][k] * v[q]; + } + uint2 out; + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] = acc[q] / (1.f + __expf(-acc[q])); + out.x = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[0])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[1])) << 16); + out.y = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[2])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[3])) << 16); + *reinterpret_cast(xbc_act + (size_t)t * MB_CONV_DIM + c) = out; + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_dtda_prefill( + const MambaLayer &L, const bf16 *zxbcdt, float *dtda, int n_tokens) { + const int total = n_tokens * MB_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_HEADS; + const int head = idx - t * MB_HEADS; + const float dt = mamba_softplus( + __bfloat162float(zxbcdt[(size_t)t * MB_IN_PROJ + MB_INNER + + MB_CONV_DIM + head]) + L.dt_bias[head]); + dtda[(size_t)idx * 2] = dt; + dtda[(size_t)idx * 2 + 1] = __expf(L.A[head] * dt); + } +} +template +__device__ __forceinline__ void mamba_ssm_prefill_impl( + const MambaLayer &L, const bf16 *xbc_act, const float *dtda, bf16 *gated, + int n_pos, int n_seq, int slot, const int *seq_slot) { + constexpr int NS = MB_STATE / 32; + constexpr int PG = MB_HEAD_DIM / PW; + constexpr bool RS = (PW % 4) == 0; + constexpr int H1 = RS ? PW / 2 : 1; + constexpr int H2 = RS ? PW / 4 : 1; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const bool keep_hi = (lane & 16) != 0; + const bool keep_hi2 = (lane & 8) != 0; + const int warps_total = gridDim.x * WARPS; + const int units = n_seq * MB_HEADS * PG; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int pg = unit % PG; + const int rest = unit / PG; + const int head = rest % MB_HEADS; + const int seq = rest / MB_HEADS; + const int p0 = pg * PW; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Dh = L.D[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p0 * MB_STATE; + float sreg[PW][NS]; + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + sreg[q][j] = st[(size_t)q * MB_STATE + lane + j * 32]; + } + auto fetch = [&](int p, float (&oB)[NS], float (&oC)[NS], + float (&oX)[PW], float2 &od) { + const size_t tt = (size_t)p * n_seq + seq; + const bf16 *xa2 = xbc_act + tt * MB_CONV_DIM; + const bf16 *Bs = xa2 + MB_INNER + group * MB_STATE; + const bf16 *Cs = xa2 + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + od = *reinterpret_cast( + dtda + (tt * MB_HEADS + head) * 2); + #pragma unroll + for (int j = 0; j < NS; ++j) { + oB[j] = __bfloat162float(Bs[lane + j * 32]); + oC[j] = __bfloat162float(Cs[lane + j * 32]); + } + if constexpr (RS) { + const uint2 *xp = reinterpret_cast( + xa2 + head * MB_HEAD_DIM + p0); + #pragma unroll + for (int q = 0; q < PW; q += 4) { + const uint2 w = xp[q >> 2]; + oX[q] = __uint_as_float(w.x << 16); + oX[q + 1] = __uint_as_float(w.x & 0xffff0000u); + oX[q + 2] = __uint_as_float(w.y << 16); + oX[q + 3] = __uint_as_float(w.y & 0xffff0000u); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) + oX[q] = __bfloat162float(xa2[head * MB_HEAD_DIM + p0 + q]); + } + }; + float pB[NS], pC[NS], pX[PW]; + float2 pd; + fetch(0, pB, pC, pX, pd); + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const float dt = pd.x; + const float dA = pd.y; + float Bv[NS], Cv[NS]; + #pragma unroll + for (int j = 0; j < NS; ++j) { + Bv[j] = pB[j] * dt; + Cv[j] = pC[j]; + } + float acc[PW], xvs[PW]; + #pragma unroll + for (int q = 0; q < PW; ++q) xvs[q] = pX[q]; + if (pos + 1 < n_pos) fetch(pos + 1, pB, pC, pX, pd); + #pragma unroll + for (int q = 0; q < PW; ++q) { + const float xv = xvs[q]; + float a = 0.f; + #pragma unroll + for (int j = 0; j < NS; ++j) { + sreg[q][j] = sreg[q][j] * dA + Bv[j] * xv; + a += sreg[q][j] * Cv[j]; + } + acc[q] = a; + } + if constexpr (RS) { + float m1[H1], t1[H1], x1[H1]; + #pragma unroll + for (int q = 0; q < H1; ++q) { + m1[q] = keep_hi ? acc[q + H1] : acc[q]; + t1[q] = keep_hi ? acc[q] : acc[q + H1]; + x1[q] = keep_hi ? xvs[q + H1] : xvs[q]; + } + #pragma unroll + for (int q = 0; q < H1; ++q) + m1[q] += __shfl_xor_sync(0xffffffff, t1[q], 16); + float m2[H2], t2[H2], x2[H2]; + #pragma unroll + for (int q = 0; q < H2; ++q) { + m2[q] = keep_hi2 ? m1[q + H2] : m1[q]; + t2[q] = keep_hi2 ? m1[q] : m1[q + H2]; + x2[q] = keep_hi2 ? x1[q + H2] : x1[q]; + } + #pragma unroll + for (int q = 0; q < H2; ++q) + m2[q] += __shfl_xor_sync(0xffffffff, t2[q], 8); + #pragma unroll + for (int q = 0; q < H2; ++q) { + #pragma unroll + for (int off = 4; off > 0; off >>= 1) + m2[q] += __shfl_xor_sync(0xffffffff, m2[q], off); + } + if ((lane & 7) == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + + p0 + (keep_hi ? H1 : 0) + (keep_hi2 ? H2 : 0); + #pragma unroll + for (int q = 0; q < H2; ++q) + dst[q] = __float2bfloat16(m2[q] + x2[q] * Dh); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc[q] += __shfl_xor_sync(0xffffffff, acc[q], off); + } + if (lane == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + p0; + #pragma unroll + for (int q = 0; q < PW; ++q) + dst[q] = __float2bfloat16(acc[q] + xvs[q] * Dh); + } + } + } + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + st[(size_t)q * MB_STATE + lane + j * 32] = sreg[q][j]; + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, + const float *dtda, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + (void)zxbcdt; + if (n_seq > 1) + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens / n_seq, n_seq, slot, seq_slot); + else + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens, 1, slot, nullptr); +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16( + held[j] * inv * __bfloat162float(L.gnorm_w[idx])); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = BATCH * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / BB_ATT_HEADS; + const int head = pair % BB_ATT_HEADS; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const int kv = head / kv_per_q; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *q = base + head * BB_HEAD_DIM; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (head % kv_per_q == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + for (int tpos = 0; tpos < T; ++tpos) { +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new + : kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#endif + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m, s); + const float rescale = __expf(m - m_new); + const float w = __expf(s - m_new); + denom = denom * rescale + w; +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *vt = kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#else + const bf16 *vt = cur_token + ? v_new + : kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + w * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + constexpr int CHUNKS = BB_HEAD_DIM / 8; + const int total = n_tokens * BB_KV_HEADS * CHUNKS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int c = idx % CHUNKS; + const int rest = idx / CHUNKS; + const int t = rest / BB_KV_HEADS; + const int kv = rest - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM + c * 8; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM + c * 8; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv) + c * 8; + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv) + c * 8; + *reinterpret_cast(kc) = *reinterpret_cast(k_new); + *reinterpret_cast(vc) = *reinterpret_cast(v_new); + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float qv[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qv[j] = __bfloat162float(q[lane + j * 32]); + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + const int *const krow = + L.block_table + ((size_t)seq * 2 + 0) * L.max_blocks_per_seq; + const int *const vrow = + L.block_table + ((size_t)seq * 2 + 1) * L.max_blocks_per_seq; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + const int head_bias = kv * L.tokens_per_block * BB_HEAD_DIM; + int page = -1; + const bf16 *kpage = nullptr, *vpage = nullptr; + float kn[PER_LANE], vn[PER_LANE]; + { + page = 0; + kpage = L.pool_base + (size_t)(krow[0] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + (size_t)(vrow[0] & 0x7FFFFFFF) * elems_per_block; + const int at = head_bias + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + for (int s = 0; s <= pos; ++s) { + float kc[PER_LANE], vc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { kc[j] = kn[j]; vc[j] = vn[j]; } + if (s < pos) { + const int np = (s + 1) / L.tokens_per_block; + if (np != page) { + page = np; + kpage = L.pool_base + + (size_t)(krow[np] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + + (size_t)(vrow[np] & 0x7FFFFFFF) * elems_per_block; + } + const int at = head_bias + + ((s + 1) % L.tokens_per_block) * BB_HEAD_DIM + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) dot += qv[j] * kc[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[j] = acc[j] * rescale + wt * vc[j]; + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH, unsigned int *tile_masks = nullptr, + bf16 *gath = nullptr, int *gidx = nullptr, int *ecnt = nullptr, + int estride = 0, bf16 *row_stage = nullptr) { + __shared__ float score[BB_EXPERTS]; + __shared__ int s_pick[BB_TOP_K], s_slot[BB_TOP_K]; + constexpr int CH = D_MODEL / 8; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + { + const float4 *gsrc = + reinterpret_cast(normed + (size_t)row * D_MODEL); + float4 *gdst = reinterpret_cast(row_stage); + for (int i = threadIdx.x; i < CH; i += THREADS) gdst[i] = gsrc[i]; + } + if ((int)threadIdx.x < BB_EXPERTS) + moe_w[(size_t)row * BB_EXPERTS + threadIdx.x] = 0.f; + __syncthreads(); + for (int e = warp; e < BB_EXPERTS; e += WARPS) { + const bf16 *wr = gate_w + (size_t)e * D_MODEL; + float acc = 0.f; + for (int i = lane; i < D_MODEL; i += 32) + acc += __bfloat162float(row_stage[i]) * + __bfloat162float(wr[i]); + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) score[e] = 1.f / (1.f + __expf(-acc)); + } + __syncthreads(); + if (threadIdx.x == 0) { + float remaining[BB_EXPERTS]; + int pick[BB_TOP_K]; + for (int e = 0; e < BB_EXPERTS; ++e) + remaining[e] = score[e] + __bfloat162float(gate_bias[e]); + float total = 0.f; + for (int k = 0; k < BB_TOP_K; ++k) { + int best = 0; + for (int e = 1; e < BB_EXPERTS; ++e) + if (remaining[e] > remaining[best]) best = e; + pick[k] = best; + total += score[best]; + remaining[best] = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + unsigned int bits = 0u; + for (int k = 0; k < BB_TOP_K; ++k) { + moe_w[row * BB_EXPERTS + pick[k]] = score[pick[k]] * norm; + bits |= 1u << pick[k]; + s_pick[k] = pick[k]; + } + if (tile_masks) atomicOr(&tile_masks[row / BATCH_PAD], bits); + } + __syncthreads(); + if (gath) { + if ((int)threadIdx.x < BB_TOP_K) { + const int k = threadIdx.x; + const int slot = atomicAdd(&ecnt[s_pick[k]], 1); + s_slot[k] = slot; + gidx[(size_t)s_pick[k] * estride + slot] = row; + } + __syncthreads(); + const float4 *src = reinterpret_cast(row_stage); + for (int k = 0; k < BB_TOP_K; ++k) { + float4 *dst = reinterpret_cast( + gath + ((size_t)s_pick[k] * estride + s_slot[k]) * D_MODEL); + for (int i = threadIdx.x; i < CH; i += THREADS) dst[i] = src[i]; + } + } + __syncthreads(); + } +} +__device__ __forceinline__ unsigned moe_row_mask(const float *routes, int row) { + unsigned m = 0u; + const float *rw = routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) if (rw[e] != 0.f) m |= 1u << e; + return m; +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx, + unsigned mask) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + for (unsigned m = mask; m; m &= m - 1u) + value = __fadd_rn(value, parts[base + (size_t)(__ffs(m) - 1) * D_MODEL]); + return __float2bfloat16( + __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL])); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + uint64_t *keys, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int n = threadIdx.x; n < SORT_N; n += THREADS) { + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + keys[n] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + __syncthreads(); + for (int k = 2; k <= SORT_N; k <<= 1) { + if (k >= 64) __syncthreads(); + for (int j = k >> 1; j > 0; j >>= 1) { + for (int i = threadIdx.x; i < SORT_N; i += THREADS) { + const int partner = i ^ j; + if (partner > i) { + const bool descending = ((i & k) == 0); + if ((keys[i] < keys[partner]) == descending) { + const uint64_t t = keys[i]; + keys[i] = keys[partner]; + keys[partner] = t; + } + } + } + if (j >= 32) __syncthreads(); + else __syncwarp(); + } + } + __syncthreads(); + if (threadIdx.x == 0) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float v; + int n; + if (!eos_used && eos_v > av) { + v = eos_v; + n = AUDIO_EOS; + eos_used = true; + } else { + v = av; + n = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = + v + gumbel[((size_t)row * N_POS + pos) * TOP_K + j]; + if (noised > best_noised) { best_noised = noised; best_code = n; } + } + out_codes[row] = best_code; + } + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void prefill_expert_masks(const float *routes, unsigned int *masks, int rows) { + const int lane=threadIdx.x&31, warp=threadIdx.x>>5; + const int nt=(rows+BATCH_PAD-1)/BATCH_PAD; + for (int t=blockIdx.x*WARPS+warp;t(gath + g * D_MODEL + c * 8); + if (local < ecnt[e]) + *dst = *reinterpret_cast( + normed + (size_t)gidx[g] * D_MODEL + c * 8); + else + *dst = make_float4(0.f, 0.f, 0.f, 0.f); + } +} +__device__ __forceinline__ int gemm_eff(int tiles, int reps) { + if (NCONS <= 1) return 1; + const int g = (int)gridDim.x; + const int rounds1 = (tiles * reps + g - 1) / g; + if constexpr (NCONS == 3) { + const int max_eff = reps <= BB_EXPERTS ? 2 : 3; + const int rounds2 = (((tiles + 1) / 2) * reps + g - 1) / g; + if (max_eff == 2) return rounds2 < rounds1 ? 2 : 1; + const int rounds3 = (((tiles + 2) / 3) * reps + g - 1) / g; + if (rounds3 < rounds2 && rounds3 < rounds1) return 3; + return rounds2 < rounds1 ? 2 : 1; + } else { + const int roundsN = (((tiles + NCONS - 1) / NCONS) * reps + g - 1) / g; + return roundsN < rounds1 ? NCONS : 1; + } +} +__device__ __forceinline__ int gemm_groups(int tiles, int reps) { + const int eff = gemm_eff(tiles, reps); + return (tiles + eff - 1) / eff; +} +__device__ __forceinline__ int gemm_consumer_id() { + const int g = warpgroup::groupid(); + if (g == 1) return -1; + const int c = g == 0 ? 0 : g - 1; + return c < NCONS ? c : -1; +} +__device__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][BATCH_PAD]; + __shared__ int moe_gidx_stage[NCONS][BATCH_PAD]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + a_s[s], a, {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + a_s[s], a, {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / BATCH_PAD) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, BATCH_PAD> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], a_s[s]); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * BATCH_PAD; + if (epi == EPI_BIAS) { + warpgroup::store(o_sf, accum); + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + dest[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(o_sf[int2{col, row}] + + __bfloat162float(bias[base + col])); + } + warpgroup::sync(1 + cons); + continue; + } + warpgroup::store(o_s[cons], accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_s[cons][int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__device__ void gemm_narrow( + const weight_gl &w, const act_gl_n &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = PAD_N, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][PAD_N]; + __shared__ int moe_gidx_stage[NCONS][PAD_N]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + PAD_N - 1) / PAD_N; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + PAD_N - 1) / PAD_N; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / PAD_N) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, PAD_N> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], + *reinterpret_cast(&a_s[s])); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * PAD_N; + out_tile_n &o_sn = *reinterpret_cast(&o_s[cons]); + warpgroup::store(o_sn, accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_sn[int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&w_s)[STAGES * NCONS] = + alloc.allocate(); + act_tile (&a_s)[STAGES] = alloc.allocate(); + out_tile (&o_s)[NCONS] = alloc.allocate(); +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 + static_assert(NCONS * sizeof(out_tile) >= sizeof(out_tile_f), + "aliased head staging must fit the output tiles"); + out_tile_f &o_sf = *reinterpret_cast(&o_s[0]); +#else + out_tile_f (&o_sf) = alloc.allocate(); +#endif + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_WARPS + 2]; +#if MK2_PREFILL_ONLY + bf16 *const route_stage = reinterpret_cast(&o_s[0]); + static_assert(sizeof(o_s) >= D_MODEL * sizeof(bf16), + "router row stage must fit the output tiles"); +#else + __shared__ __align__(16) uint64_t sort_keys[SORT_N]; + bf16 *const route_stage = reinterpret_cast(&sort_keys[0]); + static_assert(sizeof(sort_keys) >= D_MODEL * sizeof(bf16), + "router row stage must fit the sampler workspace"); +#endif + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS * NCONS, 0); + for (int c = 0; c < CONSUMER_WARPS * NCONS; ++c) arrive(drained[s]); + } + } + int gemm_iter = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map); + grid_barrier(g.barrier); + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map); + grid_barrier(g.barrier); + mamba_gated_norm(L, bb->zxbcdt, bb->gated); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, nullptr, 1, BATCH, + bb->moe_w, BB_MOE_INTER / OUT_ATOM); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + 1, BATCH, bb->moe_w, D_MODEL / OUT_ATOM); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_deterministic(bb->moe_out, idx, + (1u << BB_EXPERTS) - 1u); + grid_barrier(g.barrier); + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int original_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + bool identical_cohort = false; + if constexpr (BATCH > 1) { + if (bb->native_pf_cohort_check && bb->native_prefill_reset && + bb->pf_n_seq == BATCH && !bb->pf_fuse_decode && + !bb->pf_layer_out && !bb->moe_w_inject && !bb->moe_w_audit) { + if (blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[0] = 0; + grid_barrier(g.barrier); + const auto *bits = reinterpret_cast(bb->prompt); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < original_total * D_MODEL; idx += gridDim.x * THREADS) { + const int canonical = (idx / (BATCH * D_MODEL)) * BATCH * D_MODEL + idx % D_MODEL; + if (bits[idx] != bits[canonical]) atomicOr(bb->native_pf_cohort_check, 1u); + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + if (bb->pf_seq_slot[row] != row || bb->slot_map[row] != row) + atomicOr(bb->native_pf_cohort_check, 1u); + grid_barrier(g.barrier); + identical_cohort = bb->native_pf_cohort_check[0] == 0; + } + } + if (!identical_cohort) { + if (bb->native_pf_cohort_check && blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[1] = 0; + grid_barrier(g.barrier); + return; + } + const int n_total = original_total / BATCH; + if (bb->native_pf_cohort_check && blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[1] = identical_cohort ? 1 : BATCH; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int reset_slots = identical_cohort ? 1 : BATCH; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * reset_slots * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % reset_slots; + const MambaLayer &L = bb->mamba[row_layer / reset_slots]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = original_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + const int src = identical_cohort ? (idx / D_MODEL) * BATCH * D_MODEL + idx % D_MODEL : idx; + bb->pf_h[idx] = bb->prompt[src]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = identical_cohort ? 1 : (bb->pf_n_seq > 0 ? bb->pf_n_seq : 1); + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T, + (!bb->pf_layer_out && layer > 0 && + bb->moe_of_layer[layer - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + if (bb->native_pf_expert_masks) + for (int i = blockIdx.x * THREADS + threadIdx.x; i < n_tok; + i += gridDim.x * THREADS) + bb->native_pf_expert_masks[i] = 0u; + if (bb->pf_moe_ecnt) + for (int e = blockIdx.x * THREADS + threadIdx.x; e < BB_EXPERTS; + e += gridDim.x * THREADS) + bb->pf_moe_ecnt[e] = 0; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + for(size_t j=blockIdx.x*THREADS+threadIdx.x;jpf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + mamba_dtda_prefill(L, bb->pf_zxbcdt, bb->pf_dtda, T); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_dtda, + bb->pf_gated, T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + for(size_t j=blockIdx.x*THREADS+threadIdx.x;jpf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, + T, bb->native_pf_expert_masks, bb->pf_moe_gath, + bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, route_stage); + } + grid_barrier(g.barrier); + int max_ecnt = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + max_ecnt = max(max_ecnt, bb->pf_moe_ecnt[e]); + const bool narrow = BATCH_PAD > PAD_N && max_ecnt <= PAD_N; + const int etile = narrow ? PAD_N : BATCH_PAD; + int nt_sum = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_sum += (bb->pf_moe_ecnt[e] + etile - 1) / etile; + const int up_steps = + gemm_groups(BB_MOE_INTER / OUT_ATOM, nt_sum) * + nt_sum * (D_MODEL / KCHUNK); + const int shup_steps = + gemm_groups(BB_SHARED_INTER / OUT_ATOM, n_tok) * + n_tok * (D_MODEL / KCHUNK); + const int down_steps = + gemm_groups(D_MODEL / OUT_ATOM, nt_sum) * + nt_sum * (BB_MOE_INTER / KCHUNK); + const int shdown_steps = + gemm_groups(D_MODEL / OUT_ATOM, n_tok) * + n_tok * (BB_SHARED_INTER / KCHUNK); + if (narrow) + gemm_narrow(L.exp_up, bb->pf_moe_gath_vn, + BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, PAD_N, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + else + gemm(L.exp_up, bb->pf_moe_gath_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, BATCH_PAD, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + up_steps, up_steps + shup_steps); + grid_barrier(g.barrier); + if (narrow) + gemm_narrow(L.exp_down, bb->pf_moe_hc_vn, + BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, PAD_N, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + else + gemm(L.exp_down, bb->pf_moe_hc_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + down_steps, down_steps + shdown_steps); + grid_barrier(g.barrier); + if (bb->pf_layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * D_MODEL; idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = + bb_moe_fold_deterministic(bb->pf_moe_out, idx, + moe_row_mask(bb->pf_moe_w, idx / D_MODEL)); + grid_barrier(g.barrier); + } + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, + identical_cohort ? bb->pf_normed : bb->pf_final, ln_scratch, + n_total, + (!bb->pf_layer_out && n_backbone_layers > 0 && + bb->moe_of_layer[n_backbone_layers - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + grid_barrier(g.barrier); + if (identical_cohort) { + const int worker = blockIdx.x * THREADS + threadIdx.x; + const int workers = gridDim.x * THREADS; + for (int i = worker; i < original_total * D_MODEL; i += workers) + bb->pf_final[i] = bb->pf_normed[(i / (BATCH * D_MODEL)) * D_MODEL + i % D_MODEL]; + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = worker; i < 14 * (BATCH - 1) * CV; i += workers) { + const int cell = i % CV, pair = i / CV; + const int row = pair % (BATCH - 1) + 1; + const MambaLayer &L = bb->mamba[pair / (BATCH - 1)]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = L.conv_state[cell]; + } + const int kv_columns = BB_KV_HEADS * BB_HEAD_DIM; + for (size_t i = worker; i < size_t(3) * 2 * (BATCH - 1) * n_total * kv_columns; i += workers) { + size_t q = i / kv_columns; + const int col = i % kv_columns; + const int pos = q % n_total; q /= n_total; + const int row = q % (BATCH - 1) + 1; q /= BATCH - 1; + const int kind = q % 2, layer = q / 2; + const AttnLayer &L = bb->attn[layer]; + const int offset = kv_offset(L, pos, col / BB_HEAD_DIM) + col % BB_HEAD_DIM; + kv_slot(L, row, kind, pos)[offset] = kv_slot(L, 0, kind, pos)[offset]; + } + grid_barrier(g.barrier); + } + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } +#if MK2_PREFILL_ONLY + return; +#else + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, bb->phoneme_b); + grid_barrier(g.barrier); + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + if (threadIdx.x == 0) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = 0; n < BB_PHONEME; ++n) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + bb->phoneme_out[row] = arg; + } + __syncthreads(); + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + for (int i = lane; i < D_HEAD; i += 32) { + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + for (int t = 0; t <= pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + for (int i = lane; i < D_HEAD; i += 32) { + float acc = 0.f; + for (int t = 0; t <= pos; ++t) + acc += sc[t] * __bfloat162float( + g.vcache[cache_base + t * D_HEAD + i]); + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + i] = + __float2bfloat16(acc); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, g.sampled_pos, + sort_keys, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { +#ifdef ABLATE_SAMPLE + __syncthreads(); +#endif + } else { + grid_barrier(g.barrier); + } + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x < BATCH && blockIdx.x == 0) { + const int code = g.sampled_pos[threadIdx.x]; + g.sampled[threadIdx.x * N_POS + pos] = code; + if (g.native_codes) + g.native_codes[threadIdx.x * N_POS + pos] = (long long)code; + } + grid_barrier(g.barrier); + } + if (g.native_seq_len && blockIdx.x == 0 && threadIdx.x < BATCH) + g.native_seq_len[threadIdx.x] += 1; + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +#endif +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 2 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK), + 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = MK2_PREFILL_ONLY + ? out->sms * out->blocks_per_sm + : std::min(out->sms * out->blocks_per_sm, D_FFN / OUT_ATOM); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + 0, + nullptr, + nullptr, + nullptr, nullptr, nullptr, nullptr, nullptr, 0, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "dedup_native_prefill_impl.inc" diff --git a/deploy/nemotron_native/kernels/b1024/regions/prefill/native/prefill/mk2_abi.cuh b/deploy/nemotron_native/kernels/b1024/regions/prefill/native/prefill/mk2_abi.cuh new file mode 100644 index 0000000..95a253b --- /dev/null +++ b/deploy/nemotron_native/kernels/b1024/regions/prefill/native/prefill/mk2_abi.cuh @@ -0,0 +1,82 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} diff --git a/deploy/nemotron_native/kernels/b1024/regions/prefill/native/prefill/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/b1024/regions/prefill/native/prefill/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/b1024/regions/prefill/native/prefill/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/b1024/regions/prefill/native/prefill/prefill.cu b/deploy/nemotron_native/kernels/b1024/regions/prefill/native/prefill/prefill.cu new file mode 100644 index 0000000..c21de0b --- /dev/null +++ b/deploy/nemotron_native/kernels/b1024/regions/prefill/native/prefill/prefill.cu @@ -0,0 +1,18 @@ +#define MK2_BATCH 1024 +#define MK2_BATCH_TILE 64 +#define MK2_PREFILL_ONLY 1 +#define STAGES 3 +#define KCHUNK 128 +#define MK2_GEMM_CONSUMERS 3 +#define MK2_THREADS 512 +#define MK2_NO_MAIN 1 +#define MK2_DIRECT_INPUT 1 +#define MK2_B1_LOCAL_BARRIERS 1 +#define MK2_SKIP_HIDDEN_SNAPSHOT 1 +#define MK2_ACTIVE_EXPERTS 1 +#define MK2_SSM_SPLITS 1 +#define PF_SSM_PW_MULTI 12 +#define PF_ACTIVE_TOKEN_EXPERTS 1 +#include +#include +#include "dedup_predictor_mk2.cuh" diff --git a/deploy/nemotron_native/kernels/b1024/regions/prefill/prefill.cu b/deploy/nemotron_native/kernels/b1024/regions/prefill/prefill.cu new file mode 100644 index 0000000..c99bc6a --- /dev/null +++ b/deploy/nemotron_native/kernels/b1024/regions/prefill/prefill.cu @@ -0,0 +1,20 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace b1024_prefill { +#include "native/prefill/prefill.cu" +} diff --git a/deploy/nemotron_native/kernels/b1024/regions/stage0/native/stage0/mk2_abi.cuh b/deploy/nemotron_native/kernels/b1024/regions/stage0/native/stage0/mk2_abi.cuh new file mode 100644 index 0000000..bf1e415 --- /dev/null +++ b/deploy/nemotron_native/kernels/b1024/regions/stage0/native/stage0/mk2_abi.cuh @@ -0,0 +1,85 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_set_native_predictor_active_mask(MK2Predictor *p, const unsigned char *mask); +int mk2_native_startup_skip_abi_version(); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *, int grid); diff --git a/deploy/nemotron_native/kernels/b1024/regions/stage0/native/stage0/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/b1024/regions/stage0/native/stage0/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/b1024/regions/stage0/native/stage0/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/b1024/regions/stage0/native/stage0/native_prefill_impl.inc b/deploy/nemotron_native/kernels/b1024/regions/stage0/native/stage0/native_prefill_impl.inc new file mode 100644 index 0000000..a254308 --- /dev/null +++ b/deploy/nemotron_native/kernels/b1024/regions/stage0/native/stage0/native_prefill_impl.inc @@ -0,0 +1,198 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 16 || BATCH == 32 || BATCH == 64)) + throw std::runtime_error("native prefill requires matching batch 1/16/32/64"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (padded > INT_MAX / ((BB_EXPERTS + 1) * D_MODEL)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + float *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(float))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(float))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} diff --git a/deploy/nemotron_native/kernels/b1024/regions/stage0/native/stage0/predictor_mk2.cu b/deploy/nemotron_native/kernels/b1024/regions/stage0/native/stage0/predictor_mk2.cu new file mode 100644 index 0000000..dbc7d10 --- /dev/null +++ b/deploy/nemotron_native/kernels/b1024/regions/stage0/native/stage0/predictor_mk2.cu @@ -0,0 +1,6111 @@ +#include "../../../../../inflight/progress.cuh" +#if MK2_INFLIGHT +#include +#endif +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#define MK2_FUSE_CONV (MK2_SSM_SPLITS == 1) +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +constexpr int BATCH_PAD = 64; +constexpr int DECODE_TILES = BATCH / BATCH_PAD; +constexpr int DECODE_ROWS = BATCH_PAD; +static_assert(BATCH == 512 || BATCH == 1024 || BATCH == 2048); +static_assert(BATCH % BATCH_PAD == 0); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +#ifndef NMULT +#define NMULT 1 +#endif +#ifndef SSM_PREFETCH_MAX +#define SSM_PREFETCH_MAX 8 +#endif +constexpr int NM_MAX = (NMULT > 1 && DECODE_TILES % NMULT == 0) ? NMULT : 1; +#ifndef MK2_OUT_MULT +#define MK2_OUT_MULT 1 +#endif +constexpr int MT = MK2_OUT_MULT; +static_assert(MT >= 1 && MT <= NM_MAX, "an item holds at most NM_MAX accumulators"); +constexpr int SLOTS = 1 + NM_MAX; +static_assert(MT + 1 <= SLOTS); +constexpr int THREADS = 256; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +constexpr int NORM_SCRATCH_STRIDE = 2 * WARPS + 2; +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_ACC_F32 }; +#define EPI_ROW(idx) ((OUT_ATOM == 64 && BATCH_PAD == 64) \ + ? (((idx) & 15) | ((((idx) >> 6) & 3) << 4)) : ((idx) / (OUT_ATOM / 8))) +#define EPI_OCT(idx) ((OUT_ATOM == 64 && BATCH_PAD == 64) \ + ? ((((idx) >> 4) & 3) | ((((idx) >> 8) & 1) << 2)) : ((idx) % (OUT_ATOM / 8))) +struct alignas(16) bf16_oct { bf16 x0, x1, x2, x3, x4, x5, x6, x7; }; +static_assert(sizeof(bf16_oct) == 16); +#include +constexpr int R9_STATE_MODE=1; +using r9_state_t=bf16; +__device__ __forceinline__ int r9_scale_index(int head,int channel,int n) {return (((head/8)*4+(head%8)/2)*4+channel/6)*128+n;} +__device__ __forceinline__ float r9_load(r9_state_t v,float scale) { +return __bfloat162float(v); +} +__device__ __forceinline__ r9_state_t r9_store(float x,float scale) { +return __float2bfloat16(x); +} +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; + const float *r8_initial_ssm = nullptr; + const unsigned int *r8_initial_ready = nullptr; + const int *r8_seq_len = nullptr; + int r8_initial_positions = 0; + const float *r9_scales = nullptr; + const int *r9_map = nullptr; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + bf16 *r6_packed = nullptr; + int *r6_map = nullptr,*r6_counts = nullptr; + act_gl *r6_views = nullptr; + int *r6_lut = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +#define BARRIER_GEN_WORD 32 +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + constexpr unsigned long long GRID_BLOCKS = 132ull; + auto *wide_state = reinterpret_cast(state); + const unsigned long long arrival = atomicAdd(wide_state, 1ull); + const unsigned int rank = (unsigned int)(arrival % GRID_BLOCKS); + if (rank != GRID_BLOCKS - 1u) { + const unsigned long long target = arrival - rank + GRID_BLOCKS; + volatile unsigned long long *counter = wide_state; + while (*counter < target) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / THREADS; + float wv[PER_THREAD]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) wv[j] = weight[threadIdx.x + j * THREADS]; + float v[PER_THREAD], v_next[PER_THREAD]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + v[j] = __bfloat162float( + src[(size_t)blockIdx.x * D_MODEL + threadIdx.x + j * THREADS]); + for (int row = blockIdx.x, row_iter = 0; row < BATCH; + row += gridDim.x, ++row_iter) { + float *row_scratch = scratch + (row_iter & 1) * NORM_SCRATCH_STRIDE; + const int next_row = row + (int)gridDim.x; + if (next_row < BATCH) { + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + v_next[j] = __bfloat162float( + src[(size_t)next_row * D_MODEL + threadIdx.x + j * THREADS]); + } + float sum = 0.f, sq = 0.f; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) { + row_scratch[warp] = sum; + row_scratch[WARPS + warp] = sq; + } + __syncthreads(); + float s = 0.f, q = 0.f; + #pragma unroll + for (int w = 0; w < WARPS; ++w) { + s += row_scratch[w]; + q += row_scratch[WARPS + w]; + } + const float mean = s / D_MODEL; + const float inv = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * wv[j]); + } + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) v[j] = v_next[j]; + } +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, bf16 *snapshot = nullptr, int snapshot_layer = -1, + float *snapshot_fp32 = nullptr, float *out_fp32 = nullptr, + const float *fold_src = nullptr, const float *fold_routes = nullptr) { + constexpr int PER_THREAD = D_MODEL / THREADS; + bf16 wv[PER_THREAD]; + extern __shared__ int __shm[]; + bf16 *const row_stage = reinterpret_cast(__shm); + constexpr int ROW_SLOT = 2 * D_MODEL; + constexpr int ROW_VECTORS = D_MODEL / 8; + if ((int)blockIdx.x < n_rows && threadIdx.x < ROW_VECTORS) { + const int i = 8 * threadIdx.x; + const size_t at = (size_t)blockIdx.x * D_MODEL + i; + const uint32_t rs = static_cast( + __cvta_generic_to_shared(row_stage + D_MODEL + i)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;" + :: "r"(rs), "l"(r + at) : "memory"); + if (!fold_src) { + const uint32_t hs = static_cast( + __cvta_generic_to_shared(row_stage + i)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;" + :: "r"(hs), "l"(h + at) : "memory"); + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + wv[j] = weight[threadIdx.x + j * THREADS]; + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); + __syncthreads(); + for (int row = blockIdx.x, row_iter = 0; row < n_rows; + row += gridDim.x, ++row_iter) { + bf16 *const current = row_stage + (row_iter & 1) * ROW_SLOT; + const int next_row = row + (int)gridDim.x; + if (next_row < n_rows) { + bf16 *const next = row_stage + ((row_iter + 1) & 1) * ROW_SLOT; + if (threadIdx.x < ROW_VECTORS) { + const int i = 8 * threadIdx.x; + const size_t at = (size_t)next_row * D_MODEL + i; + const uint32_t rs = static_cast( + __cvta_generic_to_shared(next + D_MODEL + i)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;" + :: "r"(rs), "l"(r + at) : "memory"); + if (!fold_src) { + const uint32_t hs = static_cast( + __cvta_generic_to_shared(next + i)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;" + :: "r"(hs), "l"(h + at) : "memory"); + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } + float *row_scratch = scratch + (row_iter & 1) * NORM_SCRATCH_STRIDE; + const size_t base = (size_t)row * D_MODEL; + float fv[PER_THREAD]; + if (fold_src) { + unsigned mask = 0u; + const float *rw = fold_routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) + if (rw[e] != 0.f) mask |= 1u << e; + const float *pe = fold_src + + (size_t)row * (BB_EXPERTS + 1) * D_MODEL + threadIdx.x; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) fv[j] = 0.f; + #pragma unroll + for (int e = 0; e <= BB_EXPERTS; ++e) { + const bool take = e == BB_EXPERTS || ((mask >> e) & 1u); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + fv[j] = __fadd_rn(fv[j], + take ? pe[(size_t)e * D_MODEL + j * THREADS] : 0.f); + } + } + float v[PER_THREAD]; + float sq = 0.f; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + const bf16 hv = fold_src ? __float2bfloat16(fv[j]) : current[i]; + if (fold_src) h[base + i] = hv; + if (snapshot) + snapshot[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = hv; + if (snapshot_fp32) + snapshot_fp32[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = + __bfloat162float(hv); + v[j] = __bfloat162float(hv) + + __bfloat162float(current[D_MODEL + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) row_scratch[warp] = sq; + __syncthreads(); + float s = 0.f; + #pragma unroll + for (int w = 0; w < WARPS; ++w) s += row_scratch[w]; + const float inv = rsqrtf(s / D_MODEL + RMS_EPS); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + const bf16 normed = __float2bfloat16( + __bfloat162float(wv[j]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + out[base + i] = normed; + if (out_fp32) out_fp32[base + i] = __bfloat162float(normed); + } + if (next_row < n_rows) { + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); + __syncthreads(); + } + } +} +__device__ __forceinline__ bf16 mamba_conv_channel( + const MambaLayer &L, bf16 *state_slot, const bf16 *zrow, int c) { + const unsigned ch = (unsigned)c; + const bf16 *w = L.conv_w + ch * (unsigned)MB_CONV_K; + bf16 *state = state_slot + ch * (unsigned)(MB_CONV_K - 1); + const bf16 rawb = zrow[c]; + const float raw = __bfloat162float(rawb); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = rawb; + return __float2bfloat16(acc); +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map MK2_PROGRESS_PARAMETER(progress_mask)) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + MK2_PROGRESS_GUARD(progress_mask, row); + const int c = idx - row * MB_CONV_DIM; + xbc_act[(size_t)row * MB_CONV_DIM + c] = mamba_conv_channel( + L, L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride, + zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER, c); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map MK2_PROGRESS_PARAMETER(progress_mask)) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int sb4 = (lane >> 4) & 1, sb3 = (lane >> 3) & 1; + const int sb2 = (lane >> 2) & 1, sb1 = (lane >> 1) & 1; + const int ssm_own_ch = 6 * sb4 + 3 * sb3 + 2 * sb2 + sb1; + const bool ssm_owner = (2 * sb2 + sb1) < 3 && (lane & 1) == 0; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + constexpr bool FUSE_GNORM = (split == 1); + extern __shared__ int __shm[]; + bf16 (*const gn_stage)[MB_GROUP_SIZE] = + reinterpret_cast(&__shm[0]); + constexpr bool FUSE_CONV = (MK2_FUSE_CONV != 0); + bf16 (*const conv_x)[MB_GROUP_SIZE] = gn_stage + WARPS; + constexpr int SSM_NSTATE = MB_STATE / 32; + constexpr int SSM_PF = channels < SSM_PREFETCH_MAX ? channels : SSM_PREFETCH_MAX; + constexpr int SSM_GROUP_BYTES = SSM_PF * MB_STATE * (int)sizeof(r9_state_t); + constexpr bool SSM_SMEM_OK = R9_STATE_MODE == 1 && channels == 2 * SSM_PF && + SSM_GROUP_BYTES % 512 == 0; + constexpr int SSM_RING = 4; + constexpr int SSM_GROUPS_PER_HEAD = channels / SSM_PF; + constexpr int SSM_PAIRS = MB_HEADS_PER_GROUP * SSM_GROUPS_PER_HEAD; + static_assert(!SSM_SMEM_OK || (SSM_RING & (SSM_RING - 1)) == 0, + "the flat (head, channel group) ring index is masked"); + static_assert(!SSM_SMEM_OK || SSM_PAIRS >= SSM_RING, + "the prologue must be able to fill the ring"); + static_assert(!SSM_SMEM_OK || + 2 * WARPS * MB_GROUP_SIZE * sizeof(bf16) + + WARPS * SSM_RING * SSM_GROUP_BYTES <= STAGES * SLOTS * sizeof(weight_tile)); + constexpr int SSM_F32_RING = 2; + constexpr int SSM_F32_GROUP_BYTES = SSM_PF * MB_STATE * (int)sizeof(float); + constexpr bool SSM_F32_OK = SSM_SMEM_OK && SSM_F32_GROUP_BYTES % 512 == 0 && + (SSM_F32_RING & (SSM_F32_RING - 1)) == 0 && SSM_PAIRS >= SSM_F32_RING && + SSM_F32_RING * SSM_F32_GROUP_BYTES <= SSM_RING * SSM_GROUP_BYTES; + char *const ssm_stage = reinterpret_cast(conv_x + WARPS); + const int pairs = BATCH * MB_GROUPS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_GROUPS * split); + MK2_PROGRESS_GUARD(progress_mask, row); + const int group = (pair / split) % MB_GROUPS; + const int first_channel = (pair % split) * channels; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const bool r8_initial = L.r8_initial_ssm && L.r8_initial_ready && + *L.r8_initial_ready == 1 && L.r8_seq_len[row] == L.r8_initial_positions; + r9_state_t *st = reinterpret_cast(L.ssm_state) + + (size_t)slot_map[row] * L.ssm_slot_stride; + char *const warp_ssm = + ssm_stage + (size_t)warp * SSM_RING * SSM_GROUP_BYTES; + auto issue_ssm = [&](int head_in_group, int channel0, int buffer) { + const int head = group * MB_HEADS_PER_GROUP + head_in_group; + const r9_state_t *src = st + + (size_t)(head * MB_HEAD_DIM + channel0) * MB_STATE; + const uint32_t dst = (uint32_t)__cvta_generic_to_shared( + warp_ssm + (size_t)buffer * SSM_GROUP_BYTES + lane * 16); + #pragma unroll + for (int chunk = 0; chunk < SSM_GROUP_BYTES / 512; ++chunk) + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: + "r"(dst + (uint32_t)chunk * 512), + "l"(reinterpret_cast(src) + + (size_t)chunk * 512 + lane * 16) : "memory"); + asm volatile("cp.async.commit_group;\n" ::: "memory"); + }; + auto issue_ssm_f32 = [&](int head_in_group, int channel0, int buffer) { + const int head = group * MB_HEADS_PER_GROUP + head_in_group; + const float *src = L.r8_initial_ssm + + (size_t)(head * MB_HEAD_DIM + channel0) * MB_STATE; + const uint32_t dst = (uint32_t)__cvta_generic_to_shared( + warp_ssm + (size_t)buffer * SSM_F32_GROUP_BYTES + lane * 16); + #pragma unroll + for (int chunk = 0; chunk < SSM_F32_GROUP_BYTES / 512; ++chunk) + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: + "r"(dst + (uint32_t)chunk * 512), + "l"(reinterpret_cast(src) + + (size_t)chunk * 512 + lane * 16) : "memory"); + asm volatile("cp.async.commit_group;\n" ::: "memory"); + }; + if constexpr (SSM_SMEM_OK) { + if (!r8_initial) { + #pragma unroll + for (int k = 0; k < SSM_RING; ++k) + issue_ssm(k / SSM_GROUPS_PER_HEAD, + (k % SSM_GROUPS_PER_HEAD) * SSM_PF, k); + } else if constexpr (SSM_F32_OK) { + #pragma unroll + for (int k = 0; k < SSM_F32_RING; ++k) + issue_ssm_f32(k / SSM_GROUPS_PER_HEAD, + (k % SSM_GROUPS_PER_HEAD) * SSM_PF, k); + } + } + float Braw[MB_STATE / 32]; + float Cv[MB_STATE / 32]; + const bf16 *xrow; + if constexpr (FUSE_CONV) { + bf16 *const cslot = + L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride; + const bf16 *const zrow = + zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER; + static_assert(MB_CONV_K == 4, "the state record is three 32-bit words"); + static_assert(MB_GROUP_SIZE % 64 == 0 && MB_STATE % 64 == 0, + "one paired round covers 64 channels of a warp"); + auto wide_lo = [](unsigned int v) { return __uint_as_float(v << 16); }; + auto wide_hi = [](unsigned int v) { return __uint_as_float(v & 0xffff0000u); }; + struct ConvCursor { + const uint2 *w; unsigned int *s; + const unsigned int *z; const unsigned int *b; + }; + auto conv_cursor = [&](int c) { + const unsigned int cu = (unsigned int)c; + return ConvCursor{ + reinterpret_cast( + L.conv_w + cu * (unsigned int)MB_CONV_K), + reinterpret_cast( + cslot + cu * (unsigned int)(MB_CONV_K - 1)), + reinterpret_cast(zrow + cu), + reinterpret_cast(L.conv_b + cu)}; + }; + auto conv_two = [&](uint2 w0, uint2 w1, unsigned int s0, + unsigned int s1, unsigned int s2, + unsigned int zp, unsigned int bp, float &a0, + float &a1, unsigned int &n0, unsigned int &n1, + unsigned int &n2) { + a0 = wide_lo(bp); + a0 += wide_lo(w0.x) * wide_lo(s0); + a0 += wide_hi(w0.x) * wide_hi(s0); + a0 += wide_lo(w0.y) * wide_lo(s1); + a0 += wide_hi(w0.y) * wide_lo(zp); + a0 = a0 / (1.f + __expf(-a0)); + a1 = wide_hi(bp); + a1 += wide_lo(w1.x) * wide_hi(s1); + a1 += wide_hi(w1.x) * wide_lo(s2); + a1 += wide_lo(w1.y) * wide_hi(s2); + a1 += wide_hi(w1.y) * wide_hi(zp); + a1 = a1 / (1.f + __expf(-a1)); + n0 = (s0 >> 16) | (s1 << 16); + n1 = (zp & 0xffffu) | (s2 << 16); + n2 = (s2 >> 16) | (zp & 0xffff0000u); + }; + struct ConvPairIn { + uint2 w0, w1; + unsigned int s0, s1, s2, zp, bp; + }; + auto conv_pair_fetch = [](const ConvCursor &q, int r) { + const uint2 *const wp = q.w + 64 * r; + const unsigned int *const sp = q.s + 96 * r; + return ConvPairIn{wp[0], wp[1], sp[0], sp[1], sp[2], + q.z[32 * r], q.b[32 * r]}; + }; + auto conv_pair_apply = [&](unsigned int *const sp, + const ConvPairIn &v, float &a0, + float &a1) { + unsigned int n0, n1, n2; + conv_two(v.w0, v.w1, v.s0, v.s1, v.s2, v.zp, v.bp, a0, a1, + n0, n1, n2); + sp[0] = n0; + sp[1] = n1; + sp[2] = n2; + }; + constexpr int QUAD_SPAN = 4 * 32; + static_assert(MB_INNER % 4 == 0 && MB_STATE % 4 == 0 && + MB_GROUP_SIZE % 4 == 0, + "every quad base is a four-channel boundary"); + static_assert(MB_STATE == QUAD_SPAN && + MB_GROUP_SIZE == QUAD_SPAN + 2 * 32, + "one quad round covers B and C; x adds a pair round"); + struct ConvQuad { + const uint4 *w; uint2 *s; const uint2 *z; const uint2 *b; + }; + auto conv_quad_cursor = [&](int c) { + const unsigned int cu = (unsigned int)c; + return ConvQuad{ + reinterpret_cast( + L.conv_w + cu * (unsigned int)MB_CONV_K), + reinterpret_cast( + cslot + cu * (unsigned int)(MB_CONV_K - 1)), + reinterpret_cast(zrow + cu), + reinterpret_cast(L.conv_b + cu)}; + }; + struct ConvQuadIn { + uint4 w01, w23; + uint2 sa, sb, sc, zp, bp; + }; + auto conv_quad_fetch = [](const ConvQuad &q) { + const uint2 *const sp = q.s; + return ConvQuadIn{q.w[0], q.w[1], sp[0], sp[1], sp[2], + q.z[0], q.b[0]}; + }; + auto conv_quad_apply = [&](uint2 *const sp, const ConvQuadIn &v, + float &a0, float &a1, float &a2, + float &a3) { + unsigned int n0, n1, n2, n3, n4, n5; + conv_two(make_uint2(v.w01.x, v.w01.y), + make_uint2(v.w01.z, v.w01.w), + v.sa.x, v.sa.y, v.sb.x, v.zp.x, v.bp.x, + a0, a1, n0, n1, n2); + conv_two(make_uint2(v.w23.x, v.w23.y), + make_uint2(v.w23.z, v.w23.w), + v.sb.y, v.sc.x, v.sc.y, v.zp.y, v.bp.y, + a2, a3, n3, n4, n5); + sp[0] = make_uint2(n0, n1); + sp[1] = make_uint2(n2, n3); + sp[2] = make_uint2(n4, n5); + }; + auto conv_pack = [](float a, float b) { + return (unsigned int)__bfloat16_as_ushort(__float2bfloat16(a)) | + ((unsigned int)__bfloat16_as_ushort(__float2bfloat16(b)) << 16); + }; + const ConvQuad qb = conv_quad_cursor( + MB_INNER + group * MB_STATE + 4 * lane); + const ConvQuad qc = conv_quad_cursor( + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE + 4 * lane); + static_assert(MB_GROUP_SIZE * (int)sizeof(bf16) % 8 == 0, + "each warp's staging row starts 8-byte aligned"); + const ConvQuad qx = conv_quad_cursor( + group * MB_GROUP_SIZE + 4 * lane); + const ConvCursor qt = conv_cursor( + group * MB_GROUP_SIZE + QUAD_SPAN + 2 * lane); + const ConvQuadIn vb = conv_quad_fetch(qb); + const ConvQuadIn vc = conv_quad_fetch(qc); + const ConvQuadIn vx = conv_quad_fetch(qx); + const ConvPairIn vt = conv_pair_fetch(qt, 0); + uint2 bpk, cpk; + { + float a0, a1, a2, a3; + conv_quad_apply(qb.s, vb, a0, a1, a2, a3); + bpk = make_uint2(conv_pack(a0, a1), conv_pack(a2, a3)); + conv_quad_apply(qc.s, vc, a0, a1, a2, a3); + cpk = make_uint2(conv_pack(a0, a1), conv_pack(a2, a3)); + } + { + float a0, a1, a2, a3; + conv_quad_apply(qx.s, vx, a0, a1, a2, a3); + *reinterpret_cast(&conv_x[warp][4 * lane]) = + make_uint2(conv_pack(a0, a1), conv_pack(a2, a3)); + } + { + float a0, a1; + conv_pair_apply(qt.s, vt, a0, a1); + conv_x[warp][QUAD_SPAN + 2 * lane] = __float2bfloat16(a0); + conv_x[warp][QUAD_SPAN + 2 * lane + 1] = __float2bfloat16(a1); + } + const bool qword = ((lane >> 1) & 1) != 0; + const bool qhalf = (lane & 1) != 0; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int src = (lane >> 2) + 8 * j; + const unsigned int bx = __shfl_sync(0xffffffffu, bpk.x, src); + const unsigned int by = __shfl_sync(0xffffffffu, bpk.y, src); + const unsigned int cx = __shfl_sync(0xffffffffu, cpk.x, src); + const unsigned int cy = __shfl_sync(0xffffffffu, cpk.y, src); + const unsigned int bg = qword ? by : bx; + const unsigned int cg = qword ? cy : cx; + Braw[j] = qhalf ? wide_hi(bg) : wide_lo(bg); + Cv[j] = qhalf ? wide_hi(cg) : wide_lo(cg); + } + __syncwarp(); + xrow = &conv_x[warp][0]; + } else { + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Braw[j] = __bfloat162float(Bsec[lane + j * 32]); + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + xrow = xa + (size_t)group * MB_GROUP_SIZE; + } + const int head_lane = group * MB_HEADS_PER_GROUP + + (lane < MB_HEADS_PER_GROUP ? lane : 0); + const float dt_lane = mamba_softplus( + __bfloat162float(dtrow[head_lane]) + L.dt_bias[head_lane]); + const float dA_lane = __expf(L.A[head_lane] * dt_lane); + const float D_lane = L.D[head_lane]; + #pragma unroll 1 + for (int head_in_group = 0; head_in_group < MB_HEADS_PER_GROUP; + ++head_in_group) { + const int head = group * MB_HEADS_PER_GROUP + head_in_group; + const float dt = __shfl_sync(0xffffffffu, dt_lane, head_in_group); + const float dA = __shfl_sync(0xffffffffu, dA_lane, head_in_group); + const float Dh = __shfl_sync(0xffffffffu, D_lane, head_in_group); + const float *st_in = L.r8_initial_ssm + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float Bv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) + Bv[j] = Braw[j] * dt; + constexpr int NSTATE = SSM_NSTATE; + constexpr int SSM_PREFETCH = SSM_PF; + static_assert(channels % SSM_PREFETCH == 0); + for (int g0 = 0; g0 < channels; g0 += SSM_PREFETCH) { + float sv[SSM_PREFETCH][NSTATE]; + float xvr[SSM_PREFETCH]; + if (r8_initial && SSM_F32_OK) { + asm volatile("cp.async.wait_group 1;\n" ::: "memory"); + __syncwarp(); + const int pair = head_in_group * SSM_GROUPS_PER_HEAD + + g0 / SSM_PREFETCH; + const int buffer = pair & (SSM_F32_RING - 1); + const float *const smf = reinterpret_cast( + warp_ssm + (size_t)buffer * SSM_F32_GROUP_BYTES); + static_assert(SSM_PREFETCH % 2 == 0 && MB_HEAD_DIM % 2 == 0, + "x staging vectorises in element pairs"); + const __nv_bfloat162 *const xq = + reinterpret_cast( + xrow + head_in_group * MB_HEAD_DIM + + first_channel + g0); + #pragma unroll + for (int gq = 0; gq < SSM_PREFETCH / 2; ++gq) { + const float2 xp = __bfloat1622float2(xq[gq]); + xvr[2 * gq] = xp.x; + xvr[2 * gq + 1] = xp.y; + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = smf[(size_t)gi * MB_STATE + lane + j * 32]; + } + __syncwarp(); + const int ahead = pair + SSM_F32_RING; + if (ahead < SSM_PAIRS) + issue_ssm_f32(ahead / SSM_GROUPS_PER_HEAD, + (ahead % SSM_GROUPS_PER_HEAD) * SSM_PF, + buffer); + else + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } else if (r8_initial) { + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const size_t prow = (size_t)(first_channel + g0 + gi) * MB_STATE; + xvr[gi] = __bfloat162float( + xrow[head_in_group * MB_HEAD_DIM + + first_channel + g0 + gi]); + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = st_in[prow + lane + j * 32]; + } + } else if constexpr (SSM_SMEM_OK) { + asm volatile("cp.async.wait_group 3;\n" ::: "memory"); + __syncwarp(); + const int pair = head_in_group * SSM_GROUPS_PER_HEAD + + g0 / SSM_PREFETCH; + const int buffer = pair & (SSM_RING - 1); + const r9_state_t *sm = reinterpret_cast( + warp_ssm + (size_t)buffer * SSM_GROUP_BYTES); + static_assert(SSM_PREFETCH % 2 == 0 && MB_HEAD_DIM % 2 == 0, + "x staging vectorises in element pairs"); + const __nv_bfloat162 *const xq = + reinterpret_cast( + xrow + head_in_group * MB_HEAD_DIM + + first_channel + g0); + #pragma unroll + for (int gq = 0; gq < SSM_PREFETCH / 2; ++gq) { + const float2 xp = __bfloat1622float2(xq[gq]); + xvr[2 * gq] = xp.x; + xvr[2 * gq + 1] = xp.y; + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = r9_load( + sm[(size_t)gi * MB_STATE + lane + j * 32], 1.f); + } + __syncwarp(); + const int ahead = pair + SSM_RING; + if (ahead < SSM_PAIRS) + issue_ssm(ahead / SSM_GROUPS_PER_HEAD, + (ahead % SSM_GROUPS_PER_HEAD) * SSM_PF, buffer); + else + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } else { + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + xvr[gi] = __bfloat162float( + xrow[head_in_group * MB_HEAD_DIM + + first_channel + g0 + gi]); + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + { + int p = first_channel+g0+gi; + int mapped = R9_STATE_MODE==4 ? L.r9_map[head*24+p] : head*24+p; + int n = lane+j*32; + float scale = R9_STATE_MODE==1 ? 1.f : L.r9_scales[r9_scale_index(mapped/24,mapped%24,n)]; + sv[gi][j] = r9_load(st[(size_t)mapped*128+n],scale); + } + } + } + if constexpr (SSM_PREFETCH == 12) { + float v6[6]; + #pragma unroll + for (int k = 0; k < 6; ++k) { + float pair_acc[2]; + #pragma unroll + for (int half = 0; half < 2; ++half) { + const int gi = k + 6 * half; + const int p = first_channel + g0 + gi; + const float xv = xvr[gi]; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < NSTATE; ++j) { + const int n = lane + j * 32; + const float s = sv[gi][j] * dA + Bv[j] * xv; + const int mapped = R9_STATE_MODE == 4 + ? L.r9_map[head * 24 + p] : head * 24 + p; + const float scale = R9_STATE_MODE == 1 ? 1.f + : L.r9_scales[r9_scale_index( + mapped / 24, mapped % 24, n)]; + st[(size_t)mapped * 128 + n] = r9_store(s, scale); + acc += s * Cv[j]; + } + pair_acc[half] = acc; + } + const float mine = sb4 ? pair_acc[1] : pair_acc[0]; + const float away = sb4 ? pair_acc[0] : pair_acc[1]; + v6[k] = mine + __shfl_xor_sync(0xffffffff, away, 16); + } + float v4[4]; + #pragma unroll + for (int k = 0; k < 3; ++k) { + const float mine = sb3 ? v6[3 + k] : v6[k]; + const float away = sb3 ? v6[k] : v6[3 + k]; + v4[k] = mine + __shfl_xor_sync(0xffffffff, away, 8); + } + v4[3] = 0.f; + float v2[2]; + #pragma unroll + for (int k = 0; k < 2; ++k) { + const float mine = sb2 ? v4[2 + k] : v4[k]; + const float away = sb2 ? v4[k] : v4[2 + k]; + v2[k] = mine + __shfl_xor_sync(0xffffffff, away, 4); + } + float acc = (sb1 ? v2[1] : v2[0]) + + __shfl_xor_sync(0xffffffff, sb1 ? v2[0] : v2[1], 2); + acc += __shfl_xor_sync(0xffffffff, acc, 1); + if (ssm_owner) { + const int p = first_channel + g0 + ssm_own_ch; + const float xv = __bfloat162float( + xrow[head_in_group * MB_HEAD_DIM + p]); + const bf16 outv = __float2bfloat16(acc + xv * Dh); + if constexpr (FUSE_GNORM) + gn_stage[warp][head_in_group * MB_HEAD_DIM + p] = outv; + else + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = outv; + } + } else { + float cacc[SSM_PREFETCH]; + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const int p = first_channel + g0 + gi; + const float xv = xvr[gi]; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < NSTATE; ++j) { + const int n = lane + j * 32; + const float s = sv[gi][j] * dA + Bv[j] * xv; + const int mapped = R9_STATE_MODE == 4 + ? L.r9_map[head * 24 + p] : head * 24 + p; + const float scale = R9_STATE_MODE == 1 ? 1.f + : L.r9_scales[r9_scale_index( + mapped / 24, mapped % 24, n)]; + st[(size_t)mapped * 128 + n] = r9_store(s, scale); + acc += s * Cv[j]; + } + cacc[gi] = acc; + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const int p = first_channel + g0 + gi; + float acc = cacc[gi]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) { + const bf16 outv = __float2bfloat16(acc + xvr[gi] * Dh); + if constexpr (FUSE_GNORM) + gn_stage[warp][head_in_group * MB_HEAD_DIM + p] = outv; + else + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = outv; + } + } + } + } + } + if constexpr (SSM_SMEM_OK) { + if (!r8_initial || SSM_F32_OK) { + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); + __syncwarp(); + } + } + if constexpr (FUSE_GNORM) { + __syncwarp(); + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32], wg[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float gv = __bfloat162float(gn_stage[warp][lane + j * 32]) * + (z / (1.f + __expf(-z))); + wg[j] = __bfloat162float(L.gnorm_w[idx]); + held[j] = gv; + sq += gv * gv; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16(held[j] * inv * wg[j]); + } + __syncwarp(); + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_CONV_DIM; + const int c = idx - t * MB_CONV_DIM; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v; +#ifdef PF_CONV_ZERO_STATE + if (src < 0) v = 0.f; +#else + if (src < 0) v = __bfloat162float(state[pos + k]); +#endif + else v = __bfloat162float( + zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]); + acc += __bfloat162float(w[k]) * v; + } +#ifndef PF_CONV_NO_SILU + acc = acc / (1.f + __expf(-acc)); +#endif + xbc_act[(size_t)t * MB_CONV_DIM + c] = __float2bfloat16(acc); + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + const int units = n_seq * MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int seq = unit / (MB_HEADS * MB_HEAD_DIM); + const int rem = unit - seq * (MB_HEADS * MB_HEAD_DIM); + const int head = rem / MB_HEAD_DIM; + const int p = rem - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p * MB_STATE; + float sreg[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) sreg[j] = st[lane + j * 32]; + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const bf16 *xa = xbc_act + t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + const bf16 *dtrow = zxbcdt + t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + sreg[j] = sreg[j] * dA + __bfloat162float(Bsec[n]) * dt * xv; + acc += sreg[j] * __bfloat162float(Csec[n]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) st[lane + j * 32] = sreg[j]; + } + return; + } +#if PF_SSM_SPLIT_P + const int units = MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; unit += warps_total) { + const int head = unit / MB_HEAD_DIM; + const int p_only = unit - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; +#else + for (int head = blockIdx.x * WARPS + warp; head < MB_HEADS; + head += warps_total) { + const int group = head / MB_HEADS_PER_GROUP; +#endif + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *st = L.ssm_state + (size_t)slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE +#if PF_SSM_SPLIT_P + + (size_t)p_only * MB_STATE +#endif + ; + for (int t = 0; t < n_tokens; ++t) { + const bf16 *xa = xbc_act + (size_t)t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = + zxbcdt + (size_t)t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } +#if PF_SSM_SPLIT_P + { + const int p = p_only; +#else + for (int p = 0; p < MB_HEAD_DIM; ++p) { +#endif + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { +#if PF_SSM_SPLIT_P + const size_t at = (size_t)lane + j * 32; +#else + const size_t at = (size_t)p * MB_STATE + lane + j * 32; +#endif + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } + } +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32], wg[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + wg[j] = __bfloat162float(L.gnorm_w[idx]); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16(held[j] * inv * wg[j]); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len MK2_PROGRESS_PARAMETER(progress_mask)) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + constexpr int Q_PER_WARP = 3; + constexpr int GROUPS_PER_KV = 1; + const int pairs = BATCH * BB_KV_HEADS * GROUPS_PER_KV; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x + warp * gridDim.x; pair < pairs; + pair += gridDim.x * WARPS) { + const int row = pair / (BB_KV_HEADS * GROUPS_PER_KV); + MK2_PROGRESS_GUARD(progress_mask, row); + const int within_row = pair % (BB_KV_HEADS * GROUPS_PER_KV); + const int kv = within_row / GROUPS_PER_KV; + const int q_group = within_row % GROUPS_PER_KV; + const int q_first = q_group * Q_PER_WARP; + const int q_count = q_group == 0 ? Q_PER_WARP : 1; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (q_group == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[h][j] = 0.f; + } + float m[Q_PER_WARP], denom[Q_PER_WARP]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + m[h] = -CUDART_INF_F; + denom[h] = 0.f; + } + constexpr int PF = 8; + const int tpb = L.tokens_per_block; + const size_t elems_per_block = + (size_t)tpb * BB_KV_HEADS * BB_HEAD_DIM; + const int *bt_k = L.block_table + (size_t)(row * 2) * L.max_blocks_per_seq; + const int *bt_v = bt_k + L.max_blocks_per_seq; + const int head_base = kv * tpb * BB_HEAD_DIM; + float qr[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const bf16 *q = base + (kv * kv_per_q + q_first + h) * BB_HEAD_DIM; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qr[h][j] = __bfloat162float(q[lane + j * 32]); + } + for (int page0 = 0; page0 < T; page0 += tpb) { + const int k_page = bt_k[page0 / tpb] & 0x7FFFFFFF; + const int v_page = bt_v[page0 / tpb] & 0x7FFFFFFF; + const bf16 *kbase = L.pool_base + + (size_t)k_page * elems_per_block + + head_base; + const bf16 *vbase = L.pool_base + + (size_t)v_page * elems_per_block + + head_base; + const int tend = page0 + tpb < T ? page0 + tpb : T; + for (int t0 = page0; t0 < tend; t0 += PF) { + const int navail = tend - t0 < PF ? tend - t0 : PF; + float kr[PF][PER_LANE], vr[PF][PER_LANE]; + #pragma unroll + for (int u = 0; u < PF; ++u) { + const int tpos = t0 + (u < navail ? u : navail - 1); +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = vbase + (tpos - page0) * BB_HEAD_DIM; +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new : kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = cur_token + ? v_new : vbase + (tpos - page0) * BB_HEAD_DIM; +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + kr[u][j] = __bfloat162float(kt[i]); + vr[u][j] = __bfloat162float(vt[i]); + } + } + #pragma unroll + for (int u = 0; u < PF; ++u) { + if (u >= navail) break; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + dot += qr[h][j] * kr[u][j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m[h], s); + const float rescale = __expf(m[h] - m_new); + const float w = __expf(s - m_new); + denom[h] = denom[h] * rescale + w; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[h][j] = acc[h][j] * rescale + w * vr[u][j]; + m[h] = m_new; + } + } + } + } + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const float inv = 1.f / denom[h]; + const int head = kv * kv_per_q + q_first + h; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[h][j] * inv); + } + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * BB_KV_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / BB_KV_HEADS; + const int kv = idx - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv); + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv); + for (int i = 0; i < BB_HEAD_DIM; ++i) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + for (int s = 0; s <= pos; ++s) { + const bf16 *kt = kv_slot(L, seq, 0, s) + kv_offset(L, s, kv); + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + const bf16 *vt = kv_slot(L, seq, 1, s) + kv_offset(L, s, kv); + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + wt * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH, bf16 *row_stage = nullptr) { + __shared__ float score[BB_EXPERTS]; + constexpr int CH = D_MODEL / 8; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const char *gsrc = reinterpret_cast( + normed + (size_t)row * D_MODEL); + char *gdst = reinterpret_cast(row_stage); + const uint32_t act_s = (uint32_t)__cvta_generic_to_shared(gdst); + for (int i = threadIdx.x; i < CH; i += THREADS) { + const uint32_t dst = act_s + (uint32_t)i * 16; + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: + "r"(dst), "l"(gsrc + (size_t)i * 16) : "memory"); + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + if ((int)threadIdx.x < BB_EXPERTS) + moe_w[(size_t)row * BB_EXPERTS + threadIdx.x] = 0.f; + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); + __syncthreads(); + float acc0 = 0.f, acc1 = 0.f, acc2 = 0.f; + const bf16 *w0 = gate_w + (size_t)warp * D_MODEL; + const bf16 *w1 = gate_w + (size_t)(warp + WARPS) * D_MODEL; + const bf16 *w2 = gate_w + (size_t)(warp + 2 * WARPS) * D_MODEL; + for (int i = lane; i < D_MODEL; i += 128) { + bf16 sx0, sx1; + move::lds(sx0, act_s + (uint32_t)i * sizeof(bf16)); + move::lds(sx1, act_s + (uint32_t)(i + 32) * sizeof(bf16)); + float x0 = __bfloat162float(sx0); + float x1 = __bfloat162float(sx1); + acc0 += x0 * __bfloat162float(w0[i]); + acc1 += x0 * __bfloat162float(w1[i]); + acc2 += x0 * __bfloat162float(w2[i]); + acc0 += x1 * __bfloat162float(w0[i + 32]); + acc1 += x1 * __bfloat162float(w1[i + 32]); + acc2 += x1 * __bfloat162float(w2[i + 32]); + move::lds(sx0, act_s + (uint32_t)(i + 64) * sizeof(bf16)); + move::lds(sx1, act_s + (uint32_t)(i + 96) * sizeof(bf16)); + x0 = __bfloat162float(sx0); + x1 = __bfloat162float(sx1); + acc0 += x0 * __bfloat162float(w0[i + 64]); + acc1 += x0 * __bfloat162float(w1[i + 64]); + acc2 += x0 * __bfloat162float(w2[i + 64]); + acc0 += x1 * __bfloat162float(w0[i + 96]); + acc1 += x1 * __bfloat162float(w1[i + 96]); + acc2 += x1 * __bfloat162float(w2[i + 96]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + acc0 += __shfl_xor_sync(0xffffffffu, acc0, off); + acc1 += __shfl_xor_sync(0xffffffffu, acc1, off); + acc2 += __shfl_xor_sync(0xffffffffu, acc2, off); + } + if (lane == 0) { + score[warp] = 1.f / (1.f + __expf(-acc0)); + score[warp + WARPS] = 1.f / (1.f + __expf(-acc1)); + score[warp + 2 * WARPS] = 1.f / (1.f + __expf(-acc2)); + } + __syncthreads(); + if (threadIdx.x < 32) { + float remaining = (lane < BB_EXPERTS) + ? score[lane] + __bfloat162float(gate_bias[lane]) + : -CUDART_INF_F; + int pick[BB_TOP_K]; + float total = 0.f; + #pragma unroll + for (int k = 0; k < BB_TOP_K; ++k) { + float bv = remaining; + int bi = lane; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffff, bv, off); + const int oi = __shfl_xor_sync(0xffffffff, bi, off); + if (ov > bv || (ov == bv && oi < bi)) { bv = ov; bi = oi; } + } + pick[k] = bi; + total += score[bi]; + if (lane == bi) remaining = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + #pragma unroll + for (int k = 0; k < BB_TOP_K; ++k) + if (lane == pick[k]) moe_w[row * BB_EXPERTS + lane] = score[lane] * norm; + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert <= BB_EXPERTS; ++expert) + value = __fadd_rn(value, parts[base + (size_t)expert * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ bf16 bb_moe_fold_selected( + const float *parts, const float *weights, int idx) { + const int row = idx / D_MODEL; + const size_t base = (size_t)row * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert < BB_EXPERTS; ++expert) { + const float part = weights[row * BB_EXPERTS + expert] != 0.f + ? parts[base + (size_t)expert * D_MODEL] : 0.f; + value = __fadd_rn(value, part); + } + value = __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ __forceinline__ void sample_argmax_merge( + float &value, int &rank, float other, int other_rank) { + if (other > value || (other == value && other_rank < rank)) { + value = other; + rank = other_rank; + } +} +__device__ __forceinline__ int phoneme_warp_argmax(const bf16 *lg) { + const int lane = threadIdx.x & 31; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = lane; n < BB_PHONEME; n += 32) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, best, off); + const int other_arg = __shfl_down_sync(0xffffffffu, arg, off); + sample_argmax_merge(best, arg, other, other_arg); + } + return arg; +} +__device__ __forceinline__ int sample_block_argmax(float value, int rank) { + __shared__ float warp_value[WARPS]; + __shared__ int warp_rank[WARPS]; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + if (lane == 0) { warp_value[warp] = value; warp_rank[warp] = rank; } + __syncthreads(); + if (warp == 0) { + value = lane < WARPS ? warp_value[lane] : -CUDART_INF_F; + rank = lane < WARPS ? warp_rank[lane] : INT_MAX; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + } + return rank; +} +__device__ __forceinline__ int sample_eos_insertion_rank( + const uint64_t *keys, float eos) { + int lo = 0, hi = TOP_K; + while (lo < hi) { + const int mid = (lo + hi) >> 1; + if (eos > ordered_to_float((uint32_t)(keys[mid] >> 32))) hi = mid; + else lo = mid + 1; + } + return lo; +} +__device__ __forceinline__ void sample_topk_serial( + const uint64_t *keys, float eos, const float *noise, int *out) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float value; + int code; + if (!eos_used && eos > av) { + value = eos; + code = AUDIO_EOS; + eos_used = true; + } else { + value = av; + code = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = value + noise[j]; + if (noised > best_noised) { + best_noised = noised; + best_code = code; + } + } + *out = best_code; +} +__device__ __forceinline__ void sample_topk_parallel( + const uint64_t *keys, float eos, const float *noise, int *out) { + const int j = threadIdx.x; + const bool sorted_nan = j < TOP_K && + isnan(ordered_to_float((uint32_t)(keys[j] >> 32))); + if (__syncthreads_or(sorted_nan)) { + if (threadIdx.x == 0) sample_topk_serial(keys, eos, noise, out); + return; + } + const int lane = threadIdx.x & 31; + int insertion = lane == 0 ? sample_eos_insertion_rank(keys, eos) : 0; + insertion = __shfl_sync(0xffffffffu, insertion, 0); + float value = -CUDART_INF_F; + int rank = INT_MAX; + if (j < TOP_K) { + const int a = j - (j > insertion); + const float v = j == insertion + ? eos : ordered_to_float((uint32_t)(keys[a] >> 32)); + const float noised = v + noise[j]; + if (noised > -CUDART_INF_F) { value = noised; rank = j; } + } + const int winner = sample_block_argmax(value, rank); + if (threadIdx.x == 0) { + if (winner == INT_MAX) *out = 0; + else if (winner == insertion) *out = AUDIO_EOS; + else { + const int a = winner - (winner > insertion); + *out = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + } + } +} +constexpr int SORT_ITEMS_PER_THREAD = SORT_N / THREADS; +struct sampler_storage { + uint64_t keys[2][SORT_N]; +}; +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + sampler_storage &storage, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + uint64_t thread_keys[SORT_ITEMS_PER_THREAD]; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + thread_keys[item] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + int cross_stage = 0; + for (int k = 2; k <= SORT_N; k <<= 1) { + for (int j = k >> 1; j > 0; j >>= 1) { + if (j >= 32) { + const int bank = cross_stage++ & 1; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + storage.keys[bank][n] = thread_keys[item]; + } + __syncthreads(); + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + const uint64_t other = storage.keys[bank][n ^ j]; + const bool take_hi = ((n & k) == 0) == ((n & j) == 0); + const uint64_t mine = thread_keys[item]; + thread_keys[item] = take_hi + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } else { + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + const uint64_t other = __shfl_xor_sync( + 0xffffffffu, thread_keys[item], j); + const bool take_hi = ((n & k) == 0) == ((n & j) == 0); + const uint64_t mine = thread_keys[item]; + thread_keys[item] = take_hi + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } + } + } + const int top_bank = cross_stage & 1; + if (threadIdx.x < TOP_K) + storage.keys[top_bank][threadIdx.x] = thread_keys[0]; + __syncthreads(); + sample_topk_parallel( + storage.keys[top_bank], eos_v, + gumbel + ((size_t)row * N_POS + pos) * TOP_K, + out_codes + row); + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void r6_sparse_map(const float* routes,int* map,int* counts + MK2_PROGRESS_PARAMETER(progress_mask)) { + __shared__ int r7_warp_counts[THREADS/32]; + const int lane=threadIdx.x%32, warp=threadIdx.x/32; + const unsigned lower=(1u<>1; + for(int e=0;e(packed)[(size_t)e*BATCH*row_vectors+i]= + row<0?make_uint4(0,0,0,0):reinterpret_cast(source)[(size_t)row*row_vectors+i%row_vectors]; + } + } +} +__device__ void r6_tile(int work,int output_tiles,const int* lut,int &tile,int &tok, + int mtg) { + const int groups=output_tiles/mtg; + tok=lut[1+work/groups]; + const int e=tok/DECODE_TILES; + tile=e*output_tiles+(work%groups)*mtg; +} +__device__ __forceinline__ void mk2_group_tile( + int work, int grouped, int &tile, int &tok, int &here) { + here = work < grouped ? NM_MAX : 1; + const int unit = work < grouped ? work * NM_MAX : grouped * (NM_MAX - 1) + work; + tile = unit / DECODE_TILES; + tok = unit - tile * DECODE_TILES; +} +__device__ __forceinline__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&p_s)[STAGES][SLOTS], out_tile &o_s, + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + int &tile_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = DECODE_TILES, + int rows_valid = DECODE_ROWS, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const int *r6_lut = nullptr,const int *r6_map = nullptr) { + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + constexpr int QSLOT = MT + NM_MAX; + constexpr bool WIDE = MT > 1 && NM_MAX > 1 && STAGES * SLOTS >= 2 * QSLOT; + constexpr int PSTAGES = WIDE ? 2 : STAGES; + constexpr int PSLOT = WIDE ? QSLOT : SLOTS; + constexpr int ACCN = WIDE ? MT * NM_MAX : NM_MAX; + constexpr int HEX_MT = 3, HEX_NM = 2; + constexpr int HEX_SLOT = HEX_MT + HEX_NM; + auto epi_publish = [](out_tile &st, const rt_fl<16, BATCH_PAD> &src) { + using acc_t = rt_fl<16, BATCH_PAD>; + static_assert(acc_t::height == 1, "one row block per warp"); + static_assert(out_tile::rows == BATCH_PAD && out_tile::cols == OUT_ATOM, + "staging tile must be square for an in-place transpose"); + const int t = ::kittens::laneid(); + const uint32_t base = + static_cast(__cvta_generic_to_shared(&st.data[0])); + const int row0 = 8 * (t >> 4) + (t & 7); + const int col0 = 16 * warpgroup::warpid() + 8 * ((t >> 3) & 1); + #pragma unroll + for (int j = 0; j < acc_t::width; ++j) { + bf16_2 q0 = base_types::convertor::convert( + src.tiles[0][j].data[0]); + bf16_2 q1 = base_types::convertor::convert( + src.tiles[0][j].data[1]); + bf16_2 q2 = base_types::convertor::convert( + src.tiles[0][j].data[2]); + bf16_2 q3 = base_types::convertor::convert( + src.tiles[0][j].data[3]); + move::stsm4t( + out_tile::idx(base, {16 * j + row0, col0}), q0, q1, q2, q3); + } + }; + weight_tile *const pf = &p_s[0][0]; + const bool arena10 = + reinterpret_cast(&o_sf) == + reinterpret_cast(pf + STAGES * SLOTS); + const bool two_d = !r6_lut && out_tiles_per_expert == 0 + && n_tok == DECODE_TILES && DECODE_TILES % NM_MAX == 0; + int n_tiles = compact_tiles ? __popc(active_mask)*compact_tiles : n_out/OUT_ATOM; + constexpr int QTOK = DECODE_TILES / NM_MAX; + const int out_tiles = n_tiles; + const bool tall3 = WIDE && two_d && out_tiles == 24; + const bool split72 = WIDE && two_d && DECODE_TILES == 16 && out_tiles == 72; + const bool split81 = WIDE && two_d && DECODE_TILES == 16 && out_tiles == 81; + const bool split17 = WIDE && two_d && DECODE_TILES == 16 && out_tiles == 17; + const bool mixed40 = WIDE && two_d && DECODE_TILES == 16 && out_tiles == 40; + const int n_tiles_hex_all = (out_tiles / HEX_MT) * (DECODE_TILES / HEX_NM); + const bool hex = WIDE && two_d && arena10 && DECODE_TILES % HEX_NM == 0 && + (out_tiles == 96 || out_tiles == 81 || out_tiles == 72 || + out_tiles == 40) && + (epi == EPI_STORE || epi == EPI_GELU) && !r6_lut && + out_tiles_per_expert == 0; + const int hex_tok = DECODE_TILES / HEX_NM; + const int hex_wide = (out_tiles == 72 || out_tiles == 81) + ? 128 : n_tiles_hex_all; + const int pslot = hex ? HEX_SLOT : PSLOT; + int q_units = 0, q_rem_base = 0; + int grouped = 0; + if constexpr (WIDE) { + if (two_d) { + q_units = (out_tiles / MT) * QTOK; + q_rem_base = out_tiles - out_tiles % MT; + } + } else if (two_d) { + const int g = (int)gridDim.x; + const int units = n_tiles * DECODE_TILES; + const int rounds = (units + g - 1) / g; + const int cap = units / NM_MAX; + grouped = g * (rounds / NM_MAX); + if (grouped > cap) grouped = cap; + } + const int r6_output_tiles = r6_lut ? n_out/BB_EXPERTS/OUT_ATOM : 0; + const int r6_mtg = (MT > 1 && r6_lut && r6_output_tiles % MT == 0) ? MT : 1; + const int r6_groups = r6_lut ? r6_output_tiles/r6_mtg : 0; + const int r6_rows = r6_lut ? (r6_lut[0] & 0xFFF) : 0; + const int r6_pairs = (r6_lut && WIDE && r6_mtg == MT) ? (r6_lut[0] >> 12) : 0; + const int r6_pair_items = r6_pairs * r6_groups; + const int r6_single_mtg = (r6_lut && WIDE && r6_output_tiles % 3 == 0) + ? 3 : r6_mtg; + const int r6_single_groups = r6_lut ? r6_output_tiles/r6_single_mtg : 0; + const int r6_single_items = r6_lut + ? (r6_rows - 2 * r6_pairs) * r6_single_groups : 0; + if(r6_lut) { + n_tiles=r6_pair_items+r6_single_items; + n_tok=1; + } else if(two_d) { + if constexpr (WIDE) { + if (hex) n_tiles = (out_tiles == 72) ? 256 + : (out_tiles == 81) ? 264 + : (out_tiles == 40) + ? n_tiles_hex_all + hex_tok + : n_tiles_hex_all; + else if (mixed40) n_tiles = 244; + else if (split72) n_tiles = 360; + else if (split81) n_tiles = 384; + else if (split17) n_tiles = 132; + else if (tall3) n_tiles = (out_tiles / 3) * DECODE_TILES; + else n_tiles = q_units + (out_tiles % MT) * QTOK; + } + else n_tiles = grouped + (n_tiles*DECODE_TILES - grouped*NM_MAX); + n_tok=1; } + const int r6_extra = r6_lut ? n_tiles % (int)gridDim.x : 0; + const int r6_skew = r6_lut + ? (r6_single_items < r6_extra ? r6_single_items : r6_extra) : 0; + const int iters = k_total / KCHUNK; + int tile_start = (int)blockIdx.x - tile_io % (int)gridDim.x; + if (tile_start < 0) tile_start += (int)gridDim.x; + tile_io += n_tiles; + int tiles_here = 0; + for (int t = tile_start; t < n_tiles; t += gridDim.x) ++tiles_here; + const int iter_base = iter_io; + iter_io += tiles_here * iters * n_tok; + const int my_tiles = (n_tiles - tile_start + gridDim.x - 1) / gridDim.x; + if (warpgroup::groupid() == 1) { + if (hex) { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) { + if (out_tiles == 81 && work >= hex_wide) { + const int j = work - hex_wide; + const bool pair = j >= 128; + const int nt = pair ? 80 : 48 + 2 * (j / hex_tok); + const int ntok = HEX_NM * (j % hex_tok); + if (pair) { + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + wait(drained[s], (iter / PSTAGES) & 1); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], sizeof(weight_tile) + + HEX_NM * sizeof(act_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot], w, {nt, i}, filled[s]); + #pragma unroll + for (int m = 0; m < HEX_NM; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + 1 + m], a, + {ntok + m, i}, filled[s]); + } + } + } else { + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + wait(drained[s], (iter / PSTAGES) & 1); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], 2 * sizeof(weight_tile) + + HEX_NM * sizeof(act_tile)); + #pragma unroll + for (int c = 0; c < 2; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {nt + c, i}, filled[s]); + #pragma unroll + for (int m = 0; m < HEX_NM; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + 2 + m], a, + {ntok + m, i}, filled[s]); + } + } + } + continue; + } + if (out_tiles == 40 && work >= hex_wide) { + const int ntok = HEX_NM * (work - hex_wide); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + wait(drained[s], (iter / PSTAGES) & 1); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], sizeof(weight_tile) + + HEX_NM * sizeof(act_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot], w, {out_tiles - 1, i}, filled[s]); + #pragma unroll + for (int m = 0; m < HEX_NM; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + 1 + m], a, + {ntok + m, i}, filled[s]); + } + } + continue; + } + if (work >= hex_wide) { + const int j = work - hex_wide; + const int nt = HEX_MT * ((hex_wide + (j >> 1)) / hex_tok); + const int ntok = HEX_NM * ((hex_wide + (j >> 1)) % hex_tok) + (j & 1); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + wait(drained[s], (iter / PSTAGES) & 1); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], HEX_MT * sizeof(weight_tile) + + sizeof(act_tile)); + #pragma unroll + for (int c = 0; c < HEX_MT; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {nt + c, i}, filled[s]); + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + HEX_MT], a, {ntok, i}, filled[s]); + } + } + continue; + } + const int tile_id = HEX_MT * (work / hex_tok); + const int tok = HEX_NM * (work % hex_tok); + rt_fl<16, BATCH_PAD> pacc[HEX_NM]; + #pragma unroll + for (int u = 0; u < HEX_NM; ++u) warp::zero(pacc[u]); + int pend = -1, pend_ph = 0; + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(drained[s], phase); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], HEX_MT * sizeof(weight_tile) + + HEX_NM * sizeof(act_tile)); + #pragma unroll + for (int c = 0; c < HEX_MT; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {tile_id + c, i}, filled[s]); + #pragma unroll + for (int m = 0; m < HEX_NM; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + HEX_MT + m], a, {tok + m, i}, filled[s]); + } + if (pend >= 0) { { + wait(filled[pend], pend_ph); + warpgroup::mma_ABt(pacc[0], pf[pend * pslot + HEX_MT - 1], + pf[pend * pslot + HEX_MT]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + pacc[1], pf[pend * pslot + HEX_MT - 1], + pf[pend * pslot + HEX_MT + 1]); + warpgroup::mma_async_wait(); + } } + pend = s; pend_ph = phase; + } + { + wait(filled[pend], pend_ph); + warpgroup::mma_ABt(pacc[0], pf[pend * pslot + HEX_MT - 1], + pf[pend * pslot + HEX_MT]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + pacc[1], pf[pend * pslot + HEX_MT - 1], + pf[pend * pslot + HEX_MT + 1]); + warpgroup::mma_async_wait(); + } + const int es = iter % PSTAGES; + wait(drained[es], (iter / PSTAGES) & 1); + out_tile_f &q_sf = + *reinterpret_cast(&pf[es * pslot]); + out_tile &q_s = *reinterpret_cast(&q_sf); + const int qbase = (tile_id + HEX_MT - 1) * OUT_ATOM; + constexpr int OUT_OCTS = OUT_ATOM / 8; + #pragma unroll + for (int u = 0; u < HEX_NM; ++u) { + const int grow0 = (tok + u) * BATCH_PAD; + epi_publish(q_s, pacc[u]); + warpgroup::sync(3); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); + idx < OUT_OCTS * rows_valid; idx += 128) { + const int row = idx >> 3, col = (idx & 7) << 3; + const size_t at = + (size_t)(grow0 + row) * dest_stride + qbase + col; + const bf16_oct q = *reinterpret_cast( + &q_s[int2{row, col}]); + if (epi == EPI_STORE) { + *reinterpret_cast(dest + at) = q; + } else { + *reinterpret_cast(dest + at) = bf16_oct{ + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x0))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x1))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x2))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x3))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x4))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x5))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x6))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x7)))}; + } + } + warpgroup::sync(3); + } + } + } else if (warpgroup::laneid() == 0) { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) { + int tile_id = mk2_expert_tile(work, active_mask, compact_tiles), r6_tok=0; + int nm = 1, mt = 1; + if(r6_lut) { + int rwork = work - r6_skew; + if (rwork < 0) rwork += n_tiles; + const bool pair = rwork < r6_pair_items; + const int q = pair ? rwork / r6_groups + : (rwork-r6_pair_items) / r6_single_groups; + const int group = pair ? rwork-q*r6_groups + : rwork-r6_pair_items-q*r6_single_groups; + const int base_item = pair ? 2*q : 2*r6_pairs+q; + const int ltok = r6_lut[1 + base_item]; + r6_tok = ltok; + tile_id = (ltok / DECODE_TILES) * r6_output_tiles + + group * (pair ? r6_mtg : r6_single_mtg); + mt = pair ? r6_mtg : r6_single_mtg; + nm = pair ? NM_MAX : 1; + } + else if(two_d) { + if constexpr (WIDE) { + if (mixed40) { + const bool quad = work < 132; + const int item = quad ? work : work - 132; + if (quad && item < 84) { + tile_id = 2 * (item / 7); + r6_tok = 2 * (item % 7); + } else if (quad) { + const int tail = item - 84; + tile_id = 2 * (12 + tail / 6); + r6_tok = 2 * (tail % 6); + } else if (item < 48) { + tile_id = 2 * (item / 4) + (item % 4) / 2; + r6_tok = 14 + item % 2; + } else { + const int tail = item - 48; + tile_id = 2 * (12 + tail / 8) + (tail % 8) / 4; + r6_tok = 12 + tail % 4; + } + mt = quad ? MT : 1; + nm = quad ? NM_MAX : 1; + } else if (split72) { + if (work < 264) { + tile_id = 2 * (work / 8); + r6_tok = 2 * (work % 8); + mt = MT; nm = NM_MAX; + } else { + const int i = work - 264, c = i >> 2, q = i & 3; + tile_id = 2 * (33 + (c >> 3)) + (q >> 1); + r6_tok = 2 * (c & 7) + (q & 1); + mt = 1; nm = 1; + } + } else if (split81) { + if (work < 264) { + tile_id = 2 * (work / 8); + r6_tok = 2 * (work % 8); + mt = MT; nm = NM_MAX; + } else if (work < 376) { + const int i = work - 264; + tile_id = 2 * (33 + i / 16); + r6_tok = i % 16; + mt = MT; nm = 1; + } else { + tile_id = 80; + r6_tok = 2 * (work - 376); + mt = 1; nm = NM_MAX; + } + } else if (split17) { + if (work < 8) { + tile_id = 14; + r6_tok = work; + mt = 3; nm = 1; + } else if (work < 120) { + const int i = work - 8; + tile_id = i >> 3; + r6_tok = 2 * (i & 7); + mt = 1; nm = NM_MAX; + } else { + const int i = work - 120; + tile_id = 14 + (i >> 2); + r6_tok = 8 + 2 * (i & 3); + mt = 1; nm = NM_MAX; + } + } else if (tall3) { + tile_id = 3 * (work / DECODE_TILES); + r6_tok = work % DECODE_TILES; + mt = 3; + nm = 1; + } else { + const int j = work < q_units ? work : work - q_units; + tile_id = work < q_units ? MT * (j / QTOK) + : q_rem_base + j / QTOK; + r6_tok = NM_MAX * (j % QTOK); + mt = work < q_units ? MT : 1; + nm = NM_MAX; + } + } else mk2_group_tile(work,grouped,tile_id,r6_tok,nm); + } + for (int token_iteration = 0; token_iteration < n_tok; ++token_iteration) + for (int i = 0; i < iters; ++i, ++iter) { + const int tok = (r6_lut || two_d) ? r6_tok : token_iteration; + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(drained[s], phase); + tma::expect_bytes(filled[s], + mt * sizeof(weight_tile) + nm * sizeof(act_tile)); + if (mt == 1) { + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot], w, {tile_id, i}, filled[s]); + } else if (mt == 3) { + #pragma unroll + for (int c = 0; c < 3; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {tile_id + c, i}, filled[s]); + } else { + #pragma unroll + for (int c = 0; c < MT; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {tile_id + c, i}, filled[s]); + } + const int a_col = out_tiles_per_expert && !r6_lut + ? (tile_id / out_tiles_per_expert) * act_k_tiles_per_expert + i + : i; + if (NM_MAX == 1 || nm == NM_MAX) { + #pragma unroll + for (int m = 0; m < NM_MAX; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + mt + m], a, {tok + m, a_col}, filled[s]); + } else { + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + mt], a, {tok, a_col}, filled[s]); + } + } + } + } + } else { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) + for (int token_iteration = 0; token_iteration < n_tok; ++token_iteration) { + int tile_id=mk2_expert_tile(work, active_mask, compact_tiles),tok=token_iteration; + int nm = 1, mt = 1; + if(r6_lut) { + int rwork = work - r6_skew; + if (rwork < 0) rwork += n_tiles; + const bool pair = rwork < r6_pair_items; + const int q = pair ? rwork / r6_groups + : (rwork-r6_pair_items) / r6_single_groups; + const int group = pair ? rwork-q*r6_groups + : rwork-r6_pair_items-q*r6_single_groups; + const int base_item = pair ? 2*q : 2*r6_pairs+q; + const int ltok = r6_lut[1 + base_item]; + tok = ltok; + tile_id = (ltok / DECODE_TILES) * r6_output_tiles + + group * (pair ? r6_mtg : r6_single_mtg); + mt = pair ? r6_mtg : r6_single_mtg; + nm = pair ? NM_MAX : 1; + } + else if(two_d) { + if constexpr (WIDE) { + if (hex) { + if (out_tiles == 81 && work >= hex_wide) { + const int j = work - hex_wide; + const bool pair = j >= 128; + tile_id = pair ? 80 : 48 + 2 * (j / hex_tok); + tok = HEX_NM * (j % hex_tok); + mt = pair ? 1 : 2; + nm = HEX_NM; + } else if (out_tiles == 40 && work >= hex_wide) { + tile_id = out_tiles - 1; + tok = HEX_NM * (work - hex_wide); + mt = 1; + nm = HEX_NM; + } else { + const int hw = work < hex_wide + ? work : hex_wide + ((work - hex_wide) >> 1); + tile_id = HEX_MT * (hw / hex_tok); + tok = HEX_NM * (hw % hex_tok) + + (work < hex_wide ? 0 : ((work - hex_wide) & 1)); + mt = HEX_MT; nm = work < hex_wide ? HEX_NM : 1; + } + } else if (mixed40) { + const bool quad = work < 132; + const int item = quad ? work : work - 132; + if (quad && item < 84) { + tile_id = 2 * (item / 7); + tok = 2 * (item % 7); + } else if (quad) { + const int tail = item - 84; + tile_id = 2 * (12 + tail / 6); + tok = 2 * (tail % 6); + } else if (item < 48) { + tile_id = 2 * (item / 4) + (item % 4) / 2; + tok = 14 + item % 2; + } else { + const int tail = item - 48; + tile_id = 2 * (12 + tail / 8) + (tail % 8) / 4; + tok = 12 + tail % 4; + } + mt = quad ? MT : 1; + nm = quad ? NM_MAX : 1; + } else if (split72) { + if (work < 264) { + tile_id = 2 * (work / 8); + tok = 2 * (work % 8); + mt = MT; nm = NM_MAX; + } else { + const int i = work - 264, c = i >> 2, q = i & 3; + tile_id = 2 * (33 + (c >> 3)) + (q >> 1); + tok = 2 * (c & 7) + (q & 1); + mt = 1; nm = 1; + } + } else if (split81) { + if (work < 264) { + tile_id = 2 * (work / 8); + tok = 2 * (work % 8); + mt = MT; nm = NM_MAX; + } else if (work < 376) { + const int i = work - 264; + tile_id = 2 * (33 + i / 16); + tok = i % 16; + mt = MT; nm = 1; + } else { + tile_id = 80; + tok = 2 * (work - 376); + mt = 1; nm = NM_MAX; + } + } else if (split17) { + if (work < 8) { + tile_id = 14; + tok = work; + mt = 3; nm = 1; + } else if (work < 120) { + const int i = work - 8; + tile_id = i >> 3; + tok = 2 * (i & 7); + mt = 1; nm = NM_MAX; + } else { + const int i = work - 120; + tile_id = 14 + (i >> 2); + tok = 8 + 2 * (i & 3); + mt = 1; nm = NM_MAX; + } + } else if (tall3) { + tile_id = 3 * (work / DECODE_TILES); + tok = work % DECODE_TILES; + mt = 3; + nm = 1; + } else { + const int j = work < q_units ? work : work - q_units; + tile_id = work < q_units ? MT * (j / QTOK) + : q_rem_base + j / QTOK; + tok = NM_MAX * (j % QTOK); + mt = work < q_units ? MT : 1; + nm = NM_MAX; + } + } else mk2_group_tile(work,grouped,tile_id,tok,nm); + } + int epi_hold = -1; + rt_fl<16, BATCH_PAD> accum[ACCN]; + if (hex && nm == 1) { + #pragma unroll + for (int u = 0; u < 3; ++u) warp::zero(accum[u]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + HEX_MT]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[1], pf[s * pslot + 1], pf[s * pslot + HEX_MT]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[2], pf[s * pslot + 2], pf[s * pslot + HEX_MT]); + warpgroup::mma_async_wait(); + if (i + 1 < iters) { + if (warpgroup::laneid() == 0) + arrive(drained[s], CONSUMER_WARPS); + } else epi_hold = s; + } + } else if (hex && mt == HEX_MT) { + #pragma unroll + for (int u = 0; u < ACCN; ++u) warp::zero(accum[u]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], + pf[s * pslot + HEX_MT]); + #pragma unroll + for (int u = 1; u < ACCN; ++u) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[u], pf[s * pslot + u / HEX_NM], + pf[s * pslot + HEX_MT + u % HEX_NM]); + warpgroup::mma_async_wait(); + if (i + 1 < iters) { + if (warpgroup::laneid() == 0) + arrive(drained[s], CONSUMER_WARPS); + } else epi_hold = s; + } + } else if (WIDE && mt == 3) { + #pragma unroll + for (int u = 0; u < 3; ++u) warp::zero(accum[u]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + 3]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[1], pf[s * pslot + 1], pf[s * pslot + 3]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[2], pf[s * pslot + 2], pf[s * pslot + 3]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else if (WIDE && mt == MT && nm == NM_MAX) { + #pragma unroll + for (int u = 0; u < ACCN; ++u) warp::zero(accum[u]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + MT]); + #pragma unroll + for (int u = 1; u < ACCN; ++u) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[u], pf[s * pslot + u / NM_MAX], + pf[s * pslot + MT + u % NM_MAX]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else if (MT > 1 && mt == MT) { + #pragma unroll + for (int c = 0; c < MT; ++c) warp::zero(accum[c]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + MT]); + #pragma unroll + for (int c = 1; c < MT; ++c) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[c], pf[s * pslot + c], pf[s * pslot + MT]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else if (NM_MAX == 1 || nm == NM_MAX) { + #pragma unroll + for (int m = 0; m < NM_MAX; ++m) warp::zero(accum[m]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + 1]); + #pragma unroll + for (int m = 1; m < NM_MAX; ++m) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[m], pf[s * pslot], pf[s * pslot + 1 + m]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else { + warp::zero(accum[0]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + 1]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } + out_tile_f &e_sf = epi_hold >= 0 + ? *reinterpret_cast(&pf[epi_hold * pslot]) + : o_sf; + out_tile *const e_sp = reinterpret_cast(&e_sf); + const int nacc = mt * nm; + if (epi != EPI_BIAS) { + epi_publish(e_sp[0], accum[0]); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + } + #pragma unroll + for (int u = 0; u < ACCN; ++u) { + if (u >= mt * nm) break; + const int c = (mt > 1) ? ((nm > 1) ? (u / NM_MAX) : u) : 0; + const int m = (nm > 1) ? (u % NM_MAX) : 0; + const int base = (r6_lut ? (tile_id + c)%r6_output_tiles + : (tile_id + c)) * OUT_ATOM; + const int grow0 = (tok + m) * BATCH_PAD; + if (epi == EPI_BIAS) { + const int bw = 16 * warpgroup::warpid() + + (::kittens::laneid() >> 2); + const float b0 = __bfloat162float(bias[base + bw]); + const float b1 = __bfloat162float(bias[base + bw + 8]); + #pragma unroll + for (int j = 0; j < rt_fl<16, BATCH_PAD>::width; ++j) { + accum[u].tiles[0][j].data[0].x += b0; + accum[u].tiles[0][j].data[0].y += b0; + accum[u].tiles[0][j].data[1].x += b1; + accum[u].tiles[0][j].data[1].y += b1; + accum[u].tiles[0][j].data[2].x += b0; + accum[u].tiles[0][j].data[2].y += b0; + accum[u].tiles[0][j].data[3].x += b1; + accum[u].tiles[0][j].data[3].y += b1; + } + epi_publish(e_sp[0], accum[u]); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + constexpr int OUT_OCTS = OUT_ATOM / 8; + for (int idx = warpgroup::laneid(); idx < OUT_OCTS * rows_valid; + idx += 128) { + const int row = idx >> 3, col = (idx & 7) << 3; + const size_t at = (size_t)(grow0 + row) * dest_stride + base + col; + *reinterpret_cast(dest + at) = + *reinterpret_cast( + &e_sp[0][int2{row, col}]); + } + warpgroup::sync(1); + continue; + } + if (u + 1 < ACCN && u + 1 < nacc) + epi_publish(e_sp[(u + 1) & 1], + accum[u + 1 < ACCN ? u + 1 : 0]); + const out_tile &e_sr = e_sp[u & 1]; + constexpr int OUT_OCTS = OUT_ATOM / 8; + for (int idx = warpgroup::laneid(); idx < OUT_OCTS * rows_valid; + idx += 128) { + const int row = idx >> 3; + const int col = (idx & 7) << 3; + const int grow = grow0 + row; + const size_t at = + (size_t)grow * dest_stride + base + col; + const bf16_oct src = *reinterpret_cast( + &e_sr[int2{row, col}]); + if (epi == EPI_STORE) { + *reinterpret_cast(dest + at) = src; + continue; + } + const float v0 = __bfloat162float(src.x0); + const float v1 = __bfloat162float(src.x1); + const float v2 = __bfloat162float(src.x2); + const float v3 = __bfloat162float(src.x3); + const float v4 = __bfloat162float(src.x4); + const float v5 = __bfloat162float(src.x5); + const float v6 = __bfloat162float(src.x6); + const float v7 = __bfloat162float(src.x7); + switch (epi) { + case EPI_STORE: break; + case EPI_SILU: + *reinterpret_cast(dest + at) = bf16_oct{ + __float2bfloat16(v0 / (1.f + __expf(-v0))), + __float2bfloat16(v1 / (1.f + __expf(-v1))), + __float2bfloat16(v2 / (1.f + __expf(-v2))), + __float2bfloat16(v3 / (1.f + __expf(-v3))), + __float2bfloat16(v4 / (1.f + __expf(-v4))), + __float2bfloat16(v5 / (1.f + __expf(-v5))), + __float2bfloat16(v6 / (1.f + __expf(-v6))), + __float2bfloat16(v7 / (1.f + __expf(-v7)))}; + break; + case EPI_MOE_ACC: { + const int expert = (tile_id + c) / out_tiles_per_expert; + const int local = ((tile_id + c) % out_tiles_per_expert) * OUT_ATOM + + col; + const int original_row=r6_map?r6_map[grow]:grow; + if(original_row<0) break; + const float wt = moe_w[original_row * BB_EXPERTS + expert]; + if (wt != 0.f) { + const size_t dst = + ((size_t)original_row * (BB_EXPERTS + 1) + expert) * + D_MODEL + local; + *reinterpret_cast(moe_out + dst) = make_float4( + __fmul_rn(wt, v0), __fmul_rn(wt, v1), + __fmul_rn(wt, v2), __fmul_rn(wt, v3)); + *reinterpret_cast(moe_out + dst + 4) = make_float4( + __fmul_rn(wt, v4), __fmul_rn(wt, v5), + __fmul_rn(wt, v6), __fmul_rn(wt, v7)); + } + break; + } + case EPI_ACC_F32: + *reinterpret_cast(moe_out + + ((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * + D_MODEL + base + col) = make_float4(v0, v1, v2, v3); + *reinterpret_cast(moe_out + + ((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * + D_MODEL + base + col + 4) = make_float4(v4, v5, v6, v7); + break; + case EPI_BIAS: break; + case EPI_GELU: + *reinterpret_cast(dest + at) = bf16_oct{ + __float2bfloat16(gelu_tanh_fp32(v0)), + __float2bfloat16(gelu_tanh_fp32(v1)), + __float2bfloat16(gelu_tanh_fp32(v2)), + __float2bfloat16(gelu_tanh_fp32(v3)), + __float2bfloat16(gelu_tanh_fp32(v4)), + __float2bfloat16(gelu_tanh_fp32(v5)), + __float2bfloat16(gelu_tanh_fp32(v6)), + __float2bfloat16(gelu_tanh_fp32(v7))}; + break; + case EPI_RESIDUAL: { + const bf16_oct prior = + *reinterpret_cast(residual + at); + *reinterpret_cast(residual + at) = bf16_oct{ + __float2bfloat16(__bfloat162float(prior.x0) + v0), + __float2bfloat16(__bfloat162float(prior.x1) + v1), + __float2bfloat16(__bfloat162float(prior.x2) + v2), + __float2bfloat16(__bfloat162float(prior.x3) + v3), + __float2bfloat16(__bfloat162float(prior.x4) + v4), + __float2bfloat16(__bfloat162float(prior.x5) + v5), + __float2bfloat16(__bfloat162float(prior.x6) + v6), + __float2bfloat16(__bfloat162float(prior.x7) + v7)}; + break; + } + } + } + if (u + 1 < nacc) { + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + } + } + if (epi != EPI_BIAS) warpgroup::sync(1); + if (epi_hold >= 0 && warpgroup::laneid() == 0) + arrive(drained[epi_hold], CONSUMER_WARPS); + } + (void)my_tiles; + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr, + const unsigned char *native_predictor_active_mask = nullptr + MK2_PROGRESS_PARAMETER(native_progress_mask)) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&p_s)[STAGES][SLOTS] = alloc.allocate(); + out_tile_f (&o_sf) = alloc.allocate(); + out_tile &o_s = *reinterpret_cast(&o_sf); + bf16 *const route_stage = reinterpret_cast(&o_sf); + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_SCRATCH_STRIDE]; + static_assert(STAGES * SLOTS * sizeof(weight_tile) >= sizeof(sampler_storage)); + sampler_storage &sort_storage = *reinterpret_cast(&__shm[0]); + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS, 0); + arrive(drained[s], CONSUMER_WARPS); + } + } + int gemm_iter = 0; + int gemm_tile = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + const float *fold_src = nullptr, *fold_routes = nullptr; + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + fold_src = nullptr; + fold_routes = nullptr; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + grid_barrier(g.barrier); +#endif + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); +#endif + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, BATCH, + route_stage); + } + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + } + grid_barrier(g.barrier); + r6_sparse_map(bb->moe_w,bb->r6_map,bb->r6_counts + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + r6_sparse_gather(bb->normed,bb->r6_packed,bb->r6_map,bb->r6_counts, + bb->r6_lut); + grid_barrier(g.barrier); + gemm(L.exp_up, bb->r6_views[0], BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0,0,nullptr,nullptr,DECODE_TILES,DECODE_ROWS, + nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->r6_views[1], BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + DECODE_TILES, DECODE_ROWS, nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + if (bb->layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_selected(bb->moe_out, bb->moe_w, idx); + grid_barrier(g.barrier); + } else { + fold_src = bb->moe_out; + fold_routes = bb->moe_w; + } + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int n_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = n_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + bb->pf_h[idx] = bb->prompt[idx]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = bb->pf_n_seq > 0 ? bb->pf_n_seq : 1; + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_gated, + T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, T, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * (BB_EXPERTS + 1) * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->pf_normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->pf_moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = bb_moe_fold_deterministic(bb->pf_moe_out, idx); + grid_barrier(g.barrier); + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, bb->pf_final, ln_scratch, + n_total); + grid_barrier(g.barrier); + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, bb->phoneme_b); + grid_barrier(g.barrier); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + for (int row = blockIdx.x + warp * gridDim.x; row < BATCH; + row += gridDim.x * WARPS) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + MK2_PROGRESS_GUARD(native_progress_mask, row); + const int arg = phoneme_warp_argmax(lg); + if (lane == 0) bb->phoneme_out[row] = arg; + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + __shared__ int native_run_predictor; + if (native_predictor_active_mask && bb && g.native_codes) { + if (threadIdx.x == 0) { + int any_active = 0; + for (int row = 0; row < BATCH; ++row) + any_active |= int(native_predictor_active_mask[row] != 0 && + MK2_PROGRESS_ENABLED(native_progress_mask, row)); + native_run_predictor = any_active; + } + __syncthreads(); + } + const bool skip_native_predictor = native_predictor_active_mask && bb && + g.native_codes && native_run_predictor == 0; + if (skip_native_predictor) { + for (int index = blockIdx.x * THREADS + threadIdx.x; + index < BATCH * N_POS; index += gridDim.x * THREADS) { + MK2_PROGRESS_GUARD(native_progress_mask, index / N_POS); + g.sampled[index] = 0; + g.native_codes[index] = 0; + } + grid_barrier(g.barrier); + } else { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + constexpr int DPL = D_HEAD / 32; + float qv[DPL], kn[DPL], vn[DPL]; + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + qv[j] = __bfloat162float(q[i]); + kn[j] = __bfloat162float(k_new[i]); + vn[j] = __bfloat162float(v_new[i]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += qv[j] * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + { + float acc = 0.f; + #pragma unroll + for (int j = 0; j < DPL; ++j) acc += qv[j] * kn[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[pos] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + float ov[DPL] = {}; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *vt = g.vcache + cache_base + t * D_HEAD; + const float s = sc[t]; + #pragma unroll + for (int j = 0; j < DPL; ++j) + ov[j] += s * __bfloat162float(vt[lane + j * 32]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + ov[j] += sc[pos] * vn[j]; + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + + lane + j * 32] = __float2bfloat16(ov[j]); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, + g.sampled_pos, sort_storage, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif +#ifdef ABLATE_SAMPLE + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { + __syncthreads(); + } else { + grid_barrier(g.barrier); + } +#endif + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x == 0) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + g.sampled[row * N_POS + pos] = code; + if (g.native_codes && MK2_PROGRESS_ENABLED(native_progress_mask, row)) { +#if MK2_INFLIGHT + g.native_codes[row * N_POS + pos] = + native_predictor_active_mask && !native_predictor_active_mask[row] + ? 0 : (long long)code; +#else + g.native_codes[row * N_POS + pos] = (long long)code; +#endif + } + } + } +#ifdef ABLATE_SAMPLE + grid_barrier(g.barrier); +#else + if (pos + 1 == n_pos) grid_barrier(g.barrier); +#endif + } + } + if (g.native_seq_len && blockIdx.x == 0) + for (int row = threadIdx.x; row < BATCH; row += blockDim.x) { + MK2_PROGRESS_GUARD(native_progress_mask, row); + g.native_seq_len[row] += 1; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2_startup( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr, + const unsigned char *native_predictor_active_mask = nullptr + MK2_PROGRESS_PARAMETER(native_progress_mask)) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&p_s)[STAGES][SLOTS] = alloc.allocate(); + out_tile_f (&o_sf) = alloc.allocate(); + out_tile &o_s = *reinterpret_cast(&o_sf); + bf16 *const route_stage = reinterpret_cast(&o_sf); + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_SCRATCH_STRIDE]; + static_assert(STAGES * SLOTS * sizeof(weight_tile) >= sizeof(sampler_storage)); + sampler_storage &sort_storage = *reinterpret_cast(&__shm[0]); + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS, 0); + arrive(drained[s], CONSUMER_WARPS); + } + } + int gemm_iter = 0; + int gemm_tile = 0; + __syncthreads(); + constexpr bool prefilling = false; + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + const float *fold_src = nullptr, *fold_routes = nullptr; + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + fold_src = nullptr; + fold_routes = nullptr; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + grid_barrier(g.barrier); +#endif + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); +#endif + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, BATCH, + route_stage); + } + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + } + grid_barrier(g.barrier); + r6_sparse_map(bb->moe_w,bb->r6_map,bb->r6_counts + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + r6_sparse_gather(bb->normed,bb->r6_packed,bb->r6_map,bb->r6_counts, + bb->r6_lut); + grid_barrier(g.barrier); + gemm(L.exp_up, bb->r6_views[0], BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0,0,nullptr,nullptr,DECODE_TILES,DECODE_ROWS, + nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->r6_views[1], BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + DECODE_TILES, DECODE_ROWS, nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + if (bb->layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_selected(bb->moe_out, bb->moe_w, idx); + grid_barrier(g.barrier); + } else { + fold_src = bb->moe_out; + fold_routes = bb->moe_w; + } + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if constexpr (false) { + const int n_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = n_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + bb->pf_h[idx] = bb->prompt[idx]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = bb->pf_n_seq > 0 ? bb->pf_n_seq : 1; + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_gated, + T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, T, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * (BB_EXPERTS + 1) * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->pf_normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->pf_moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = bb_moe_fold_deterministic(bb->pf_moe_out, idx); + grid_barrier(g.barrier); + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, bb->pf_final, ln_scratch, + n_total); + grid_barrier(g.barrier); + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, bb->phoneme_b); + grid_barrier(g.barrier); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + for (int row = blockIdx.x + warp * gridDim.x; row < BATCH; + row += gridDim.x * WARPS) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + MK2_PROGRESS_GUARD(native_progress_mask, row); + const int arg = phoneme_warp_argmax(lg); + if (lane == 0) bb->phoneme_out[row] = arg; + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + __shared__ int native_run_predictor; + if (native_predictor_active_mask && bb && g.native_codes) { + if (threadIdx.x == 0) { + int any_active = 0; + for (int row = 0; row < BATCH; ++row) + any_active |= int(native_predictor_active_mask[row] != 0 && + MK2_PROGRESS_ENABLED(native_progress_mask, row)); + native_run_predictor = any_active; + } + __syncthreads(); + } + const bool skip_native_predictor = native_predictor_active_mask && bb && + g.native_codes && native_run_predictor == 0; + if constexpr (true) { + for (int index = blockIdx.x * THREADS + threadIdx.x; + index < BATCH * N_POS; index += gridDim.x * THREADS) { + MK2_PROGRESS_GUARD(native_progress_mask, index / N_POS); + g.sampled[index] = 0; + g.native_codes[index] = 0; + } + grid_barrier(g.barrier); + } else { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + constexpr int DPL = D_HEAD / 32; + float qv[DPL], kn[DPL], vn[DPL]; + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + qv[j] = __bfloat162float(q[i]); + kn[j] = __bfloat162float(k_new[i]); + vn[j] = __bfloat162float(v_new[i]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += qv[j] * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + { + float acc = 0.f; + #pragma unroll + for (int j = 0; j < DPL; ++j) acc += qv[j] * kn[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[pos] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + float ov[DPL] = {}; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *vt = g.vcache + cache_base + t * D_HEAD; + const float s = sc[t]; + #pragma unroll + for (int j = 0; j < DPL; ++j) + ov[j] += s * __bfloat162float(vt[lane + j * 32]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + ov[j] += sc[pos] * vn[j]; + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + + lane + j * 32] = __float2bfloat16(ov[j]); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, + g.sampled_pos, sort_storage, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif +#ifdef ABLATE_SAMPLE + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { + __syncthreads(); + } else { + grid_barrier(g.barrier); + } +#endif + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x == 0) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + g.sampled[row * N_POS + pos] = code; + if (g.native_codes && MK2_PROGRESS_ENABLED(native_progress_mask, row)) { +#if MK2_INFLIGHT + g.native_codes[row * N_POS + pos] = + native_predictor_active_mask && !native_predictor_active_mask[row] + ? 0 : (long long)code; +#else + g.native_codes[row * N_POS + pos] = (long long)code; +#endif + } + } + } +#ifdef ABLATE_SAMPLE + grid_barrier(g.barrier); +#else + if (pos + 1 == n_pos) grid_barrier(g.barrier); +#endif + } + } + if (g.native_seq_len && blockIdx.x == 0) + for (int row = threadIdx.x; row < BATCH; row += blockDim.x) { + MK2_PROGRESS_GUARD(native_progress_mask, row); + g.native_seq_len[row] += 1; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 64 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * SLOTS * sizeof(weight_tile) + + sizeof(out_tile_f) + 2048; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; + const unsigned char *native_predictor_active_mask = nullptr; + MK2_PROGRESS_MEMBER + int maximum_grid = 0; + bool grid_reset_pending = false; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * SLOTS * sizeof(weight_tile) + + sizeof(out_tile_f) + 2048), 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaFuncSetAttribute(predictor_mk2_startup, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = std::min(out->sms * out->blocks_per_sm, 4096); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 64 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->maximum_grid = info.sms * info.blocks_per_sm; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + mk2_check(owned.alloc(&hb.r6_packed,(size_t)BB_EXPERTS*BATCH*D_MODEL*sizeof(bf16))); + mk2_check(owned.alloc(&hb.r6_map,(size_t)BB_EXPERTS*BATCH*sizeof(int))); + mk2_check(owned.alloc(&hb.r6_counts,BB_EXPERTS*sizeof(int))); + mk2_check(owned.alloc(&hb.r6_lut,(1+BB_EXPERTS*DECODE_TILES)*sizeof(int))); + act_gl sparse_views[]={ + act_gl{hb.r6_packed,nullptr,nullptr,(size_t)BB_EXPERTS*BATCH,(size_t)D_MODEL}, + act_gl{moe_h,nullptr,nullptr,(size_t)BB_EXPERTS*BATCH,(size_t)BB_MOE_INTER}}; + mk2_check(owned.alloc(&hb.r6_views,sizeof(sparse_views))); + mk2_check(cudaMemcpy(hb.r6_views,sparse_views,sizeof(sparse_views),cudaMemcpyHostToDevice)); + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.maximum_grid = pred->maximum_grid; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native, bool startup = false) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(s, &capture_status)); + if (st->pred.grid_reset_pending || capture_status != cudaStreamCaptureStatusNone) { + mk2_check(cudaMemsetAsync(st->pred.g.barrier, 0, 64 * sizeof(unsigned int), s)); + st->pred.grid_reset_pending = false; + } + if (startup) { + predictor_mk2_startup<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS, nullptr, + st->pred.native_predictor_active_mask + MK2_PROGRESS_ARGUMENT(st->pred.native_progress_mask)); + } else { + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS, nullptr, + native ? st->pred.native_predictor_active_mask : nullptr + MK2_PROGRESS_ARGUMENT(native ? st->pred.native_progress_mask : nullptr)); + } + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" int mk2_run_native_startup_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream) { + return mk2_enqueue_stage0(p, b, stream, true, true); +} +extern "C" int mk2_set_native_predictor_active_mask(MK2Predictor *p, + const unsigned char *mask) { + try { + if (!p || reinterpret_cast(p)->kind != 1) + throw std::runtime_error("startup mask requires a full Stage0 handle"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && st->pred.native_predictor_active_mask != mask) + throw std::runtime_error("startup mask pointer cannot change after first Stage0 enqueue"); + if (mask) + mk2_stage0_device_pointer(mask, st->pred.device, "native_predictor_active_mask"); + st->pred.native_predictor_active_mask = mask; + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown startup mask binding failure"; return 1; } +} +extern "C" int mk2_native_startup_skip_abi_version() { return 1; } +#include "../../../../../inflight/runtime.cuh" +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "native_prefill_impl.inc" +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *p, int grid) { + if (!p) { g_last_error = "grid setter requires full Stage0 handle"; return 1; } + auto *st = reinterpret_cast(p); + if (st->kind != 1 || grid < 1 || grid > st->maximum_grid) { + g_last_error = "requested Stage0 grid outside independently resident capacity"; + return 1; + } + if (st->grid != grid) st->grid_reset_pending = true; + st->grid = grid; + g_last_error.clear(); + return 0; +} +extern "C" int mk2_r8_bind_initial_state(MK2Predictor *p, const void *const *initial, const void *ready, const void *seq_len, int positions) { + try { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || st->bindings_initialized || !initial || !ready || !seq_len || positions <= 0) + throw std::runtime_error("invalid immutable initial-state binding"); + for (int i = 0; i < 14; ++i) { + mk2_stage0_device_pointer(initial[i], st->pred.device, "initial state"); + st->mb_host[i].r8_initial_ssm = static_cast(initial[i]); + st->mb_host[i].r8_initial_ready = static_cast(ready); + st->mb_host[i].r8_seq_len = static_cast(seq_len); + st->mb_host[i].r8_initial_positions = positions; + } + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" int mk2_r9_state_bind(void* p,const void*const* scales,const void*const* maps) {auto*st=(MK2Stage0State*)p;if(st->bindings_initialized)return 1;for(int i=0;i<14;++i){st->mb_host[i].r9_scales=(const float*)scales[i];st->mb_host[i].r9_map=(const int*)maps[i];}return 0;} +extern "C" int mk2_r9_state_bytes(){return sizeof(r9_state_t);} diff --git a/deploy/nemotron_native/kernels/b1024/regions/stage0/stage0.cu b/deploy/nemotron_native/kernels/b1024/regions/stage0/stage0.cu new file mode 100644 index 0000000..37e02b0 --- /dev/null +++ b/deploy/nemotron_native/kernels/b1024/regions/stage0/stage0.cu @@ -0,0 +1,35 @@ +#define MK2_BATCH 1024 +#define STAGES 3 +#define KCHUNK 128 +#define MK2_NO_MAIN 1 +#define MK2_DIRECT_INPUT 1 +#define MK2_B1_LOCAL_BARRIERS 1 +#define MK2_SKIP_HIDDEN_SNAPSHOT 1 +#define MK2_ACTIVE_EXPERTS 1 +#define MK2_SSM_SPLITS 1 +#define NMULT 2 +#define MK2_OUT_MULT 2 +#define SSM_PREFETCH_MAX 12 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#if MK2_INFLIGHT +#include +#endif +#include +#include +#include +#include +#include +#include +#include +#include +namespace b1024_stage0 { +#include "native/stage0/predictor_mk2.cu" +} diff --git a/deploy/nemotron_native/kernels/b1024/stage0.cu b/deploy/nemotron_native/kernels/b1024/stage0.cu new file mode 100644 index 0000000..6c04cd7 --- /dev/null +++ b/deploy/nemotron_native/kernels/b1024/stage0.cu @@ -0,0 +1 @@ +#include "regions/stage0/stage0.cu" diff --git a/deploy/nemotron_native/kernels/b128/prefill.cu b/deploy/nemotron_native/kernels/b128/prefill.cu new file mode 100644 index 0000000..6d21d7e --- /dev/null +++ b/deploy/nemotron_native/kernels/b128/prefill.cu @@ -0,0 +1 @@ +#include "regions/prefill/prefill.cu" diff --git a/deploy/nemotron_native/kernels/b128/regions/prefill/native/prefill/dedup_native_prefill_impl.inc b/deploy/nemotron_native/kernels/b128/regions/prefill/native/prefill/dedup_native_prefill_impl.inc new file mode 100644 index 0000000..fe4c125 --- /dev/null +++ b/deploy/nemotron_native/kernels/b128/regions/prefill/native/prefill/dedup_native_prefill_impl.inc @@ -0,0 +1,267 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 16 || BATCH == 32 || BATCH == 64 || BATCH == 128 || BATCH == 256 || BATCH == 512 || BATCH == 1024 || BATCH == 2048)) + throw std::runtime_error("native prefill requires a matching supported physical batch"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((std::uint64_t(c.prompt_positions) + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (rows > (INT_MAX - 132 * THREADS) / D_MODEL) + throw std::runtime_error("native prefill full-width prompt exceeds safe device loop bounds"); + if (padded > (INT_MAX - 132 * THREADS) / (BB_EXPERTS * BB_MOE_INTER)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((size_t(config->prompt_positions) + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + float *r9_workspace = nullptr; + mk2_check(owned.alloc(&r9_workspace,size_t(BATCH)*ssm_elements*sizeof(float))); + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + r9_state_t *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(r9_state_t))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(r9_state_t))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = r9_workspace; + st->mb_host[i].r9_private = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + mk2_check(owned.alloc(&st->mb_host[i].r8_canonical_ssm, ssm_elements * sizeof(float))); + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_dtda, 2 * MB_HEADS); + { + const size_t estride = padded; + const size_t grows = size_t(BB_EXPERTS) * estride; + mk2_check(owned.alloc(&hb.pf_moe_gath, grows * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_gath, 0, grows * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_hc, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_hc, 0, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_gidx, grows * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_gidx, 0, grows * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ecnt, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ecnt, 0, BB_EXPERTS * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ntiles, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ntiles, 0, BB_EXPERTS * sizeof(int))); + hb.pf_moe_estride = int(estride); + hb.pf_moe_gath_v = act_gl{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_v = act_gl{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + hb.pf_moe_gath_vn = act_gl_n{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_vn = act_gl_n{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + scratch_bytes += grows * (D_MODEL + BB_MOE_INTER) * sizeof(bf16); + } + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(owned.alloc(&hb.native_pf_cohort_check, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(hb.native_pf_cohort_check, 0, 2 * sizeof(unsigned int))); +#if PF_ACTIVE_TOKEN_EXPERTS + mk2_check(owned.alloc(&hb.native_pf_expert_masks, (padded/BATCH_PAD)*sizeof(unsigned int))); +#endif + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} +extern "C" int mk2_native_prefill_validate(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null compact prefill handle"); + auto *st = reinterpret_cast(p->base); + mk2_check(cudaStreamSynchronize(static_cast(stream))); + unsigned int status[2] = {0, 0}; + mk2_check(cudaMemcpy(status, st->host_bb.native_pf_cohort_check, + sizeof(status), cudaMemcpyDeviceToHost)); + if (status[0] != 0 || status[1] != 1) + throw std::runtime_error("compact prefill requires bit-identical prompts and identity slots"); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown compact prefill validation failure"; return 1; } +} +extern "C" int mk2_native_prefill_workspace_sequences() { return 1; } +extern "C" int mk2_native_prefill_effective_sequences(MK2NativePrefill *p, int *out) { + try { + if (!p || !out) throw std::runtime_error("null effective-sequence query"); + auto *st = reinterpret_cast(p->base); + mk2_check(cudaMemcpy(out, st->host_bb.native_pf_cohort_check + 1, + sizeof(int), cudaMemcpyDeviceToHost)); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" int mk2_r8_initial_state(MK2NativePrefill *p, void **initial, void **ready) { + if (!p || !initial || !ready) return 1; + auto *st = reinterpret_cast(p->base); + for (int i = 0; i < 14; ++i) initial[i] = st->mb_host[i].r8_canonical_ssm; + *ready = st->host_bb.native_pf_cohort_check + 1; + return 0; +} +extern "C" int mk2_r9_prefill_bind(MK2NativePrefill*p,const void*const*scales,const void*const*maps){if(!p)return 1;auto*st=(MK2Stage0State*)p->base;for(int i=0;i<14;++i){st->mb_host[i].r9_scales=(const float*)scales[i];st->mb_host[i].r9_map=(const int*)maps[i];}return (int)cudaMemcpy((void*)st->host_bb.mamba,st->mb_host.data(),14*sizeof(MambaLayer),cudaMemcpyHostToDevice);} diff --git a/deploy/nemotron_native/kernels/b128/regions/prefill/native/prefill/dedup_predictor_mk2.cuh b/deploy/nemotron_native/kernels/b128/regions/prefill/native/prefill/dedup_predictor_mk2.cuh new file mode 100644 index 0000000..f98dd53 --- /dev/null +++ b/deploy/nemotron_native/kernels/b128/regions/prefill/native/prefill/dedup_predictor_mk2.cuh @@ -0,0 +1,4600 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef PF_ACTIVE_TOKEN_EXPERTS +#define PF_ACTIVE_TOKEN_EXPERTS 0 +#endif +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +#ifndef MK2_PREFILL_ONLY +#define MK2_PREFILL_ONLY 0 +#endif +#ifndef MK2_BATCH_TILE +#define MK2_BATCH_TILE (((MK2_BATCH + 15) / 16) * 16) +#endif +constexpr int BATCH_PAD = MK2_BATCH_TILE; +static_assert(BATCH >= 1 && (MK2_PREFILL_ONLY || BATCH <= BATCH_PAD)); +static_assert(BATCH_PAD == 16 || BATCH_PAD == 32 || BATCH_PAD == 64 || BATCH_PAD == 128, + "native prefill tile must be 16, 32, 64 or 128"); +static_assert(MK2_PREFILL_ONLY || BATCH_PAD == ((BATCH + 15) / 16) * 16, + "a wider token tile requires a prefill-only library"); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_PW_MULTI +#define PF_SSM_PW_MULTI 4 +#endif +#ifndef PF_SSM_PW_SINGLE +#define PF_SSM_PW_SINGLE 2 +#endif +static_assert(MB_HEAD_DIM % PF_SSM_PW_MULTI == 0 && + MB_HEAD_DIM % PF_SSM_PW_SINGLE == 0, + "head_dim must split evenly into scan channel groups"); +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef MK2_SMEM_SLACK +#define MK2_SMEM_SLACK 1024 +#endif +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +#ifndef MK2_THREADS +#define MK2_THREADS 256 +#endif +constexpr int THREADS = MK2_THREADS; +constexpr int NORM_THREADS = 256; +constexpr int NORM_WARPS = NORM_THREADS / 32; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +#ifndef MK2_GEMM_CONSUMERS +#define MK2_GEMM_CONSUMERS 1 +#endif +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 +#define MK2_HEAD_STAGE_BYTES 0 +#else +#define MK2_HEAD_STAGE_BYTES sizeof(out_tile_f) +#endif +constexpr int NCONS = MK2_GEMM_CONSUMERS; +static_assert(NCONS >= 1 && NCONS <= 3, "consumers occupy warpgroups 0, 2 and 3"); +static_assert(THREADS >= (NCONS + 1) * 128, "one producer plus NCONS consumers"); +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +constexpr int PAD_N = 32; +using act_tile_n = st_bf; +using out_tile_n = st_bf; +using act_gl_n = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_MOE_ACC_C, EPI_ACC_F32 }; +#include +constexpr int R9_STATE_MODE=1; +using r9_state_t=bf16; +__device__ __forceinline__ int r9_scale_index(int head,int channel,int n) {return (((head/8)*4+(head%8)/2)*4+channel/6)*128+n;} +__device__ __forceinline__ float r9_load(r9_state_t v,float scale) { +return __bfloat162float(v); +} +__device__ __forceinline__ r9_state_t r9_store(float x,float scale) { +return __float2bfloat16(x); +} +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; + float *r8_canonical_ssm = nullptr; + r9_state_t *r9_private = nullptr; + const float *r9_scales = nullptr; + const int *r9_map = nullptr; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + unsigned int *native_pf_expert_masks = nullptr; + float *pf_dtda = nullptr; + bf16 *pf_moe_gath = nullptr; + bf16 *pf_moe_hc = nullptr; + int *pf_moe_gidx = nullptr; + int *pf_moe_ecnt = nullptr; + int *pf_moe_ntiles = nullptr; + int pf_moe_estride = 0; + act_gl pf_moe_gath_v; + act_gl pf_moe_hc_v; + act_gl_n pf_moe_gath_vn; + act_gl_n pf_moe_hc_vn; + unsigned int *native_pf_cohort_check = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + volatile unsigned int *generation = state + 1; + const unsigned int seen = *generation; + if (atomicAdd(state, 1u) == gridDim.x - 1) { + atomicExch(state, 0u); + atomicAdd(state + 1, 1u); + } else { + while (*generation == seen) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const bf16 *in = src + (size_t)row * D_MODEL; + float v[PER_THREAD]; + float sum = 0.f, sq = 0.f; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (threadIdx.x < NORM_THREADS) { + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + v[j] = __bfloat162float(in[threadIdx.x + j * NORM_THREADS]); + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + if (lane == 0) { scratch[warp] = sum; scratch[NORM_WARPS + warp] = sq; } + } + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f, q = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) { s += scratch[w]; q += scratch[NORM_WARPS + w]; } + scratch[2 * NORM_WARPS] = s / D_MODEL; + scratch[2 * NORM_WARPS + 1] = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + } + __syncthreads(); + const float mean = scratch[2 * NORM_WARPS], inv = scratch[2 * NORM_WARPS + 1]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + if (threadIdx.x >= NORM_THREADS) break; + const int i = threadIdx.x + j * NORM_THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * weight[i]); + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, + int idx, unsigned mask); +__device__ __forceinline__ void norm_sub_sync(int sub) { + if (THREADS == NORM_THREADS) { __syncthreads(); return; } + asm volatile("bar.sync %0, %1;" :: "r"(6 + sub), "n"(NORM_THREADS) : "memory"); +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, const float *fold_src = nullptr, + const float *fold_routes = nullptr) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + constexpr int NSUB = THREADS / NORM_THREADS; + if (threadIdx.x >= NSUB * NORM_THREADS) return; + const int sub = threadIdx.x / NORM_THREADS; + const int tid = threadIdx.x - sub * NORM_THREADS; + float *const sc = scratch + sub * (NORM_WARPS + 1); + for (int row = blockIdx.x * NSUB + sub; row < n_rows; + row += gridDim.x * NSUB) { + const size_t base = (size_t)row * D_MODEL; + unsigned fold_mask = (1u << BB_EXPERTS) - 1u; + if (fold_src && fold_routes) { + fold_mask = 0u; + const float *rw = fold_routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) + if (rw[e] != 0.f) fold_mask |= 1u << e; + } + float v[PER_THREAD]; + float sq = 0.f; + const int warp = tid >> 5, lane = tid & 31; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + bf16 hv; + if (fold_src) { + hv = bb_moe_fold_deterministic(fold_src, (int)(base + i), fold_mask); + h[base + i] = hv; + } else { + hv = h[base + i]; + } + v[j] = __bfloat162float(hv) + __bfloat162float(r[base + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + if (lane == 0) sc[warp] = sq; + norm_sub_sync(sub); + if (tid == 0) { + float s = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) s += sc[w]; + sc[NORM_WARPS] = rsqrtf(s / D_MODEL + RMS_EPS); + } + norm_sub_sync(sub); + const float inv = sc[NORM_WARPS]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + out[base + i] = __float2bfloat16( + __bfloat162float(weight[i]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + } + norm_sub_sync(sub); + } +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + const int c = idx - row * MB_CONV_DIM; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + const float raw = __bfloat162float( + zxbcdt[(size_t)row * MB_IN_PROJ + MB_INNER + c]); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + xbc_act[(size_t)row * MB_CONV_DIM + c] = __float2bfloat16(acc); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + const int pairs = BATCH * MB_HEADS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_HEADS * split); + const int head = (pair / split) % MB_HEADS; + const int first_channel = (pair % split) * channels; + const int group = head / MB_HEADS_PER_GROUP; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + L.dt_bias[head]); + const float dA = __expf(L.A[head] * dt); + const float Dh = L.D[head]; + float *st = L.ssm_state + (size_t)slot_map[row] * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + for (int p = first_channel; p < first_channel + channels; ++p) { + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + const size_t at = (size_t)p * MB_STATE + n; + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + static_assert(MB_CONV_K == 4, "packed conv taps assume a 4-tap kernel"); + constexpr int CPT = 4; + static_assert(MB_CONV_DIM % CPT == 0, "channel vector width"); + constexpr int CGROUPS = MB_CONV_DIM / CPT; + const int total = n_tokens * CGROUPS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / CGROUPS; + const int c = (idx - t * CGROUPS) * CPT; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + float wv[CPT][MB_CONV_K]; + { + const float4 *wsrc = reinterpret_cast( + L.conv_w + (size_t)c * MB_CONV_K); + #pragma unroll + for (int j = 0; j < CPT / 2; ++j) { + const float4 w4 = wsrc[j]; + const unsigned uu0 = __float_as_uint(w4.x); + const unsigned uu1 = __float_as_uint(w4.y); + const unsigned uu2 = __float_as_uint(w4.z); + const unsigned uu3 = __float_as_uint(w4.w); + wv[2 * j][0] = __uint_as_float(uu0 << 16); + wv[2 * j][1] = __uint_as_float(uu0 & 0xffff0000u); + wv[2 * j][2] = __uint_as_float(uu1 << 16); + wv[2 * j][3] = __uint_as_float(uu1 & 0xffff0000u); + wv[2 * j + 1][0] = __uint_as_float(uu2 << 16); + wv[2 * j + 1][1] = __uint_as_float(uu2 & 0xffff0000u); + wv[2 * j + 1][2] = __uint_as_float(uu3 << 16); + wv[2 * j + 1][3] = __uint_as_float(uu3 & 0xffff0000u); + } + } + float acc[CPT]; + { + const uint2 b2 = *reinterpret_cast(L.conv_b + c); + acc[0] = __uint_as_float(b2.x << 16); + acc[1] = __uint_as_float(b2.x & 0xffff0000u); + acc[2] = __uint_as_float(b2.y << 16); + acc[3] = __uint_as_float(b2.y & 0xffff0000u); + } + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v[CPT]; + if (src < 0) { + #pragma unroll + for (int q = 0; q < CPT; ++q) + v[q] = __bfloat162float( + state[(size_t)q * (MB_CONV_K - 1) + pos + k]); + } else { + const uint2 p = *reinterpret_cast( + zxbcdt + (size_t)src * MB_IN_PROJ + MB_INNER + c); + v[0] = __uint_as_float(p.x << 16); + v[1] = __uint_as_float(p.x & 0xffff0000u); + v[2] = __uint_as_float(p.y << 16); + v[3] = __uint_as_float(p.y & 0xffff0000u); + } + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] += wv[q][k] * v[q]; + } + uint2 out; + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] = acc[q] / (1.f + __expf(-acc[q])); + out.x = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[0])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[1])) << 16); + out.y = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[2])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[3])) << 16); + *reinterpret_cast(xbc_act + (size_t)t * MB_CONV_DIM + c) = out; + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_dtda_prefill( + const MambaLayer &L, const bf16 *zxbcdt, float *dtda, int n_tokens) { + const int total = n_tokens * MB_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_HEADS; + const int head = idx - t * MB_HEADS; + const float dt = mamba_softplus( + __bfloat162float(zxbcdt[(size_t)t * MB_IN_PROJ + MB_INNER + + MB_CONV_DIM + head]) + L.dt_bias[head]); + dtda[(size_t)idx * 2] = dt; + dtda[(size_t)idx * 2 + 1] = __expf(L.A[head] * dt); + } +} +template +__device__ __forceinline__ void mamba_ssm_prefill_impl( + const MambaLayer &L, const bf16 *xbc_act, const float *dtda, bf16 *gated, + int n_pos, int n_seq, int slot, const int *seq_slot) { + constexpr int NS = MB_STATE / 32; + constexpr int PG = MB_HEAD_DIM / PW; + constexpr bool RS = (PW % 4) == 0; + constexpr int H1 = RS ? PW / 2 : 1; + constexpr int H2 = RS ? PW / 4 : 1; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const bool keep_hi = (lane & 16) != 0; + const bool keep_hi2 = (lane & 8) != 0; + const int warps_total = gridDim.x * WARPS; + const int units = n_seq * MB_HEADS * PG; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int pg = unit % PG; + const int rest = unit / PG; + const int head = rest % MB_HEADS; + const int seq = rest / MB_HEADS; + const int p0 = pg * PW; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Dh = L.D[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p0 * MB_STATE; + float sreg[PW][NS]; + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + sreg[q][j] = st[(size_t)q * MB_STATE + lane + j * 32]; + } + auto fetch = [&](int p, float (&oB)[NS], float (&oC)[NS], + float (&oX)[PW], float2 &od) { + const size_t tt = (size_t)p * n_seq + seq; + const bf16 *xa2 = xbc_act + tt * MB_CONV_DIM; + const bf16 *Bs = xa2 + MB_INNER + group * MB_STATE; + const bf16 *Cs = xa2 + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + od = *reinterpret_cast( + dtda + (tt * MB_HEADS + head) * 2); + #pragma unroll + for (int j = 0; j < NS; ++j) { + oB[j] = __bfloat162float(Bs[lane + j * 32]); + oC[j] = __bfloat162float(Cs[lane + j * 32]); + } + if constexpr (RS) { + const uint2 *xp = reinterpret_cast( + xa2 + head * MB_HEAD_DIM + p0); + #pragma unroll + for (int q = 0; q < PW; q += 4) { + const uint2 w = xp[q >> 2]; + oX[q] = __uint_as_float(w.x << 16); + oX[q + 1] = __uint_as_float(w.x & 0xffff0000u); + oX[q + 2] = __uint_as_float(w.y << 16); + oX[q + 3] = __uint_as_float(w.y & 0xffff0000u); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) + oX[q] = __bfloat162float(xa2[head * MB_HEAD_DIM + p0 + q]); + } + }; + float pB[NS], pC[NS], pX[PW]; + float2 pd; + fetch(0, pB, pC, pX, pd); + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const float dt = pd.x; + const float dA = pd.y; + float Bv[NS], Cv[NS]; + #pragma unroll + for (int j = 0; j < NS; ++j) { + Bv[j] = pB[j] * dt; + Cv[j] = pC[j]; + } + float acc[PW], xvs[PW]; + #pragma unroll + for (int q = 0; q < PW; ++q) xvs[q] = pX[q]; + if (pos + 1 < n_pos) fetch(pos + 1, pB, pC, pX, pd); + #pragma unroll + for (int q = 0; q < PW; ++q) { + const float xv = xvs[q]; + float a = 0.f; + #pragma unroll + for (int j = 0; j < NS; ++j) { + sreg[q][j] = sreg[q][j] * dA + Bv[j] * xv; + a += sreg[q][j] * Cv[j]; + } + acc[q] = a; + } + if constexpr (RS) { + float m1[H1], t1[H1], x1[H1]; + #pragma unroll + for (int q = 0; q < H1; ++q) { + m1[q] = keep_hi ? acc[q + H1] : acc[q]; + t1[q] = keep_hi ? acc[q] : acc[q + H1]; + x1[q] = keep_hi ? xvs[q + H1] : xvs[q]; + } + #pragma unroll + for (int q = 0; q < H1; ++q) + m1[q] += __shfl_xor_sync(0xffffffff, t1[q], 16); + float m2[H2], t2[H2], x2[H2]; + #pragma unroll + for (int q = 0; q < H2; ++q) { + m2[q] = keep_hi2 ? m1[q + H2] : m1[q]; + t2[q] = keep_hi2 ? m1[q] : m1[q + H2]; + x2[q] = keep_hi2 ? x1[q + H2] : x1[q]; + } + #pragma unroll + for (int q = 0; q < H2; ++q) + m2[q] += __shfl_xor_sync(0xffffffff, t2[q], 8); + #pragma unroll + for (int q = 0; q < H2; ++q) { + #pragma unroll + for (int off = 4; off > 0; off >>= 1) + m2[q] += __shfl_xor_sync(0xffffffff, m2[q], off); + } + if ((lane & 7) == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + + p0 + (keep_hi ? H1 : 0) + (keep_hi2 ? H2 : 0); + #pragma unroll + for (int q = 0; q < H2; ++q) + dst[q] = __float2bfloat16(m2[q] + x2[q] * Dh); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc[q] += __shfl_xor_sync(0xffffffff, acc[q], off); + } + if (lane == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + p0; + #pragma unroll + for (int q = 0; q < PW; ++q) + dst[q] = __float2bfloat16(acc[q] + xvs[q] * Dh); + } + } + } + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + st[(size_t)q * MB_STATE + lane + j * 32] = sreg[q][j]; + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, + const float *dtda, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + (void)zxbcdt; + if (n_seq > 1) + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens / n_seq, n_seq, slot, seq_slot); + else + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens, 1, slot, nullptr); +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16( + held[j] * inv * __bfloat162float(L.gnorm_w[idx])); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = BATCH * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / BB_ATT_HEADS; + const int head = pair % BB_ATT_HEADS; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const int kv = head / kv_per_q; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *q = base + head * BB_HEAD_DIM; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (head % kv_per_q == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + for (int tpos = 0; tpos < T; ++tpos) { +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new + : kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#endif + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m, s); + const float rescale = __expf(m - m_new); + const float w = __expf(s - m_new); + denom = denom * rescale + w; +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *vt = kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#else + const bf16 *vt = cur_token + ? v_new + : kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + w * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + constexpr int CHUNKS = BB_HEAD_DIM / 8; + const int total = n_tokens * BB_KV_HEADS * CHUNKS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int c = idx % CHUNKS; + const int rest = idx / CHUNKS; + const int t = rest / BB_KV_HEADS; + const int kv = rest - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM + c * 8; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM + c * 8; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv) + c * 8; + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv) + c * 8; + *reinterpret_cast(kc) = *reinterpret_cast(k_new); + *reinterpret_cast(vc) = *reinterpret_cast(v_new); + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float qv[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qv[j] = __bfloat162float(q[lane + j * 32]); + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + const int *const krow = + L.block_table + ((size_t)seq * 2 + 0) * L.max_blocks_per_seq; + const int *const vrow = + L.block_table + ((size_t)seq * 2 + 1) * L.max_blocks_per_seq; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + const int head_bias = kv * L.tokens_per_block * BB_HEAD_DIM; + int page = -1; + const bf16 *kpage = nullptr, *vpage = nullptr; + float kn[PER_LANE], vn[PER_LANE]; + { + page = 0; + kpage = L.pool_base + (size_t)(krow[0] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + (size_t)(vrow[0] & 0x7FFFFFFF) * elems_per_block; + const int at = head_bias + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + for (int s = 0; s <= pos; ++s) { + float kc[PER_LANE], vc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { kc[j] = kn[j]; vc[j] = vn[j]; } + if (s < pos) { + const int np = (s + 1) / L.tokens_per_block; + if (np != page) { + page = np; + kpage = L.pool_base + + (size_t)(krow[np] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + + (size_t)(vrow[np] & 0x7FFFFFFF) * elems_per_block; + } + const int at = head_bias + + ((s + 1) % L.tokens_per_block) * BB_HEAD_DIM + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) dot += qv[j] * kc[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[j] = acc[j] * rescale + wt * vc[j]; + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH, unsigned int *tile_masks = nullptr, + bf16 *gath = nullptr, int *gidx = nullptr, int *ecnt = nullptr, + int estride = 0, bf16 *row_stage = nullptr) { + __shared__ float score[BB_EXPERTS]; + __shared__ int s_pick[BB_TOP_K], s_slot[BB_TOP_K]; + constexpr int CH = D_MODEL / 8; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + { + const float4 *gsrc = + reinterpret_cast(normed + (size_t)row * D_MODEL); + float4 *gdst = reinterpret_cast(row_stage); + for (int i = threadIdx.x; i < CH; i += THREADS) gdst[i] = gsrc[i]; + } + if ((int)threadIdx.x < BB_EXPERTS) + moe_w[(size_t)row * BB_EXPERTS + threadIdx.x] = 0.f; + __syncthreads(); + for (int e = warp; e < BB_EXPERTS; e += WARPS) { + const bf16 *wr = gate_w + (size_t)e * D_MODEL; + float acc = 0.f; + for (int i = lane; i < D_MODEL; i += 32) + acc += __bfloat162float(row_stage[i]) * + __bfloat162float(wr[i]); + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) score[e] = 1.f / (1.f + __expf(-acc)); + } + __syncthreads(); + if (threadIdx.x == 0) { + float remaining[BB_EXPERTS]; + int pick[BB_TOP_K]; + for (int e = 0; e < BB_EXPERTS; ++e) + remaining[e] = score[e] + __bfloat162float(gate_bias[e]); + float total = 0.f; + for (int k = 0; k < BB_TOP_K; ++k) { + int best = 0; + for (int e = 1; e < BB_EXPERTS; ++e) + if (remaining[e] > remaining[best]) best = e; + pick[k] = best; + total += score[best]; + remaining[best] = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + unsigned int bits = 0u; + for (int k = 0; k < BB_TOP_K; ++k) { + moe_w[row * BB_EXPERTS + pick[k]] = score[pick[k]] * norm; + bits |= 1u << pick[k]; + s_pick[k] = pick[k]; + } + if (tile_masks) atomicOr(&tile_masks[row / BATCH_PAD], bits); + } + __syncthreads(); + if (gath) { + if ((int)threadIdx.x < BB_TOP_K) { + const int k = threadIdx.x; + const int slot = atomicAdd(&ecnt[s_pick[k]], 1); + s_slot[k] = slot; + gidx[(size_t)s_pick[k] * estride + slot] = row; + } + __syncthreads(); + const float4 *src = reinterpret_cast(row_stage); + for (int k = 0; k < BB_TOP_K; ++k) { + float4 *dst = reinterpret_cast( + gath + ((size_t)s_pick[k] * estride + s_slot[k]) * D_MODEL); + for (int i = threadIdx.x; i < CH; i += THREADS) dst[i] = src[i]; + } + } + __syncthreads(); + } +} +__device__ __forceinline__ unsigned moe_row_mask(const float *routes, int row) { + unsigned m = 0u; + const float *rw = routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) if (rw[e] != 0.f) m |= 1u << e; + return m; +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx, + unsigned mask) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + for (unsigned m = mask; m; m &= m - 1u) + value = __fadd_rn(value, parts[base + (size_t)(__ffs(m) - 1) * D_MODEL]); + return __float2bfloat16( + __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL])); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + uint64_t *keys, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int n = threadIdx.x; n < SORT_N; n += THREADS) { + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + keys[n] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + __syncthreads(); + for (int k = 2; k <= SORT_N; k <<= 1) { + if (k >= 64) __syncthreads(); + for (int j = k >> 1; j > 0; j >>= 1) { + for (int i = threadIdx.x; i < SORT_N; i += THREADS) { + const int partner = i ^ j; + if (partner > i) { + const bool descending = ((i & k) == 0); + if ((keys[i] < keys[partner]) == descending) { + const uint64_t t = keys[i]; + keys[i] = keys[partner]; + keys[partner] = t; + } + } + } + if (j >= 32) __syncthreads(); + else __syncwarp(); + } + } + __syncthreads(); + if (threadIdx.x == 0) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float v; + int n; + if (!eos_used && eos_v > av) { + v = eos_v; + n = AUDIO_EOS; + eos_used = true; + } else { + v = av; + n = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = + v + gumbel[((size_t)row * N_POS + pos) * TOP_K + j]; + if (noised > best_noised) { best_noised = noised; best_code = n; } + } + out_codes[row] = best_code; + } + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void prefill_expert_masks(const float *routes, unsigned int *masks, int rows) { + const int lane=threadIdx.x&31, warp=threadIdx.x>>5; + const int nt=(rows+BATCH_PAD-1)/BATCH_PAD; + for (int t=blockIdx.x*WARPS+warp;t(gath + g * D_MODEL + c * 8); + if (local < ecnt[e]) + *dst = *reinterpret_cast( + normed + (size_t)gidx[g] * D_MODEL + c * 8); + else + *dst = make_float4(0.f, 0.f, 0.f, 0.f); + } +} +__device__ __forceinline__ int gemm_eff(int tiles, int reps) { + if (NCONS <= 1) return 1; + const int g = (int)gridDim.x; + const int rounds1 = (tiles * reps + g - 1) / g; + if constexpr (NCONS == 3) { + const int max_eff = reps <= BB_EXPERTS ? 2 : 3; + const int rounds2 = (((tiles + 1) / 2) * reps + g - 1) / g; + if (max_eff == 2) return rounds2 < rounds1 ? 2 : 1; + const int rounds3 = (((tiles + 2) / 3) * reps + g - 1) / g; + if (rounds3 < rounds2 && rounds3 < rounds1) return 3; + return rounds2 < rounds1 ? 2 : 1; + } else { + const int roundsN = (((tiles + NCONS - 1) / NCONS) * reps + g - 1) / g; + return roundsN < rounds1 ? NCONS : 1; + } +} +__device__ __forceinline__ int gemm_groups(int tiles, int reps) { + const int eff = gemm_eff(tiles, reps); + return (tiles + eff - 1) / eff; +} +__device__ __forceinline__ int gemm_consumer_id() { + const int g = warpgroup::groupid(); + if (g == 1) return -1; + const int c = g == 0 ? 0 : g - 1; + return c < NCONS ? c : -1; +} +__device__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][BATCH_PAD]; + __shared__ int moe_gidx_stage[NCONS][BATCH_PAD]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + a_s[s], a, {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + a_s[s], a, {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / BATCH_PAD) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, BATCH_PAD> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], a_s[s]); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * BATCH_PAD; + if (epi == EPI_BIAS) { + warpgroup::store(o_sf, accum); + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + dest[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(o_sf[int2{col, row}] + + __bfloat162float(bias[base + col])); + } + warpgroup::sync(1 + cons); + continue; + } + warpgroup::store(o_s[cons], accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_s[cons][int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__device__ void gemm_narrow( + const weight_gl &w, const act_gl_n &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = PAD_N, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][PAD_N]; + __shared__ int moe_gidx_stage[NCONS][PAD_N]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + PAD_N - 1) / PAD_N; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + PAD_N - 1) / PAD_N; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / PAD_N) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, PAD_N> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], + *reinterpret_cast(&a_s[s])); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * PAD_N; + out_tile_n &o_sn = *reinterpret_cast(&o_s[cons]); + warpgroup::store(o_sn, accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_sn[int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&w_s)[STAGES * NCONS] = + alloc.allocate(); + act_tile (&a_s)[STAGES] = alloc.allocate(); + out_tile (&o_s)[NCONS] = alloc.allocate(); +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 + static_assert(NCONS * sizeof(out_tile) >= sizeof(out_tile_f), + "aliased head staging must fit the output tiles"); + out_tile_f &o_sf = *reinterpret_cast(&o_s[0]); +#else + out_tile_f (&o_sf) = alloc.allocate(); +#endif + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_WARPS + 2]; +#if MK2_PREFILL_ONLY + bf16 *const route_stage = reinterpret_cast(&o_s[0]); + static_assert(sizeof(o_s) >= D_MODEL * sizeof(bf16), + "router row stage must fit the output tiles"); +#else + __shared__ __align__(16) uint64_t sort_keys[SORT_N]; + bf16 *const route_stage = reinterpret_cast(&sort_keys[0]); + static_assert(sizeof(sort_keys) >= D_MODEL * sizeof(bf16), + "router row stage must fit the sampler workspace"); +#endif + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS * NCONS, 0); + for (int c = 0; c < CONSUMER_WARPS * NCONS; ++c) arrive(drained[s]); + } + } + int gemm_iter = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map); + grid_barrier(g.barrier); + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map); + grid_barrier(g.barrier); + mamba_gated_norm(L, bb->zxbcdt, bb->gated); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, nullptr, 1, BATCH, + bb->moe_w, BB_MOE_INTER / OUT_ATOM); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + 1, BATCH, bb->moe_w, D_MODEL / OUT_ATOM); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_deterministic(bb->moe_out, idx, + (1u << BB_EXPERTS) - 1u); + grid_barrier(g.barrier); + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int original_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + bool identical_cohort = false; + if constexpr (BATCH > 1) { + if (bb->native_pf_cohort_check && bb->native_prefill_reset && + bb->pf_n_seq == BATCH && !bb->pf_fuse_decode && + !bb->pf_layer_out && !bb->moe_w_inject && !bb->moe_w_audit) { + if (blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[0] = 0; + grid_barrier(g.barrier); + const auto *bits = reinterpret_cast(bb->prompt); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < original_total * D_MODEL; idx += gridDim.x * THREADS) { + const int canonical = (idx / (BATCH * D_MODEL)) * BATCH * D_MODEL + idx % D_MODEL; + if (bits[idx] != bits[canonical]) atomicOr(bb->native_pf_cohort_check, 1u); + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + if (bb->pf_seq_slot[row] != row || bb->slot_map[row] != row) + atomicOr(bb->native_pf_cohort_check, 1u); + grid_barrier(g.barrier); + identical_cohort = bb->native_pf_cohort_check[0] == 0; + } + } + if (!identical_cohort) { + if (bb->native_pf_cohort_check && blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[1] = 0; + grid_barrier(g.barrier); + return; + } + const int n_total = original_total / BATCH; + if (bb->native_pf_cohort_check && blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[1] = identical_cohort ? 1 : BATCH; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int reset_slots = identical_cohort ? 1 : BATCH; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * reset_slots * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % reset_slots; + const MambaLayer &L = bb->mamba[row_layer / reset_slots]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = original_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + const int src = identical_cohort ? (idx / D_MODEL) * BATCH * D_MODEL + idx % D_MODEL : idx; + bb->pf_h[idx] = bb->prompt[src]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = identical_cohort ? 1 : (bb->pf_n_seq > 0 ? bb->pf_n_seq : 1); + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T, + (!bb->pf_layer_out && layer > 0 && + bb->moe_of_layer[layer - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + if (bb->native_pf_expert_masks) + for (int i = blockIdx.x * THREADS + threadIdx.x; i < n_tok; + i += gridDim.x * THREADS) + bb->native_pf_expert_masks[i] = 0u; + if (bb->pf_moe_ecnt) + for (int e = blockIdx.x * THREADS + threadIdx.x; e < BB_EXPERTS; + e += gridDim.x * THREADS) + bb->pf_moe_ecnt[e] = 0; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + for(size_t j=blockIdx.x*THREADS+threadIdx.x;jpf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + mamba_dtda_prefill(L, bb->pf_zxbcdt, bb->pf_dtda, T); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_dtda, + bb->pf_gated, T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + for(size_t j=blockIdx.x*THREADS+threadIdx.x;jpf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, + T, bb->native_pf_expert_masks, bb->pf_moe_gath, + bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, route_stage); + } + grid_barrier(g.barrier); + int max_ecnt = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + max_ecnt = max(max_ecnt, bb->pf_moe_ecnt[e]); + const bool narrow = BATCH_PAD > PAD_N && max_ecnt <= PAD_N; + const int etile = narrow ? PAD_N : BATCH_PAD; + int nt_sum = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_sum += (bb->pf_moe_ecnt[e] + etile - 1) / etile; + const int up_steps = + gemm_groups(BB_MOE_INTER / OUT_ATOM, nt_sum) * + nt_sum * (D_MODEL / KCHUNK); + const int shup_steps = + gemm_groups(BB_SHARED_INTER / OUT_ATOM, n_tok) * + n_tok * (D_MODEL / KCHUNK); + const int down_steps = + gemm_groups(D_MODEL / OUT_ATOM, nt_sum) * + nt_sum * (BB_MOE_INTER / KCHUNK); + const int shdown_steps = + gemm_groups(D_MODEL / OUT_ATOM, n_tok) * + n_tok * (BB_SHARED_INTER / KCHUNK); + if (narrow) + gemm_narrow(L.exp_up, bb->pf_moe_gath_vn, + BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, PAD_N, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + else + gemm(L.exp_up, bb->pf_moe_gath_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, BATCH_PAD, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + up_steps, up_steps + shup_steps); + grid_barrier(g.barrier); + if (narrow) + gemm_narrow(L.exp_down, bb->pf_moe_hc_vn, + BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, PAD_N, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + else + gemm(L.exp_down, bb->pf_moe_hc_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + down_steps, down_steps + shdown_steps); + grid_barrier(g.barrier); + if (bb->pf_layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * D_MODEL; idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = + bb_moe_fold_deterministic(bb->pf_moe_out, idx, + moe_row_mask(bb->pf_moe_w, idx / D_MODEL)); + grid_barrier(g.barrier); + } + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, + identical_cohort ? bb->pf_normed : bb->pf_final, ln_scratch, + n_total, + (!bb->pf_layer_out && n_backbone_layers > 0 && + bb->moe_of_layer[n_backbone_layers - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + grid_barrier(g.barrier); + if (identical_cohort) { + const int worker = blockIdx.x * THREADS + threadIdx.x; + const int workers = gridDim.x * THREADS; + for (int i = worker; i < original_total * D_MODEL; i += workers) + bb->pf_final[i] = bb->pf_normed[(i / (BATCH * D_MODEL)) * D_MODEL + i % D_MODEL]; + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = worker; i < 14 * (BATCH - 1) * CV; i += workers) { + const int cell = i % CV, pair = i / CV; + const int row = pair % (BATCH - 1) + 1; + const MambaLayer &L = bb->mamba[pair / (BATCH - 1)]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = L.conv_state[cell]; + } + const int kv_columns = BB_KV_HEADS * BB_HEAD_DIM; + for (size_t i = worker; i < size_t(3) * 2 * (BATCH - 1) * n_total * kv_columns; i += workers) { + size_t q = i / kv_columns; + const int col = i % kv_columns; + const int pos = q % n_total; q /= n_total; + const int row = q % (BATCH - 1) + 1; q /= BATCH - 1; + const int kind = q % 2, layer = q / 2; + const AttnLayer &L = bb->attn[layer]; + const int offset = kv_offset(L, pos, col / BB_HEAD_DIM) + col % BB_HEAD_DIM; + kv_slot(L, row, kind, pos)[offset] = kv_slot(L, 0, kind, pos)[offset]; + } + grid_barrier(g.barrier); + } + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } +#if MK2_PREFILL_ONLY + return; +#else + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, bb->phoneme_b); + grid_barrier(g.barrier); + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + if (threadIdx.x == 0) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = 0; n < BB_PHONEME; ++n) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + bb->phoneme_out[row] = arg; + } + __syncthreads(); + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + for (int i = lane; i < D_HEAD; i += 32) { + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + for (int t = 0; t <= pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + for (int i = lane; i < D_HEAD; i += 32) { + float acc = 0.f; + for (int t = 0; t <= pos; ++t) + acc += sc[t] * __bfloat162float( + g.vcache[cache_base + t * D_HEAD + i]); + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + i] = + __float2bfloat16(acc); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, g.sampled_pos, + sort_keys, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { +#ifdef ABLATE_SAMPLE + __syncthreads(); +#endif + } else { + grid_barrier(g.barrier); + } + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x < BATCH && blockIdx.x == 0) { + const int code = g.sampled_pos[threadIdx.x]; + g.sampled[threadIdx.x * N_POS + pos] = code; + if (g.native_codes) + g.native_codes[threadIdx.x * N_POS + pos] = (long long)code; + } + grid_barrier(g.barrier); + } + if (g.native_seq_len && blockIdx.x == 0 && threadIdx.x < BATCH) + g.native_seq_len[threadIdx.x] += 1; + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +#endif +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 2 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK), + 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = MK2_PREFILL_ONLY + ? out->sms * out->blocks_per_sm + : std::min(out->sms * out->blocks_per_sm, D_FFN / OUT_ATOM); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + 0, + nullptr, + nullptr, + nullptr, nullptr, nullptr, nullptr, nullptr, 0, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "dedup_native_prefill_impl.inc" diff --git a/deploy/nemotron_native/kernels/b128/regions/prefill/native/prefill/mk2_abi.cuh b/deploy/nemotron_native/kernels/b128/regions/prefill/native/prefill/mk2_abi.cuh new file mode 100644 index 0000000..95a253b --- /dev/null +++ b/deploy/nemotron_native/kernels/b128/regions/prefill/native/prefill/mk2_abi.cuh @@ -0,0 +1,82 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} diff --git a/deploy/nemotron_native/kernels/b128/regions/prefill/native/prefill/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/b128/regions/prefill/native/prefill/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/b128/regions/prefill/native/prefill/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/b128/regions/prefill/native/prefill/prefill.cu b/deploy/nemotron_native/kernels/b128/regions/prefill/native/prefill/prefill.cu new file mode 100644 index 0000000..3033076 --- /dev/null +++ b/deploy/nemotron_native/kernels/b128/regions/prefill/native/prefill/prefill.cu @@ -0,0 +1,18 @@ +#define MK2_BATCH 128 +#define MK2_BATCH_TILE 64 +#define MK2_PREFILL_ONLY 1 +#define STAGES 3 +#define KCHUNK 128 +#define MK2_GEMM_CONSUMERS 3 +#define MK2_THREADS 512 +#define MK2_NO_MAIN 1 +#define MK2_DIRECT_INPUT 1 +#define MK2_B1_LOCAL_BARRIERS 1 +#define MK2_SKIP_HIDDEN_SNAPSHOT 1 +#define MK2_ACTIVE_EXPERTS 1 +#define MK2_SSM_SPLITS 1 +#define PF_SSM_PW_MULTI 12 +#define PF_ACTIVE_TOKEN_EXPERTS 1 +#include +#include +#include "dedup_predictor_mk2.cuh" diff --git a/deploy/nemotron_native/kernels/b128/regions/prefill/prefill.cu b/deploy/nemotron_native/kernels/b128/regions/prefill/prefill.cu new file mode 100644 index 0000000..3dce994 --- /dev/null +++ b/deploy/nemotron_native/kernels/b128/regions/prefill/prefill.cu @@ -0,0 +1,20 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace b128_prefill { +#include "native/prefill/prefill.cu" +} diff --git a/deploy/nemotron_native/kernels/b128/regions/stage0/native/stage0/mk2_abi.cuh b/deploy/nemotron_native/kernels/b128/regions/stage0/native/stage0/mk2_abi.cuh new file mode 100644 index 0000000..bf1e415 --- /dev/null +++ b/deploy/nemotron_native/kernels/b128/regions/stage0/native/stage0/mk2_abi.cuh @@ -0,0 +1,85 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_set_native_predictor_active_mask(MK2Predictor *p, const unsigned char *mask); +int mk2_native_startup_skip_abi_version(); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *, int grid); diff --git a/deploy/nemotron_native/kernels/b128/regions/stage0/native/stage0/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/b128/regions/stage0/native/stage0/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/b128/regions/stage0/native/stage0/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/b128/regions/stage0/native/stage0/native_prefill_impl.inc b/deploy/nemotron_native/kernels/b128/regions/stage0/native/stage0/native_prefill_impl.inc new file mode 100644 index 0000000..a254308 --- /dev/null +++ b/deploy/nemotron_native/kernels/b128/regions/stage0/native/stage0/native_prefill_impl.inc @@ -0,0 +1,198 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 16 || BATCH == 32 || BATCH == 64)) + throw std::runtime_error("native prefill requires matching batch 1/16/32/64"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (padded > INT_MAX / ((BB_EXPERTS + 1) * D_MODEL)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + float *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(float))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(float))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} diff --git a/deploy/nemotron_native/kernels/b128/regions/stage0/native/stage0/predictor_mk2.cu b/deploy/nemotron_native/kernels/b128/regions/stage0/native/stage0/predictor_mk2.cu new file mode 100644 index 0000000..ddf223f --- /dev/null +++ b/deploy/nemotron_native/kernels/b128/regions/stage0/native/stage0/predictor_mk2.cu @@ -0,0 +1,6107 @@ +#include "../../../../../inflight/progress.cuh" +#if MK2_INFLIGHT +#include +#endif +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#define MK2_FUSE_CONV (MK2_SSM_SPLITS == 1) +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +constexpr int BATCH_PAD = 64; +constexpr int DECODE_TILES = BATCH / BATCH_PAD; +constexpr int DECODE_ROWS = BATCH_PAD; +static_assert(BATCH == 128 || BATCH == 256 || BATCH == 512 || BATCH == 1024 || BATCH == 2048); +static_assert(BATCH % BATCH_PAD == 0); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +#ifndef NMULT +#define NMULT 1 +#endif +#ifndef SSM_PREFETCH_MAX +#define SSM_PREFETCH_MAX 8 +#endif +constexpr int NM_MAX = (NMULT > 1 && DECODE_TILES % NMULT == 0) ? NMULT : 1; +#ifndef MK2_OUT_MULT +#define MK2_OUT_MULT 1 +#endif +constexpr int MT = MK2_OUT_MULT; +static_assert(MT >= 1 && MT <= NM_MAX, "an item holds at most NM_MAX accumulators"); +constexpr int SLOTS = 1 + NM_MAX; +static_assert(MT + 1 <= SLOTS); +constexpr int THREADS = 256; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +constexpr int NORM_SCRATCH_STRIDE = 2 * WARPS + 2; +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_ACC_F32 }; +#define EPI_ROW(idx) ((OUT_ATOM == 64 && BATCH_PAD == 64) \ + ? (((idx) & 15) | ((((idx) >> 6) & 3) << 4)) : ((idx) / (OUT_ATOM / 8))) +#define EPI_OCT(idx) ((OUT_ATOM == 64 && BATCH_PAD == 64) \ + ? ((((idx) >> 4) & 3) | ((((idx) >> 8) & 1) << 2)) : ((idx) % (OUT_ATOM / 8))) +struct alignas(16) bf16_oct { bf16 x0, x1, x2, x3, x4, x5, x6, x7; }; +static_assert(sizeof(bf16_oct) == 16); +#include +constexpr int R9_STATE_MODE=1; +using r9_state_t=bf16; +__device__ __forceinline__ int r9_scale_index(int head,int channel,int n) {return (((head/8)*4+(head%8)/2)*4+channel/6)*128+n;} +__device__ __forceinline__ float r9_load(r9_state_t v,float scale) { +return __bfloat162float(v); +} +__device__ __forceinline__ r9_state_t r9_store(float x,float scale) { +return __float2bfloat16(x); +} +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; + const float *r8_initial_ssm = nullptr; + const unsigned int *r8_initial_ready = nullptr; + const int *r8_seq_len = nullptr; + int r8_initial_positions = 0; + const float *r9_scales = nullptr; + const int *r9_map = nullptr; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + bf16 *r6_packed = nullptr; + int *r6_map = nullptr,*r6_counts = nullptr; + act_gl *r6_views = nullptr; + int *r6_lut = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +#define BARRIER_GEN_WORD 32 +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + constexpr unsigned long long GRID_BLOCKS = 132ull; + auto *wide_state = reinterpret_cast(state); + const unsigned long long arrival = atomicAdd(wide_state, 1ull); + const unsigned int rank = (unsigned int)(arrival % GRID_BLOCKS); + if (rank != GRID_BLOCKS - 1u) { + const unsigned long long target = arrival - rank + GRID_BLOCKS; + volatile unsigned long long *counter = wide_state; + while (*counter < target) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / THREADS; + float wv[PER_THREAD]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) wv[j] = weight[threadIdx.x + j * THREADS]; + float v[PER_THREAD], v_next[PER_THREAD]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + v[j] = __bfloat162float( + src[(size_t)blockIdx.x * D_MODEL + threadIdx.x + j * THREADS]); + for (int row = blockIdx.x, row_iter = 0; row < BATCH; + row += gridDim.x, ++row_iter) { + float *row_scratch = scratch + (row_iter & 1) * NORM_SCRATCH_STRIDE; + const int next_row = row + (int)gridDim.x; + if (next_row < BATCH) { + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + v_next[j] = __bfloat162float( + src[(size_t)next_row * D_MODEL + threadIdx.x + j * THREADS]); + } + float sum = 0.f, sq = 0.f; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) { + row_scratch[warp] = sum; + row_scratch[WARPS + warp] = sq; + } + __syncthreads(); + float s = 0.f, q = 0.f; + #pragma unroll + for (int w = 0; w < WARPS; ++w) { + s += row_scratch[w]; + q += row_scratch[WARPS + w]; + } + const float mean = s / D_MODEL; + const float inv = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * wv[j]); + } + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) v[j] = v_next[j]; + } +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, bf16 *snapshot = nullptr, int snapshot_layer = -1, + float *snapshot_fp32 = nullptr, float *out_fp32 = nullptr, + const float *fold_src = nullptr, const float *fold_routes = nullptr) { + constexpr int PER_THREAD = D_MODEL / THREADS; + bf16 wv[PER_THREAD]; + extern __shared__ int __shm[]; + bf16 *const row_stage = reinterpret_cast(__shm); + constexpr int ROW_SLOT = 2 * D_MODEL; + constexpr int ROW_VECTORS = D_MODEL / 8; + if ((int)blockIdx.x < n_rows && threadIdx.x < ROW_VECTORS) { + const int i = 8 * threadIdx.x; + const size_t at = (size_t)blockIdx.x * D_MODEL + i; + const uint32_t rs = static_cast( + __cvta_generic_to_shared(row_stage + D_MODEL + i)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;" + :: "r"(rs), "l"(r + at) : "memory"); + if (!fold_src) { + const uint32_t hs = static_cast( + __cvta_generic_to_shared(row_stage + i)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;" + :: "r"(hs), "l"(h + at) : "memory"); + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + wv[j] = weight[threadIdx.x + j * THREADS]; + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); + __syncthreads(); + for (int row = blockIdx.x, row_iter = 0; row < n_rows; + row += gridDim.x, ++row_iter) { + bf16 *const current = row_stage + (row_iter & 1) * ROW_SLOT; + const int next_row = row + (int)gridDim.x; + if (next_row < n_rows) { + bf16 *const next = row_stage + ((row_iter + 1) & 1) * ROW_SLOT; + if (threadIdx.x < ROW_VECTORS) { + const int i = 8 * threadIdx.x; + const size_t at = (size_t)next_row * D_MODEL + i; + const uint32_t rs = static_cast( + __cvta_generic_to_shared(next + D_MODEL + i)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;" + :: "r"(rs), "l"(r + at) : "memory"); + if (!fold_src) { + const uint32_t hs = static_cast( + __cvta_generic_to_shared(next + i)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;" + :: "r"(hs), "l"(h + at) : "memory"); + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } + float *row_scratch = scratch + (row_iter & 1) * NORM_SCRATCH_STRIDE; + const size_t base = (size_t)row * D_MODEL; + float fv[PER_THREAD]; + if (fold_src) { + unsigned mask = 0u; + const float *rw = fold_routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) + if (rw[e] != 0.f) mask |= 1u << e; + const float *pe = fold_src + + (size_t)row * (BB_EXPERTS + 1) * D_MODEL + threadIdx.x; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) fv[j] = 0.f; + #pragma unroll + for (int e = 0; e <= BB_EXPERTS; ++e) { + const bool take = e == BB_EXPERTS || ((mask >> e) & 1u); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + fv[j] = __fadd_rn(fv[j], + take ? pe[(size_t)e * D_MODEL + j * THREADS] : 0.f); + } + } + float v[PER_THREAD]; + float sq = 0.f; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + const bf16 hv = fold_src ? __float2bfloat16(fv[j]) : current[i]; + if (fold_src) h[base + i] = hv; + if (snapshot) + snapshot[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = hv; + if (snapshot_fp32) + snapshot_fp32[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = + __bfloat162float(hv); + v[j] = __bfloat162float(hv) + + __bfloat162float(current[D_MODEL + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) row_scratch[warp] = sq; + __syncthreads(); + float s = 0.f; + #pragma unroll + for (int w = 0; w < WARPS; ++w) s += row_scratch[w]; + const float inv = rsqrtf(s / D_MODEL + RMS_EPS); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + const bf16 normed = __float2bfloat16( + __bfloat162float(wv[j]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + out[base + i] = normed; + if (out_fp32) out_fp32[base + i] = __bfloat162float(normed); + } + if (next_row < n_rows) { + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); + __syncthreads(); + } + } +} +__device__ __forceinline__ bf16 mamba_conv_channel( + const MambaLayer &L, bf16 *state_slot, const bf16 *zrow, int c) { + const unsigned ch = (unsigned)c; + const bf16 *w = L.conv_w + ch * (unsigned)MB_CONV_K; + bf16 *state = state_slot + ch * (unsigned)(MB_CONV_K - 1); + const bf16 rawb = zrow[c]; + const float raw = __bfloat162float(rawb); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = rawb; + return __float2bfloat16(acc); +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map MK2_PROGRESS_PARAMETER(progress_mask)) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + MK2_PROGRESS_GUARD(progress_mask, row); + const int c = idx - row * MB_CONV_DIM; + xbc_act[(size_t)row * MB_CONV_DIM + c] = mamba_conv_channel( + L, L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride, + zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER, c); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map MK2_PROGRESS_PARAMETER(progress_mask)) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int sb4 = (lane >> 4) & 1, sb3 = (lane >> 3) & 1; + const int sb2 = (lane >> 2) & 1, sb1 = (lane >> 1) & 1; + const int ssm_own_ch = 6 * sb4 + 3 * sb3 + 2 * sb2 + sb1; + const bool ssm_owner = (2 * sb2 + sb1) < 3 && (lane & 1) == 0; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + constexpr bool FUSE_GNORM = (split == 1); + extern __shared__ int __shm[]; + bf16 (*const gn_stage)[MB_GROUP_SIZE] = + reinterpret_cast(&__shm[0]); + constexpr bool FUSE_CONV = (MK2_FUSE_CONV != 0); + bf16 (*const conv_x)[MB_GROUP_SIZE] = gn_stage + WARPS; + constexpr int SSM_NSTATE = MB_STATE / 32; + constexpr int SSM_PF = channels < SSM_PREFETCH_MAX ? channels : SSM_PREFETCH_MAX; + constexpr int SSM_GROUP_BYTES = SSM_PF * MB_STATE * (int)sizeof(r9_state_t); + constexpr bool SSM_SMEM_OK = R9_STATE_MODE == 1 && channels == 2 * SSM_PF && + SSM_GROUP_BYTES % 512 == 0; + constexpr int SSM_RING = 4; + constexpr int SSM_GROUPS_PER_HEAD = channels / SSM_PF; + constexpr int SSM_PAIRS = MB_HEADS_PER_GROUP * SSM_GROUPS_PER_HEAD; + static_assert(!SSM_SMEM_OK || (SSM_RING & (SSM_RING - 1)) == 0, + "the flat (head, channel group) ring index is masked"); + static_assert(!SSM_SMEM_OK || SSM_PAIRS >= SSM_RING, + "the prologue must be able to fill the ring"); + static_assert(!SSM_SMEM_OK || + 2 * WARPS * MB_GROUP_SIZE * sizeof(bf16) + + WARPS * SSM_RING * SSM_GROUP_BYTES <= STAGES * SLOTS * sizeof(weight_tile)); + constexpr int SSM_F32_RING = 2; + constexpr int SSM_F32_GROUP_BYTES = SSM_PF * MB_STATE * (int)sizeof(float); + constexpr bool SSM_F32_OK = SSM_SMEM_OK && SSM_F32_GROUP_BYTES % 512 == 0 && + (SSM_F32_RING & (SSM_F32_RING - 1)) == 0 && SSM_PAIRS >= SSM_F32_RING && + SSM_F32_RING * SSM_F32_GROUP_BYTES <= SSM_RING * SSM_GROUP_BYTES; + char *const ssm_stage = reinterpret_cast(conv_x + WARPS); + const int pairs = BATCH * MB_GROUPS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_GROUPS * split); + MK2_PROGRESS_GUARD(progress_mask, row); + const int group = (pair / split) % MB_GROUPS; + const int first_channel = (pair % split) * channels; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const bool r8_initial = L.r8_initial_ssm && L.r8_initial_ready && + *L.r8_initial_ready == 1 && L.r8_seq_len[row] == L.r8_initial_positions; + r9_state_t *st = reinterpret_cast(L.ssm_state) + + (size_t)slot_map[row] * L.ssm_slot_stride; + char *const warp_ssm = + ssm_stage + (size_t)warp * SSM_RING * SSM_GROUP_BYTES; + auto issue_ssm = [&](int head_in_group, int channel0, int buffer) { + const int head = group * MB_HEADS_PER_GROUP + head_in_group; + const r9_state_t *src = st + + (size_t)(head * MB_HEAD_DIM + channel0) * MB_STATE; + const uint32_t dst = (uint32_t)__cvta_generic_to_shared( + warp_ssm + (size_t)buffer * SSM_GROUP_BYTES + lane * 16); + #pragma unroll + for (int chunk = 0; chunk < SSM_GROUP_BYTES / 512; ++chunk) + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: + "r"(dst + (uint32_t)chunk * 512), + "l"(reinterpret_cast(src) + + (size_t)chunk * 512 + lane * 16) : "memory"); + asm volatile("cp.async.commit_group;\n" ::: "memory"); + }; + auto issue_ssm_f32 = [&](int head_in_group, int channel0, int buffer) { + const int head = group * MB_HEADS_PER_GROUP + head_in_group; + const float *src = L.r8_initial_ssm + + (size_t)(head * MB_HEAD_DIM + channel0) * MB_STATE; + const uint32_t dst = (uint32_t)__cvta_generic_to_shared( + warp_ssm + (size_t)buffer * SSM_F32_GROUP_BYTES + lane * 16); + #pragma unroll + for (int chunk = 0; chunk < SSM_F32_GROUP_BYTES / 512; ++chunk) + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: + "r"(dst + (uint32_t)chunk * 512), + "l"(reinterpret_cast(src) + + (size_t)chunk * 512 + lane * 16) : "memory"); + asm volatile("cp.async.commit_group;\n" ::: "memory"); + }; + if constexpr (SSM_SMEM_OK) { + if (!r8_initial) { + #pragma unroll + for (int k = 0; k < SSM_RING; ++k) + issue_ssm(k / SSM_GROUPS_PER_HEAD, + (k % SSM_GROUPS_PER_HEAD) * SSM_PF, k); + } else if constexpr (SSM_F32_OK) { + #pragma unroll + for (int k = 0; k < SSM_F32_RING; ++k) + issue_ssm_f32(k / SSM_GROUPS_PER_HEAD, + (k % SSM_GROUPS_PER_HEAD) * SSM_PF, k); + } + } + float Braw[MB_STATE / 32]; + float Cv[MB_STATE / 32]; + const bf16 *xrow; + if constexpr (FUSE_CONV) { + bf16 *const cslot = + L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride; + const bf16 *const zrow = + zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER; + static_assert(MB_CONV_K == 4, "the state record is three 32-bit words"); + static_assert(MB_GROUP_SIZE % 64 == 0 && MB_STATE % 64 == 0, + "one paired round covers 64 channels of a warp"); + auto wide_lo = [](unsigned int v) { return __uint_as_float(v << 16); }; + auto wide_hi = [](unsigned int v) { return __uint_as_float(v & 0xffff0000u); }; + struct ConvCursor { + const uint2 *w; unsigned int *s; + const unsigned int *z; const unsigned int *b; + }; + auto conv_cursor = [&](int c) { + const unsigned int cu = (unsigned int)c; + return ConvCursor{ + reinterpret_cast( + L.conv_w + cu * (unsigned int)MB_CONV_K), + reinterpret_cast( + cslot + cu * (unsigned int)(MB_CONV_K - 1)), + reinterpret_cast(zrow + cu), + reinterpret_cast(L.conv_b + cu)}; + }; + auto conv_two = [&](uint2 w0, uint2 w1, unsigned int s0, + unsigned int s1, unsigned int s2, + unsigned int zp, unsigned int bp, float &a0, + float &a1, unsigned int &n0, unsigned int &n1, + unsigned int &n2) { + a0 = wide_lo(bp); + a0 += wide_lo(w0.x) * wide_lo(s0); + a0 += wide_hi(w0.x) * wide_hi(s0); + a0 += wide_lo(w0.y) * wide_lo(s1); + a0 += wide_hi(w0.y) * wide_lo(zp); + a0 = a0 / (1.f + __expf(-a0)); + a1 = wide_hi(bp); + a1 += wide_lo(w1.x) * wide_hi(s1); + a1 += wide_hi(w1.x) * wide_lo(s2); + a1 += wide_lo(w1.y) * wide_hi(s2); + a1 += wide_hi(w1.y) * wide_hi(zp); + a1 = a1 / (1.f + __expf(-a1)); + n0 = (s0 >> 16) | (s1 << 16); + n1 = (zp & 0xffffu) | (s2 << 16); + n2 = (s2 >> 16) | (zp & 0xffff0000u); + }; + struct ConvPairIn { + uint2 w0, w1; + unsigned int s0, s1, s2, zp, bp; + }; + auto conv_pair_fetch = [](const ConvCursor &q, int r) { + const uint2 *const wp = q.w + 64 * r; + const unsigned int *const sp = q.s + 96 * r; + return ConvPairIn{wp[0], wp[1], sp[0], sp[1], sp[2], + q.z[32 * r], q.b[32 * r]}; + }; + auto conv_pair_apply = [&](unsigned int *const sp, + const ConvPairIn &v, float &a0, + float &a1) { + unsigned int n0, n1, n2; + conv_two(v.w0, v.w1, v.s0, v.s1, v.s2, v.zp, v.bp, a0, a1, + n0, n1, n2); + sp[0] = n0; + sp[1] = n1; + sp[2] = n2; + }; + constexpr int QUAD_SPAN = 4 * 32; + static_assert(MB_INNER % 4 == 0 && MB_STATE % 4 == 0 && + MB_GROUP_SIZE % 4 == 0, + "every quad base is a four-channel boundary"); + static_assert(MB_STATE == QUAD_SPAN && + MB_GROUP_SIZE == QUAD_SPAN + 2 * 32, + "one quad round covers B and C; x adds a pair round"); + struct ConvQuad { + const uint4 *w; uint2 *s; const uint2 *z; const uint2 *b; + }; + auto conv_quad_cursor = [&](int c) { + const unsigned int cu = (unsigned int)c; + return ConvQuad{ + reinterpret_cast( + L.conv_w + cu * (unsigned int)MB_CONV_K), + reinterpret_cast( + cslot + cu * (unsigned int)(MB_CONV_K - 1)), + reinterpret_cast(zrow + cu), + reinterpret_cast(L.conv_b + cu)}; + }; + struct ConvQuadIn { + uint4 w01, w23; + uint2 sa, sb, sc, zp, bp; + }; + auto conv_quad_fetch = [](const ConvQuad &q) { + const uint2 *const sp = q.s; + return ConvQuadIn{q.w[0], q.w[1], sp[0], sp[1], sp[2], + q.z[0], q.b[0]}; + }; + auto conv_quad_apply = [&](uint2 *const sp, const ConvQuadIn &v, + float &a0, float &a1, float &a2, + float &a3) { + unsigned int n0, n1, n2, n3, n4, n5; + conv_two(make_uint2(v.w01.x, v.w01.y), + make_uint2(v.w01.z, v.w01.w), + v.sa.x, v.sa.y, v.sb.x, v.zp.x, v.bp.x, + a0, a1, n0, n1, n2); + conv_two(make_uint2(v.w23.x, v.w23.y), + make_uint2(v.w23.z, v.w23.w), + v.sb.y, v.sc.x, v.sc.y, v.zp.y, v.bp.y, + a2, a3, n3, n4, n5); + sp[0] = make_uint2(n0, n1); + sp[1] = make_uint2(n2, n3); + sp[2] = make_uint2(n4, n5); + }; + auto conv_pack = [](float a, float b) { + return (unsigned int)__bfloat16_as_ushort(__float2bfloat16(a)) | + ((unsigned int)__bfloat16_as_ushort(__float2bfloat16(b)) << 16); + }; + const ConvQuad qb = conv_quad_cursor( + MB_INNER + group * MB_STATE + 4 * lane); + const ConvQuad qc = conv_quad_cursor( + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE + 4 * lane); + static_assert(MB_GROUP_SIZE * (int)sizeof(bf16) % 8 == 0, + "each warp's staging row starts 8-byte aligned"); + const ConvQuad qx = conv_quad_cursor( + group * MB_GROUP_SIZE + 4 * lane); + const ConvCursor qt = conv_cursor( + group * MB_GROUP_SIZE + QUAD_SPAN + 2 * lane); + const ConvQuadIn vb = conv_quad_fetch(qb); + const ConvQuadIn vc = conv_quad_fetch(qc); + const ConvQuadIn vx = conv_quad_fetch(qx); + const ConvPairIn vt = conv_pair_fetch(qt, 0); + uint2 bpk, cpk; + { + float a0, a1, a2, a3; + conv_quad_apply(qb.s, vb, a0, a1, a2, a3); + bpk = make_uint2(conv_pack(a0, a1), conv_pack(a2, a3)); + conv_quad_apply(qc.s, vc, a0, a1, a2, a3); + cpk = make_uint2(conv_pack(a0, a1), conv_pack(a2, a3)); + } + { + float a0, a1, a2, a3; + conv_quad_apply(qx.s, vx, a0, a1, a2, a3); + *reinterpret_cast(&conv_x[warp][4 * lane]) = + make_uint2(conv_pack(a0, a1), conv_pack(a2, a3)); + } + { + float a0, a1; + conv_pair_apply(qt.s, vt, a0, a1); + conv_x[warp][QUAD_SPAN + 2 * lane] = __float2bfloat16(a0); + conv_x[warp][QUAD_SPAN + 2 * lane + 1] = __float2bfloat16(a1); + } + const bool qword = ((lane >> 1) & 1) != 0; + const bool qhalf = (lane & 1) != 0; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int src = (lane >> 2) + 8 * j; + const unsigned int bx = __shfl_sync(0xffffffffu, bpk.x, src); + const unsigned int by = __shfl_sync(0xffffffffu, bpk.y, src); + const unsigned int cx = __shfl_sync(0xffffffffu, cpk.x, src); + const unsigned int cy = __shfl_sync(0xffffffffu, cpk.y, src); + const unsigned int bg = qword ? by : bx; + const unsigned int cg = qword ? cy : cx; + Braw[j] = qhalf ? wide_hi(bg) : wide_lo(bg); + Cv[j] = qhalf ? wide_hi(cg) : wide_lo(cg); + } + __syncwarp(); + xrow = &conv_x[warp][0]; + } else { + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Braw[j] = __bfloat162float(Bsec[lane + j * 32]); + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + xrow = xa + (size_t)group * MB_GROUP_SIZE; + } + const int head_lane = group * MB_HEADS_PER_GROUP + + (lane < MB_HEADS_PER_GROUP ? lane : 0); + const float dt_lane = mamba_softplus( + __bfloat162float(dtrow[head_lane]) + L.dt_bias[head_lane]); + const float dA_lane = __expf(L.A[head_lane] * dt_lane); + const float D_lane = L.D[head_lane]; + #pragma unroll 1 + for (int head_in_group = 0; head_in_group < MB_HEADS_PER_GROUP; + ++head_in_group) { + const int head = group * MB_HEADS_PER_GROUP + head_in_group; + const float dt = __shfl_sync(0xffffffffu, dt_lane, head_in_group); + const float dA = __shfl_sync(0xffffffffu, dA_lane, head_in_group); + const float Dh = __shfl_sync(0xffffffffu, D_lane, head_in_group); + const float *st_in = L.r8_initial_ssm + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float Bv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) + Bv[j] = Braw[j] * dt; + constexpr int NSTATE = SSM_NSTATE; + constexpr int SSM_PREFETCH = SSM_PF; + static_assert(channels % SSM_PREFETCH == 0); + for (int g0 = 0; g0 < channels; g0 += SSM_PREFETCH) { + float sv[SSM_PREFETCH][NSTATE]; + float xvr[SSM_PREFETCH]; + if (r8_initial && SSM_F32_OK) { + asm volatile("cp.async.wait_group 1;\n" ::: "memory"); + __syncwarp(); + const int pair = head_in_group * SSM_GROUPS_PER_HEAD + + g0 / SSM_PREFETCH; + const int buffer = pair & (SSM_F32_RING - 1); + const float *const smf = reinterpret_cast( + warp_ssm + (size_t)buffer * SSM_F32_GROUP_BYTES); + static_assert(SSM_PREFETCH % 2 == 0 && MB_HEAD_DIM % 2 == 0, + "x staging vectorises in element pairs"); + const __nv_bfloat162 *const xq = + reinterpret_cast( + xrow + head_in_group * MB_HEAD_DIM + + first_channel + g0); + #pragma unroll + for (int gq = 0; gq < SSM_PREFETCH / 2; ++gq) { + const float2 xp = __bfloat1622float2(xq[gq]); + xvr[2 * gq] = xp.x; + xvr[2 * gq + 1] = xp.y; + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = smf[(size_t)gi * MB_STATE + lane + j * 32]; + } + __syncwarp(); + const int ahead = pair + SSM_F32_RING; + if (ahead < SSM_PAIRS) + issue_ssm_f32(ahead / SSM_GROUPS_PER_HEAD, + (ahead % SSM_GROUPS_PER_HEAD) * SSM_PF, + buffer); + else + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } else if (r8_initial) { + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const size_t prow = (size_t)(first_channel + g0 + gi) * MB_STATE; + xvr[gi] = __bfloat162float( + xrow[head_in_group * MB_HEAD_DIM + + first_channel + g0 + gi]); + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = st_in[prow + lane + j * 32]; + } + } else if constexpr (SSM_SMEM_OK) { + asm volatile("cp.async.wait_group 3;\n" ::: "memory"); + __syncwarp(); + const int pair = head_in_group * SSM_GROUPS_PER_HEAD + + g0 / SSM_PREFETCH; + const int buffer = pair & (SSM_RING - 1); + const r9_state_t *sm = reinterpret_cast( + warp_ssm + (size_t)buffer * SSM_GROUP_BYTES); + static_assert(SSM_PREFETCH % 2 == 0 && MB_HEAD_DIM % 2 == 0, + "x staging vectorises in element pairs"); + const __nv_bfloat162 *const xq = + reinterpret_cast( + xrow + head_in_group * MB_HEAD_DIM + + first_channel + g0); + #pragma unroll + for (int gq = 0; gq < SSM_PREFETCH / 2; ++gq) { + const float2 xp = __bfloat1622float2(xq[gq]); + xvr[2 * gq] = xp.x; + xvr[2 * gq + 1] = xp.y; + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = r9_load( + sm[(size_t)gi * MB_STATE + lane + j * 32], 1.f); + } + __syncwarp(); + const int ahead = pair + SSM_RING; + if (ahead < SSM_PAIRS) + issue_ssm(ahead / SSM_GROUPS_PER_HEAD, + (ahead % SSM_GROUPS_PER_HEAD) * SSM_PF, buffer); + else + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } else { + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + xvr[gi] = __bfloat162float( + xrow[head_in_group * MB_HEAD_DIM + + first_channel + g0 + gi]); + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + { + int p = first_channel+g0+gi; + int mapped = R9_STATE_MODE==4 ? L.r9_map[head*24+p] : head*24+p; + int n = lane+j*32; + float scale = R9_STATE_MODE==1 ? 1.f : L.r9_scales[r9_scale_index(mapped/24,mapped%24,n)]; + sv[gi][j] = r9_load(st[(size_t)mapped*128+n],scale); + } + } + } + if constexpr (SSM_PREFETCH == 12) { + float v6[6]; + #pragma unroll + for (int k = 0; k < 6; ++k) { + float pair_acc[2]; + #pragma unroll + for (int half = 0; half < 2; ++half) { + const int gi = k + 6 * half; + const int p = first_channel + g0 + gi; + const float xv = xvr[gi]; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < NSTATE; ++j) { + const int n = lane + j * 32; + const float s = sv[gi][j] * dA + Bv[j] * xv; + const int mapped = R9_STATE_MODE == 4 + ? L.r9_map[head * 24 + p] : head * 24 + p; + const float scale = R9_STATE_MODE == 1 ? 1.f + : L.r9_scales[r9_scale_index( + mapped / 24, mapped % 24, n)]; + st[(size_t)mapped * 128 + n] = r9_store(s, scale); + acc += s * Cv[j]; + } + pair_acc[half] = acc; + } + const float mine = sb4 ? pair_acc[1] : pair_acc[0]; + const float away = sb4 ? pair_acc[0] : pair_acc[1]; + v6[k] = mine + __shfl_xor_sync(0xffffffff, away, 16); + } + float v4[4]; + #pragma unroll + for (int k = 0; k < 3; ++k) { + const float mine = sb3 ? v6[3 + k] : v6[k]; + const float away = sb3 ? v6[k] : v6[3 + k]; + v4[k] = mine + __shfl_xor_sync(0xffffffff, away, 8); + } + v4[3] = 0.f; + float v2[2]; + #pragma unroll + for (int k = 0; k < 2; ++k) { + const float mine = sb2 ? v4[2 + k] : v4[k]; + const float away = sb2 ? v4[k] : v4[2 + k]; + v2[k] = mine + __shfl_xor_sync(0xffffffff, away, 4); + } + float acc = (sb1 ? v2[1] : v2[0]) + + __shfl_xor_sync(0xffffffff, sb1 ? v2[0] : v2[1], 2); + acc += __shfl_xor_sync(0xffffffff, acc, 1); + if (ssm_owner) { + const int p = first_channel + g0 + ssm_own_ch; + const float xv = __bfloat162float( + xrow[head_in_group * MB_HEAD_DIM + p]); + const bf16 outv = __float2bfloat16(acc + xv * Dh); + if constexpr (FUSE_GNORM) + gn_stage[warp][head_in_group * MB_HEAD_DIM + p] = outv; + else + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = outv; + } + } else { + float cacc[SSM_PREFETCH]; + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const int p = first_channel + g0 + gi; + const float xv = xvr[gi]; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < NSTATE; ++j) { + const int n = lane + j * 32; + const float s = sv[gi][j] * dA + Bv[j] * xv; + const int mapped = R9_STATE_MODE == 4 + ? L.r9_map[head * 24 + p] : head * 24 + p; + const float scale = R9_STATE_MODE == 1 ? 1.f + : L.r9_scales[r9_scale_index( + mapped / 24, mapped % 24, n)]; + st[(size_t)mapped * 128 + n] = r9_store(s, scale); + acc += s * Cv[j]; + } + cacc[gi] = acc; + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const int p = first_channel + g0 + gi; + float acc = cacc[gi]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) { + const bf16 outv = __float2bfloat16(acc + xvr[gi] * Dh); + if constexpr (FUSE_GNORM) + gn_stage[warp][head_in_group * MB_HEAD_DIM + p] = outv; + else + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = outv; + } + } + } + } + } + if constexpr (SSM_SMEM_OK) { + if (!r8_initial || SSM_F32_OK) { + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); + __syncwarp(); + } + } + if constexpr (FUSE_GNORM) { + __syncwarp(); + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32], wg[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float gv = __bfloat162float(gn_stage[warp][lane + j * 32]) * + (z / (1.f + __expf(-z))); + wg[j] = __bfloat162float(L.gnorm_w[idx]); + held[j] = gv; + sq += gv * gv; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16(held[j] * inv * wg[j]); + } + __syncwarp(); + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_CONV_DIM; + const int c = idx - t * MB_CONV_DIM; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v; +#ifdef PF_CONV_ZERO_STATE + if (src < 0) v = 0.f; +#else + if (src < 0) v = __bfloat162float(state[pos + k]); +#endif + else v = __bfloat162float( + zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]); + acc += __bfloat162float(w[k]) * v; + } +#ifndef PF_CONV_NO_SILU + acc = acc / (1.f + __expf(-acc)); +#endif + xbc_act[(size_t)t * MB_CONV_DIM + c] = __float2bfloat16(acc); + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + const int units = n_seq * MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int seq = unit / (MB_HEADS * MB_HEAD_DIM); + const int rem = unit - seq * (MB_HEADS * MB_HEAD_DIM); + const int head = rem / MB_HEAD_DIM; + const int p = rem - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p * MB_STATE; + float sreg[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) sreg[j] = st[lane + j * 32]; + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const bf16 *xa = xbc_act + t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + const bf16 *dtrow = zxbcdt + t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + sreg[j] = sreg[j] * dA + __bfloat162float(Bsec[n]) * dt * xv; + acc += sreg[j] * __bfloat162float(Csec[n]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) st[lane + j * 32] = sreg[j]; + } + return; + } +#if PF_SSM_SPLIT_P + const int units = MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; unit += warps_total) { + const int head = unit / MB_HEAD_DIM; + const int p_only = unit - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; +#else + for (int head = blockIdx.x * WARPS + warp; head < MB_HEADS; + head += warps_total) { + const int group = head / MB_HEADS_PER_GROUP; +#endif + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *st = L.ssm_state + (size_t)slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE +#if PF_SSM_SPLIT_P + + (size_t)p_only * MB_STATE +#endif + ; + for (int t = 0; t < n_tokens; ++t) { + const bf16 *xa = xbc_act + (size_t)t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = + zxbcdt + (size_t)t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } +#if PF_SSM_SPLIT_P + { + const int p = p_only; +#else + for (int p = 0; p < MB_HEAD_DIM; ++p) { +#endif + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { +#if PF_SSM_SPLIT_P + const size_t at = (size_t)lane + j * 32; +#else + const size_t at = (size_t)p * MB_STATE + lane + j * 32; +#endif + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } + } +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32], wg[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + wg[j] = __bfloat162float(L.gnorm_w[idx]); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16(held[j] * inv * wg[j]); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len MK2_PROGRESS_PARAMETER(progress_mask)) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + constexpr int Q_PER_WARP = 3; + constexpr int GROUPS_PER_KV = 1; + const int pairs = BATCH * BB_KV_HEADS * GROUPS_PER_KV; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x + warp * gridDim.x; pair < pairs; + pair += gridDim.x * WARPS) { + const int row = pair / (BB_KV_HEADS * GROUPS_PER_KV); + MK2_PROGRESS_GUARD(progress_mask, row); + const int within_row = pair % (BB_KV_HEADS * GROUPS_PER_KV); + const int kv = within_row / GROUPS_PER_KV; + const int q_group = within_row % GROUPS_PER_KV; + const int q_first = q_group * Q_PER_WARP; + const int q_count = q_group == 0 ? Q_PER_WARP : 1; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (q_group == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[h][j] = 0.f; + } + float m[Q_PER_WARP], denom[Q_PER_WARP]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + m[h] = -CUDART_INF_F; + denom[h] = 0.f; + } + constexpr int PF = 8; + const int tpb = L.tokens_per_block; + const size_t elems_per_block = + (size_t)tpb * BB_KV_HEADS * BB_HEAD_DIM; + const int *bt_k = L.block_table + (size_t)(row * 2) * L.max_blocks_per_seq; + const int *bt_v = bt_k + L.max_blocks_per_seq; + const int head_base = kv * tpb * BB_HEAD_DIM; + float qr[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const bf16 *q = base + (kv * kv_per_q + q_first + h) * BB_HEAD_DIM; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qr[h][j] = __bfloat162float(q[lane + j * 32]); + } + for (int page0 = 0; page0 < T; page0 += tpb) { + const int k_page = bt_k[page0 / tpb] & 0x7FFFFFFF; + const int v_page = bt_v[page0 / tpb] & 0x7FFFFFFF; + const bf16 *kbase = L.pool_base + + (size_t)k_page * elems_per_block + + head_base; + const bf16 *vbase = L.pool_base + + (size_t)v_page * elems_per_block + + head_base; + const int tend = page0 + tpb < T ? page0 + tpb : T; + for (int t0 = page0; t0 < tend; t0 += PF) { + const int navail = tend - t0 < PF ? tend - t0 : PF; + float kr[PF][PER_LANE], vr[PF][PER_LANE]; + #pragma unroll + for (int u = 0; u < PF; ++u) { + const int tpos = t0 + (u < navail ? u : navail - 1); +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = vbase + (tpos - page0) * BB_HEAD_DIM; +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new : kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = cur_token + ? v_new : vbase + (tpos - page0) * BB_HEAD_DIM; +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + kr[u][j] = __bfloat162float(kt[i]); + vr[u][j] = __bfloat162float(vt[i]); + } + } + #pragma unroll + for (int u = 0; u < PF; ++u) { + if (u >= navail) break; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + dot += qr[h][j] * kr[u][j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m[h], s); + const float rescale = __expf(m[h] - m_new); + const float w = __expf(s - m_new); + denom[h] = denom[h] * rescale + w; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[h][j] = acc[h][j] * rescale + w * vr[u][j]; + m[h] = m_new; + } + } + } + } + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const float inv = 1.f / denom[h]; + const int head = kv * kv_per_q + q_first + h; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[h][j] * inv); + } + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * BB_KV_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / BB_KV_HEADS; + const int kv = idx - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv); + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv); + for (int i = 0; i < BB_HEAD_DIM; ++i) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + for (int s = 0; s <= pos; ++s) { + const bf16 *kt = kv_slot(L, seq, 0, s) + kv_offset(L, s, kv); + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + const bf16 *vt = kv_slot(L, seq, 1, s) + kv_offset(L, s, kv); + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + wt * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH, bf16 *row_stage = nullptr) { + __shared__ float score[BB_EXPERTS]; + constexpr int CH = D_MODEL / 8; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const char *gsrc = reinterpret_cast( + normed + (size_t)row * D_MODEL); + char *gdst = reinterpret_cast(row_stage); + const uint32_t act_s = (uint32_t)__cvta_generic_to_shared(gdst); + for (int i = threadIdx.x; i < CH; i += THREADS) { + const uint32_t dst = act_s + (uint32_t)i * 16; + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: + "r"(dst), "l"(gsrc + (size_t)i * 16) : "memory"); + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + if ((int)threadIdx.x < BB_EXPERTS) + moe_w[(size_t)row * BB_EXPERTS + threadIdx.x] = 0.f; + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); + __syncthreads(); + float acc0 = 0.f, acc1 = 0.f, acc2 = 0.f; + const bf16 *w0 = gate_w + (size_t)warp * D_MODEL; + const bf16 *w1 = gate_w + (size_t)(warp + WARPS) * D_MODEL; + const bf16 *w2 = gate_w + (size_t)(warp + 2 * WARPS) * D_MODEL; + for (int i = lane; i < D_MODEL; i += 128) { + bf16 sx0, sx1; + move::lds(sx0, act_s + (uint32_t)i * sizeof(bf16)); + move::lds(sx1, act_s + (uint32_t)(i + 32) * sizeof(bf16)); + float x0 = __bfloat162float(sx0); + float x1 = __bfloat162float(sx1); + acc0 += x0 * __bfloat162float(w0[i]); + acc1 += x0 * __bfloat162float(w1[i]); + acc2 += x0 * __bfloat162float(w2[i]); + acc0 += x1 * __bfloat162float(w0[i + 32]); + acc1 += x1 * __bfloat162float(w1[i + 32]); + acc2 += x1 * __bfloat162float(w2[i + 32]); + move::lds(sx0, act_s + (uint32_t)(i + 64) * sizeof(bf16)); + move::lds(sx1, act_s + (uint32_t)(i + 96) * sizeof(bf16)); + x0 = __bfloat162float(sx0); + x1 = __bfloat162float(sx1); + acc0 += x0 * __bfloat162float(w0[i + 64]); + acc1 += x0 * __bfloat162float(w1[i + 64]); + acc2 += x0 * __bfloat162float(w2[i + 64]); + acc0 += x1 * __bfloat162float(w0[i + 96]); + acc1 += x1 * __bfloat162float(w1[i + 96]); + acc2 += x1 * __bfloat162float(w2[i + 96]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + acc0 += __shfl_xor_sync(0xffffffffu, acc0, off); + acc1 += __shfl_xor_sync(0xffffffffu, acc1, off); + acc2 += __shfl_xor_sync(0xffffffffu, acc2, off); + } + if (lane == 0) { + score[warp] = 1.f / (1.f + __expf(-acc0)); + score[warp + WARPS] = 1.f / (1.f + __expf(-acc1)); + score[warp + 2 * WARPS] = 1.f / (1.f + __expf(-acc2)); + } + __syncthreads(); + if (threadIdx.x < 32) { + float remaining = (lane < BB_EXPERTS) + ? score[lane] + __bfloat162float(gate_bias[lane]) + : -CUDART_INF_F; + int pick[BB_TOP_K]; + float total = 0.f; + #pragma unroll + for (int k = 0; k < BB_TOP_K; ++k) { + float bv = remaining; + int bi = lane; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffff, bv, off); + const int oi = __shfl_xor_sync(0xffffffff, bi, off); + if (ov > bv || (ov == bv && oi < bi)) { bv = ov; bi = oi; } + } + pick[k] = bi; + total += score[bi]; + if (lane == bi) remaining = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + #pragma unroll + for (int k = 0; k < BB_TOP_K; ++k) + if (lane == pick[k]) moe_w[row * BB_EXPERTS + lane] = score[lane] * norm; + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert <= BB_EXPERTS; ++expert) + value = __fadd_rn(value, parts[base + (size_t)expert * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ bf16 bb_moe_fold_selected( + const float *parts, const float *weights, int idx) { + const int row = idx / D_MODEL; + const size_t base = (size_t)row * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert < BB_EXPERTS; ++expert) { + const float part = weights[row * BB_EXPERTS + expert] != 0.f + ? parts[base + (size_t)expert * D_MODEL] : 0.f; + value = __fadd_rn(value, part); + } + value = __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ __forceinline__ void sample_argmax_merge( + float &value, int &rank, float other, int other_rank) { + if (other > value || (other == value && other_rank < rank)) { + value = other; + rank = other_rank; + } +} +__device__ __forceinline__ int phoneme_warp_argmax(const bf16 *lg) { + const int lane = threadIdx.x & 31; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = lane; n < BB_PHONEME; n += 32) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, best, off); + const int other_arg = __shfl_down_sync(0xffffffffu, arg, off); + sample_argmax_merge(best, arg, other, other_arg); + } + return arg; +} +__device__ __forceinline__ int sample_block_argmax(float value, int rank) { + __shared__ float warp_value[WARPS]; + __shared__ int warp_rank[WARPS]; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + if (lane == 0) { warp_value[warp] = value; warp_rank[warp] = rank; } + __syncthreads(); + if (warp == 0) { + value = lane < WARPS ? warp_value[lane] : -CUDART_INF_F; + rank = lane < WARPS ? warp_rank[lane] : INT_MAX; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + } + return rank; +} +__device__ __forceinline__ int sample_eos_insertion_rank( + const uint64_t *keys, float eos) { + int lo = 0, hi = TOP_K; + while (lo < hi) { + const int mid = (lo + hi) >> 1; + if (eos > ordered_to_float((uint32_t)(keys[mid] >> 32))) hi = mid; + else lo = mid + 1; + } + return lo; +} +__device__ __forceinline__ void sample_topk_serial( + const uint64_t *keys, float eos, const float *noise, int *out) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float value; + int code; + if (!eos_used && eos > av) { + value = eos; + code = AUDIO_EOS; + eos_used = true; + } else { + value = av; + code = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = value + noise[j]; + if (noised > best_noised) { + best_noised = noised; + best_code = code; + } + } + *out = best_code; +} +__device__ __forceinline__ void sample_topk_parallel( + const uint64_t *keys, float eos, const float *noise, int *out) { + const int j = threadIdx.x; + const bool sorted_nan = j < TOP_K && + isnan(ordered_to_float((uint32_t)(keys[j] >> 32))); + if (__syncthreads_or(sorted_nan)) { + if (threadIdx.x == 0) sample_topk_serial(keys, eos, noise, out); + return; + } + const int lane = threadIdx.x & 31; + int insertion = lane == 0 ? sample_eos_insertion_rank(keys, eos) : 0; + insertion = __shfl_sync(0xffffffffu, insertion, 0); + float value = -CUDART_INF_F; + int rank = INT_MAX; + if (j < TOP_K) { + const int a = j - (j > insertion); + const float v = j == insertion + ? eos : ordered_to_float((uint32_t)(keys[a] >> 32)); + const float noised = v + noise[j]; + if (noised > -CUDART_INF_F) { value = noised; rank = j; } + } + const int winner = sample_block_argmax(value, rank); + if (threadIdx.x == 0) { + if (winner == INT_MAX) *out = 0; + else if (winner == insertion) *out = AUDIO_EOS; + else { + const int a = winner - (winner > insertion); + *out = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + } + } +} +constexpr int SORT_ITEMS_PER_THREAD = SORT_N / THREADS; +struct sampler_storage { + uint64_t keys[2][SORT_N]; +}; +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + sampler_storage &storage, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + uint64_t thread_keys[SORT_ITEMS_PER_THREAD]; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + thread_keys[item] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + int cross_stage = 0; + for (int k = 2; k <= SORT_N; k <<= 1) { + for (int j = k >> 1; j > 0; j >>= 1) { + if (j >= 32) { + const int bank = cross_stage++ & 1; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + storage.keys[bank][n] = thread_keys[item]; + } + __syncthreads(); + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + const uint64_t other = storage.keys[bank][n ^ j]; + const bool take_hi = ((n & k) == 0) == ((n & j) == 0); + const uint64_t mine = thread_keys[item]; + thread_keys[item] = take_hi + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } else { + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + const uint64_t other = __shfl_xor_sync( + 0xffffffffu, thread_keys[item], j); + const bool take_hi = ((n & k) == 0) == ((n & j) == 0); + const uint64_t mine = thread_keys[item]; + thread_keys[item] = take_hi + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } + } + } + const int top_bank = cross_stage & 1; + if (threadIdx.x < TOP_K) + storage.keys[top_bank][threadIdx.x] = thread_keys[0]; + __syncthreads(); + sample_topk_parallel( + storage.keys[top_bank], eos_v, + gumbel + ((size_t)row * N_POS + pos) * TOP_K, + out_codes + row); + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void r6_sparse_map(const float* routes,int* map,int* counts) { + __shared__ int r7_warp_counts[THREADS/32]; + const int lane=threadIdx.x%32, warp=threadIdx.x/32; + const unsigned lower=(1u<>1; + for(int e=0;e(packed)[(size_t)e*BATCH*row_vectors+i]= + row<0?make_uint4(0,0,0,0):reinterpret_cast(source)[(size_t)row*row_vectors+i%row_vectors]; + } + } +} +__device__ void r6_tile(int work,int output_tiles,const int* lut,int &tile,int &tok, + int mtg) { + const int groups=output_tiles/mtg; + tok=lut[1+work/groups]; + const int e=tok/DECODE_TILES; + tile=e*output_tiles+(work%groups)*mtg; +} +__device__ __forceinline__ void mk2_group_tile( + int work, int grouped, int &tile, int &tok, int &here) { + here = work < grouped ? NM_MAX : 1; + const int unit = work < grouped ? work * NM_MAX : grouped * (NM_MAX - 1) + work; + tile = unit / DECODE_TILES; + tok = unit - tile * DECODE_TILES; +} +__device__ __forceinline__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&p_s)[STAGES][SLOTS], out_tile &o_s, + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + int &tile_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = DECODE_TILES, + int rows_valid = DECODE_ROWS, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const int *r6_lut = nullptr,const int *r6_map = nullptr) { + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + constexpr int QSLOT = MT + NM_MAX; + constexpr bool WIDE = MT > 1 && NM_MAX > 1 && STAGES * SLOTS >= 2 * QSLOT; + constexpr int PSTAGES = WIDE ? 2 : STAGES; + constexpr int PSLOT = WIDE ? QSLOT : SLOTS; + constexpr int ACCN = WIDE ? MT * NM_MAX : NM_MAX; + constexpr int HEX_MT = 3, HEX_NM = 2; + constexpr int HEX_SLOT = HEX_MT + HEX_NM; + auto epi_publish = [](out_tile &st, const rt_fl<16, BATCH_PAD> &src) { + using acc_t = rt_fl<16, BATCH_PAD>; + static_assert(acc_t::height == 1, "one row block per warp"); + static_assert(out_tile::rows == BATCH_PAD && out_tile::cols == OUT_ATOM, + "staging tile must be square for an in-place transpose"); + const int t = ::kittens::laneid(); + const uint32_t base = + static_cast(__cvta_generic_to_shared(&st.data[0])); + const int row0 = 8 * (t >> 4) + (t & 7); + const int col0 = 16 * warpgroup::warpid() + 8 * ((t >> 3) & 1); + #pragma unroll + for (int j = 0; j < acc_t::width; ++j) { + bf16_2 q0 = base_types::convertor::convert( + src.tiles[0][j].data[0]); + bf16_2 q1 = base_types::convertor::convert( + src.tiles[0][j].data[1]); + bf16_2 q2 = base_types::convertor::convert( + src.tiles[0][j].data[2]); + bf16_2 q3 = base_types::convertor::convert( + src.tiles[0][j].data[3]); + move::stsm4t( + out_tile::idx(base, {16 * j + row0, col0}), q0, q1, q2, q3); + } + }; + weight_tile *const pf = &p_s[0][0]; + const bool arena10 = + reinterpret_cast(&o_sf) == + reinterpret_cast(pf + STAGES * SLOTS); + const bool two_d = !r6_lut && out_tiles_per_expert == 0 + && n_tok == DECODE_TILES && DECODE_TILES % NM_MAX == 0; + int n_tiles = compact_tiles ? __popc(active_mask)*compact_tiles : n_out/OUT_ATOM; + constexpr int QTOK = DECODE_TILES / NM_MAX; + const int out_tiles = n_tiles; + const bool tall3 = WIDE && two_d && out_tiles == 24; + const bool split72 = WIDE && two_d && DECODE_TILES == 16 && out_tiles == 72; + const bool split81 = WIDE && two_d && DECODE_TILES == 16 && out_tiles == 81; + const bool split17 = WIDE && two_d && DECODE_TILES == 16 && out_tiles == 17; + const bool mixed40 = WIDE && two_d && DECODE_TILES == 16 && out_tiles == 40; + const int n_tiles_hex_all = (out_tiles / HEX_MT) * (DECODE_TILES / HEX_NM); + const bool hex = WIDE && two_d && arena10 && DECODE_TILES == 16 && DECODE_TILES % HEX_NM == 0 && + (out_tiles == 96 || out_tiles == 81 || out_tiles == 72 || + out_tiles == 40) && + (epi == EPI_STORE || epi == EPI_GELU) && !r6_lut && + out_tiles_per_expert == 0; + const int hex_tok = DECODE_TILES / HEX_NM; + const int hex_wide = (out_tiles == 72 || out_tiles == 81) + ? 128 : n_tiles_hex_all; + const int pslot = hex ? HEX_SLOT : PSLOT; + int q_units = 0, q_rem_base = 0; + int grouped = 0; + if constexpr (WIDE) { + if (two_d) { + q_units = (out_tiles / MT) * QTOK; + q_rem_base = out_tiles - out_tiles % MT; + } + } else if (two_d) { + const int g = (int)gridDim.x; + const int units = n_tiles * DECODE_TILES; + const int rounds = (units + g - 1) / g; + const int cap = units / NM_MAX; + grouped = g * (rounds / NM_MAX); + if (grouped > cap) grouped = cap; + } + const int r6_output_tiles = r6_lut ? n_out/BB_EXPERTS/OUT_ATOM : 0; + const int r6_mtg = (MT > 1 && r6_lut && r6_output_tiles % MT == 0) ? MT : 1; + const int r6_groups = r6_lut ? r6_output_tiles/r6_mtg : 0; + const int r6_rows = r6_lut ? (r6_lut[0] & 0xFFF) : 0; + const int r6_pairs = (r6_lut && WIDE && r6_mtg == MT) ? (r6_lut[0] >> 12) : 0; + const int r6_pair_items = r6_pairs * r6_groups; + const int r6_single_mtg = (r6_lut && WIDE && r6_output_tiles % 3 == 0) + ? 3 : r6_mtg; + const int r6_single_groups = r6_lut ? r6_output_tiles/r6_single_mtg : 0; + const int r6_single_items = r6_lut + ? (r6_rows - 2 * r6_pairs) * r6_single_groups : 0; + if(r6_lut) { + n_tiles=r6_pair_items+r6_single_items; + n_tok=1; + } else if(two_d) { + if constexpr (WIDE) { + if (hex) n_tiles = (out_tiles == 72) ? 256 + : (out_tiles == 81) ? 264 + : (out_tiles == 40) + ? n_tiles_hex_all + hex_tok + : n_tiles_hex_all; + else if (mixed40) n_tiles = 244; + else if (split72) n_tiles = 360; + else if (split81) n_tiles = 384; + else if (split17) n_tiles = 132; + else if (tall3) n_tiles = (out_tiles / 3) * DECODE_TILES; + else n_tiles = q_units + (out_tiles % MT) * QTOK; + } + else n_tiles = grouped + (n_tiles*DECODE_TILES - grouped*NM_MAX); + n_tok=1; } + const int r6_extra = r6_lut ? n_tiles % (int)gridDim.x : 0; + const int r6_skew = r6_lut + ? (r6_single_items < r6_extra ? r6_single_items : r6_extra) : 0; + const int iters = k_total / KCHUNK; + int tile_start = (int)blockIdx.x - tile_io % (int)gridDim.x; + if (tile_start < 0) tile_start += (int)gridDim.x; + tile_io += n_tiles; + int tiles_here = 0; + for (int t = tile_start; t < n_tiles; t += gridDim.x) ++tiles_here; + const int iter_base = iter_io; + iter_io += tiles_here * iters * n_tok; + const int my_tiles = (n_tiles - tile_start + gridDim.x - 1) / gridDim.x; + if (warpgroup::groupid() == 1) { + if (hex) { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) { + if (out_tiles == 81 && work >= hex_wide) { + const int j = work - hex_wide; + const bool pair = j >= 128; + const int nt = pair ? 80 : 48 + 2 * (j / hex_tok); + const int ntok = HEX_NM * (j % hex_tok); + if (pair) { + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + wait(drained[s], (iter / PSTAGES) & 1); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], sizeof(weight_tile) + + HEX_NM * sizeof(act_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot], w, {nt, i}, filled[s]); + #pragma unroll + for (int m = 0; m < HEX_NM; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + 1 + m], a, + {ntok + m, i}, filled[s]); + } + } + } else { + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + wait(drained[s], (iter / PSTAGES) & 1); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], 2 * sizeof(weight_tile) + + HEX_NM * sizeof(act_tile)); + #pragma unroll + for (int c = 0; c < 2; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {nt + c, i}, filled[s]); + #pragma unroll + for (int m = 0; m < HEX_NM; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + 2 + m], a, + {ntok + m, i}, filled[s]); + } + } + } + continue; + } + if (out_tiles == 40 && work >= hex_wide) { + const int ntok = HEX_NM * (work - hex_wide); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + wait(drained[s], (iter / PSTAGES) & 1); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], sizeof(weight_tile) + + HEX_NM * sizeof(act_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot], w, {out_tiles - 1, i}, filled[s]); + #pragma unroll + for (int m = 0; m < HEX_NM; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + 1 + m], a, + {ntok + m, i}, filled[s]); + } + } + continue; + } + if (work >= hex_wide) { + const int j = work - hex_wide; + const int nt = HEX_MT * ((hex_wide + (j >> 1)) / hex_tok); + const int ntok = HEX_NM * ((hex_wide + (j >> 1)) % hex_tok) + (j & 1); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + wait(drained[s], (iter / PSTAGES) & 1); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], HEX_MT * sizeof(weight_tile) + + sizeof(act_tile)); + #pragma unroll + for (int c = 0; c < HEX_MT; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {nt + c, i}, filled[s]); + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + HEX_MT], a, {ntok, i}, filled[s]); + } + } + continue; + } + const int tile_id = HEX_MT * (work / hex_tok); + const int tok = HEX_NM * (work % hex_tok); + rt_fl<16, BATCH_PAD> pacc[HEX_NM]; + #pragma unroll + for (int u = 0; u < HEX_NM; ++u) warp::zero(pacc[u]); + int pend = -1, pend_ph = 0; + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(drained[s], phase); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], HEX_MT * sizeof(weight_tile) + + HEX_NM * sizeof(act_tile)); + #pragma unroll + for (int c = 0; c < HEX_MT; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {tile_id + c, i}, filled[s]); + #pragma unroll + for (int m = 0; m < HEX_NM; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + HEX_MT + m], a, {tok + m, i}, filled[s]); + } + if (pend >= 0) { { + wait(filled[pend], pend_ph); + warpgroup::mma_ABt(pacc[0], pf[pend * pslot + HEX_MT - 1], + pf[pend * pslot + HEX_MT]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + pacc[1], pf[pend * pslot + HEX_MT - 1], + pf[pend * pslot + HEX_MT + 1]); + warpgroup::mma_async_wait(); + } } + pend = s; pend_ph = phase; + } + { + wait(filled[pend], pend_ph); + warpgroup::mma_ABt(pacc[0], pf[pend * pslot + HEX_MT - 1], + pf[pend * pslot + HEX_MT]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + pacc[1], pf[pend * pslot + HEX_MT - 1], + pf[pend * pslot + HEX_MT + 1]); + warpgroup::mma_async_wait(); + } + const int es = iter % PSTAGES; + wait(drained[es], (iter / PSTAGES) & 1); + out_tile_f &q_sf = + *reinterpret_cast(&pf[es * pslot]); + out_tile &q_s = *reinterpret_cast(&q_sf); + const int qbase = (tile_id + HEX_MT - 1) * OUT_ATOM; + constexpr int OUT_OCTS = OUT_ATOM / 8; + #pragma unroll + for (int u = 0; u < HEX_NM; ++u) { + const int grow0 = (tok + u) * BATCH_PAD; + epi_publish(q_s, pacc[u]); + warpgroup::sync(3); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); + idx < OUT_OCTS * rows_valid; idx += 128) { + const int row = idx >> 3, col = (idx & 7) << 3; + const size_t at = + (size_t)(grow0 + row) * dest_stride + qbase + col; + const bf16_oct q = *reinterpret_cast( + &q_s[int2{row, col}]); + if (epi == EPI_STORE) { + *reinterpret_cast(dest + at) = q; + } else { + *reinterpret_cast(dest + at) = bf16_oct{ + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x0))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x1))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x2))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x3))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x4))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x5))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x6))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x7)))}; + } + } + warpgroup::sync(3); + } + } + } else if (warpgroup::laneid() == 0) { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) { + int tile_id = mk2_expert_tile(work, active_mask, compact_tiles), r6_tok=0; + int nm = 1, mt = 1; + if(r6_lut) { + int rwork = work - r6_skew; + if (rwork < 0) rwork += n_tiles; + const bool pair = rwork < r6_pair_items; + const int q = pair ? rwork / r6_groups + : (rwork-r6_pair_items) / r6_single_groups; + const int group = pair ? rwork-q*r6_groups + : rwork-r6_pair_items-q*r6_single_groups; + const int base_item = pair ? 2*q : 2*r6_pairs+q; + const int ltok = r6_lut[1 + base_item]; + r6_tok = ltok; + tile_id = (ltok / DECODE_TILES) * r6_output_tiles + + group * (pair ? r6_mtg : r6_single_mtg); + mt = pair ? r6_mtg : r6_single_mtg; + nm = pair ? NM_MAX : 1; + } + else if(two_d) { + if constexpr (WIDE) { + if (mixed40) { + const bool quad = work < 132; + const int item = quad ? work : work - 132; + if (quad && item < 84) { + tile_id = 2 * (item / 7); + r6_tok = 2 * (item % 7); + } else if (quad) { + const int tail = item - 84; + tile_id = 2 * (12 + tail / 6); + r6_tok = 2 * (tail % 6); + } else if (item < 48) { + tile_id = 2 * (item / 4) + (item % 4) / 2; + r6_tok = 14 + item % 2; + } else { + const int tail = item - 48; + tile_id = 2 * (12 + tail / 8) + (tail % 8) / 4; + r6_tok = 12 + tail % 4; + } + mt = quad ? MT : 1; + nm = quad ? NM_MAX : 1; + } else if (split72) { + if (work < 264) { + tile_id = 2 * (work / 8); + r6_tok = 2 * (work % 8); + mt = MT; nm = NM_MAX; + } else { + const int i = work - 264, c = i >> 2, q = i & 3; + tile_id = 2 * (33 + (c >> 3)) + (q >> 1); + r6_tok = 2 * (c & 7) + (q & 1); + mt = 1; nm = 1; + } + } else if (split81) { + if (work < 264) { + tile_id = 2 * (work / 8); + r6_tok = 2 * (work % 8); + mt = MT; nm = NM_MAX; + } else if (work < 376) { + const int i = work - 264; + tile_id = 2 * (33 + i / 16); + r6_tok = i % 16; + mt = MT; nm = 1; + } else { + tile_id = 80; + r6_tok = 2 * (work - 376); + mt = 1; nm = NM_MAX; + } + } else if (split17) { + if (work < 8) { + tile_id = 14; + r6_tok = work; + mt = 3; nm = 1; + } else if (work < 120) { + const int i = work - 8; + tile_id = i >> 3; + r6_tok = 2 * (i & 7); + mt = 1; nm = NM_MAX; + } else { + const int i = work - 120; + tile_id = 14 + (i >> 2); + r6_tok = 8 + 2 * (i & 3); + mt = 1; nm = NM_MAX; + } + } else if (tall3) { + tile_id = 3 * (work / DECODE_TILES); + r6_tok = work % DECODE_TILES; + mt = 3; + nm = 1; + } else { + const int j = work < q_units ? work : work - q_units; + tile_id = work < q_units ? MT * (j / QTOK) + : q_rem_base + j / QTOK; + r6_tok = NM_MAX * (j % QTOK); + mt = work < q_units ? MT : 1; + nm = NM_MAX; + } + } else mk2_group_tile(work,grouped,tile_id,r6_tok,nm); + } + for (int token_iteration = 0; token_iteration < n_tok; ++token_iteration) + for (int i = 0; i < iters; ++i, ++iter) { + const int tok = (r6_lut || two_d) ? r6_tok : token_iteration; + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(drained[s], phase); + tma::expect_bytes(filled[s], + mt * sizeof(weight_tile) + nm * sizeof(act_tile)); + if (mt == 1) { + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot], w, {tile_id, i}, filled[s]); + } else if (mt == 3) { + #pragma unroll + for (int c = 0; c < 3; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {tile_id + c, i}, filled[s]); + } else { + #pragma unroll + for (int c = 0; c < MT; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {tile_id + c, i}, filled[s]); + } + const int a_col = out_tiles_per_expert && !r6_lut + ? (tile_id / out_tiles_per_expert) * act_k_tiles_per_expert + i + : i; + if (NM_MAX == 1 || nm == NM_MAX) { + #pragma unroll + for (int m = 0; m < NM_MAX; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + mt + m], a, {tok + m, a_col}, filled[s]); + } else { + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + mt], a, {tok, a_col}, filled[s]); + } + } + } + } + } else { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) + for (int token_iteration = 0; token_iteration < n_tok; ++token_iteration) { + int tile_id=mk2_expert_tile(work, active_mask, compact_tiles),tok=token_iteration; + int nm = 1, mt = 1; + if(r6_lut) { + int rwork = work - r6_skew; + if (rwork < 0) rwork += n_tiles; + const bool pair = rwork < r6_pair_items; + const int q = pair ? rwork / r6_groups + : (rwork-r6_pair_items) / r6_single_groups; + const int group = pair ? rwork-q*r6_groups + : rwork-r6_pair_items-q*r6_single_groups; + const int base_item = pair ? 2*q : 2*r6_pairs+q; + const int ltok = r6_lut[1 + base_item]; + tok = ltok; + tile_id = (ltok / DECODE_TILES) * r6_output_tiles + + group * (pair ? r6_mtg : r6_single_mtg); + mt = pair ? r6_mtg : r6_single_mtg; + nm = pair ? NM_MAX : 1; + } + else if(two_d) { + if constexpr (WIDE) { + if (hex) { + if (out_tiles == 81 && work >= hex_wide) { + const int j = work - hex_wide; + const bool pair = j >= 128; + tile_id = pair ? 80 : 48 + 2 * (j / hex_tok); + tok = HEX_NM * (j % hex_tok); + mt = pair ? 1 : 2; + nm = HEX_NM; + } else if (out_tiles == 40 && work >= hex_wide) { + tile_id = out_tiles - 1; + tok = HEX_NM * (work - hex_wide); + mt = 1; + nm = HEX_NM; + } else { + const int hw = work < hex_wide + ? work : hex_wide + ((work - hex_wide) >> 1); + tile_id = HEX_MT * (hw / hex_tok); + tok = HEX_NM * (hw % hex_tok) + + (work < hex_wide ? 0 : ((work - hex_wide) & 1)); + mt = HEX_MT; nm = work < hex_wide ? HEX_NM : 1; + } + } else if (mixed40) { + const bool quad = work < 132; + const int item = quad ? work : work - 132; + if (quad && item < 84) { + tile_id = 2 * (item / 7); + tok = 2 * (item % 7); + } else if (quad) { + const int tail = item - 84; + tile_id = 2 * (12 + tail / 6); + tok = 2 * (tail % 6); + } else if (item < 48) { + tile_id = 2 * (item / 4) + (item % 4) / 2; + tok = 14 + item % 2; + } else { + const int tail = item - 48; + tile_id = 2 * (12 + tail / 8) + (tail % 8) / 4; + tok = 12 + tail % 4; + } + mt = quad ? MT : 1; + nm = quad ? NM_MAX : 1; + } else if (split72) { + if (work < 264) { + tile_id = 2 * (work / 8); + tok = 2 * (work % 8); + mt = MT; nm = NM_MAX; + } else { + const int i = work - 264, c = i >> 2, q = i & 3; + tile_id = 2 * (33 + (c >> 3)) + (q >> 1); + tok = 2 * (c & 7) + (q & 1); + mt = 1; nm = 1; + } + } else if (split81) { + if (work < 264) { + tile_id = 2 * (work / 8); + tok = 2 * (work % 8); + mt = MT; nm = NM_MAX; + } else if (work < 376) { + const int i = work - 264; + tile_id = 2 * (33 + i / 16); + tok = i % 16; + mt = MT; nm = 1; + } else { + tile_id = 80; + tok = 2 * (work - 376); + mt = 1; nm = NM_MAX; + } + } else if (split17) { + if (work < 8) { + tile_id = 14; + tok = work; + mt = 3; nm = 1; + } else if (work < 120) { + const int i = work - 8; + tile_id = i >> 3; + tok = 2 * (i & 7); + mt = 1; nm = NM_MAX; + } else { + const int i = work - 120; + tile_id = 14 + (i >> 2); + tok = 8 + 2 * (i & 3); + mt = 1; nm = NM_MAX; + } + } else if (tall3) { + tile_id = 3 * (work / DECODE_TILES); + tok = work % DECODE_TILES; + mt = 3; + nm = 1; + } else { + const int j = work < q_units ? work : work - q_units; + tile_id = work < q_units ? MT * (j / QTOK) + : q_rem_base + j / QTOK; + tok = NM_MAX * (j % QTOK); + mt = work < q_units ? MT : 1; + nm = NM_MAX; + } + } else mk2_group_tile(work,grouped,tile_id,tok,nm); + } + int epi_hold = -1; + rt_fl<16, BATCH_PAD> accum[ACCN]; + if (hex && nm == 1) { + #pragma unroll + for (int u = 0; u < 3; ++u) warp::zero(accum[u]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + HEX_MT]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[1], pf[s * pslot + 1], pf[s * pslot + HEX_MT]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[2], pf[s * pslot + 2], pf[s * pslot + HEX_MT]); + warpgroup::mma_async_wait(); + if (i + 1 < iters) { + if (warpgroup::laneid() == 0) + arrive(drained[s], CONSUMER_WARPS); + } else epi_hold = s; + } + } else if (hex && mt == HEX_MT) { + #pragma unroll + for (int u = 0; u < ACCN; ++u) warp::zero(accum[u]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], + pf[s * pslot + HEX_MT]); + #pragma unroll + for (int u = 1; u < ACCN; ++u) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[u], pf[s * pslot + u / HEX_NM], + pf[s * pslot + HEX_MT + u % HEX_NM]); + warpgroup::mma_async_wait(); + if (i + 1 < iters) { + if (warpgroup::laneid() == 0) + arrive(drained[s], CONSUMER_WARPS); + } else epi_hold = s; + } + } else if (WIDE && mt == 3) { + #pragma unroll + for (int u = 0; u < 3; ++u) warp::zero(accum[u]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + 3]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[1], pf[s * pslot + 1], pf[s * pslot + 3]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[2], pf[s * pslot + 2], pf[s * pslot + 3]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else if (WIDE && mt == MT && nm == NM_MAX) { + #pragma unroll + for (int u = 0; u < ACCN; ++u) warp::zero(accum[u]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + MT]); + #pragma unroll + for (int u = 1; u < ACCN; ++u) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[u], pf[s * pslot + u / NM_MAX], + pf[s * pslot + MT + u % NM_MAX]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else if (MT > 1 && mt == MT) { + #pragma unroll + for (int c = 0; c < MT; ++c) warp::zero(accum[c]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + MT]); + #pragma unroll + for (int c = 1; c < MT; ++c) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[c], pf[s * pslot + c], pf[s * pslot + MT]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else if (NM_MAX == 1 || nm == NM_MAX) { + #pragma unroll + for (int m = 0; m < NM_MAX; ++m) warp::zero(accum[m]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + 1]); + #pragma unroll + for (int m = 1; m < NM_MAX; ++m) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[m], pf[s * pslot], pf[s * pslot + 1 + m]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else { + warp::zero(accum[0]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + 1]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } + out_tile_f &e_sf = epi_hold >= 0 + ? *reinterpret_cast(&pf[epi_hold * pslot]) + : o_sf; + out_tile *const e_sp = reinterpret_cast(&e_sf); + const int nacc = mt * nm; + if (epi != EPI_BIAS) { + epi_publish(e_sp[0], accum[0]); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + } + #pragma unroll + for (int u = 0; u < ACCN; ++u) { + if (u >= mt * nm) break; + const int c = (mt > 1) ? ((nm > 1) ? (u / NM_MAX) : u) : 0; + const int m = (nm > 1) ? (u % NM_MAX) : 0; + const int base = (r6_lut ? (tile_id + c)%r6_output_tiles + : (tile_id + c)) * OUT_ATOM; + const int grow0 = (tok + m) * BATCH_PAD; + if (epi == EPI_BIAS) { + const int bw = 16 * warpgroup::warpid() + + (::kittens::laneid() >> 2); + const float b0 = __bfloat162float(bias[base + bw]); + const float b1 = __bfloat162float(bias[base + bw + 8]); + #pragma unroll + for (int j = 0; j < rt_fl<16, BATCH_PAD>::width; ++j) { + accum[u].tiles[0][j].data[0].x += b0; + accum[u].tiles[0][j].data[0].y += b0; + accum[u].tiles[0][j].data[1].x += b1; + accum[u].tiles[0][j].data[1].y += b1; + accum[u].tiles[0][j].data[2].x += b0; + accum[u].tiles[0][j].data[2].y += b0; + accum[u].tiles[0][j].data[3].x += b1; + accum[u].tiles[0][j].data[3].y += b1; + } + epi_publish(e_sp[0], accum[u]); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + constexpr int OUT_OCTS = OUT_ATOM / 8; + for (int idx = warpgroup::laneid(); idx < OUT_OCTS * rows_valid; + idx += 128) { + const int row = idx >> 3, col = (idx & 7) << 3; + const size_t at = (size_t)(grow0 + row) * dest_stride + base + col; + *reinterpret_cast(dest + at) = + *reinterpret_cast( + &e_sp[0][int2{row, col}]); + } + warpgroup::sync(1); + continue; + } + if (u + 1 < ACCN && u + 1 < nacc) + epi_publish(e_sp[(u + 1) & 1], + accum[u + 1 < ACCN ? u + 1 : 0]); + const out_tile &e_sr = e_sp[u & 1]; + constexpr int OUT_OCTS = OUT_ATOM / 8; + for (int idx = warpgroup::laneid(); idx < OUT_OCTS * rows_valid; + idx += 128) { + const int row = idx >> 3; + const int col = (idx & 7) << 3; + const int grow = grow0 + row; + const size_t at = + (size_t)grow * dest_stride + base + col; + const bf16_oct src = *reinterpret_cast( + &e_sr[int2{row, col}]); + if (epi == EPI_STORE) { + *reinterpret_cast(dest + at) = src; + continue; + } + const float v0 = __bfloat162float(src.x0); + const float v1 = __bfloat162float(src.x1); + const float v2 = __bfloat162float(src.x2); + const float v3 = __bfloat162float(src.x3); + const float v4 = __bfloat162float(src.x4); + const float v5 = __bfloat162float(src.x5); + const float v6 = __bfloat162float(src.x6); + const float v7 = __bfloat162float(src.x7); + switch (epi) { + case EPI_STORE: break; + case EPI_SILU: + *reinterpret_cast(dest + at) = bf16_oct{ + __float2bfloat16(v0 / (1.f + __expf(-v0))), + __float2bfloat16(v1 / (1.f + __expf(-v1))), + __float2bfloat16(v2 / (1.f + __expf(-v2))), + __float2bfloat16(v3 / (1.f + __expf(-v3))), + __float2bfloat16(v4 / (1.f + __expf(-v4))), + __float2bfloat16(v5 / (1.f + __expf(-v5))), + __float2bfloat16(v6 / (1.f + __expf(-v6))), + __float2bfloat16(v7 / (1.f + __expf(-v7)))}; + break; + case EPI_MOE_ACC: { + const int expert = (tile_id + c) / out_tiles_per_expert; + const int local = ((tile_id + c) % out_tiles_per_expert) * OUT_ATOM + + col; + const int original_row=r6_map?r6_map[grow]:grow; + if(original_row<0) break; + const float wt = moe_w[original_row * BB_EXPERTS + expert]; + if (wt != 0.f) { + const size_t dst = + ((size_t)original_row * (BB_EXPERTS + 1) + expert) * + D_MODEL + local; + *reinterpret_cast(moe_out + dst) = make_float4( + __fmul_rn(wt, v0), __fmul_rn(wt, v1), + __fmul_rn(wt, v2), __fmul_rn(wt, v3)); + *reinterpret_cast(moe_out + dst + 4) = make_float4( + __fmul_rn(wt, v4), __fmul_rn(wt, v5), + __fmul_rn(wt, v6), __fmul_rn(wt, v7)); + } + break; + } + case EPI_ACC_F32: + *reinterpret_cast(moe_out + + ((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * + D_MODEL + base + col) = make_float4(v0, v1, v2, v3); + *reinterpret_cast(moe_out + + ((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * + D_MODEL + base + col + 4) = make_float4(v4, v5, v6, v7); + break; + case EPI_BIAS: break; + case EPI_GELU: + *reinterpret_cast(dest + at) = bf16_oct{ + __float2bfloat16(gelu_tanh_fp32(v0)), + __float2bfloat16(gelu_tanh_fp32(v1)), + __float2bfloat16(gelu_tanh_fp32(v2)), + __float2bfloat16(gelu_tanh_fp32(v3)), + __float2bfloat16(gelu_tanh_fp32(v4)), + __float2bfloat16(gelu_tanh_fp32(v5)), + __float2bfloat16(gelu_tanh_fp32(v6)), + __float2bfloat16(gelu_tanh_fp32(v7))}; + break; + case EPI_RESIDUAL: { + const bf16_oct prior = + *reinterpret_cast(residual + at); + *reinterpret_cast(residual + at) = bf16_oct{ + __float2bfloat16(__bfloat162float(prior.x0) + v0), + __float2bfloat16(__bfloat162float(prior.x1) + v1), + __float2bfloat16(__bfloat162float(prior.x2) + v2), + __float2bfloat16(__bfloat162float(prior.x3) + v3), + __float2bfloat16(__bfloat162float(prior.x4) + v4), + __float2bfloat16(__bfloat162float(prior.x5) + v5), + __float2bfloat16(__bfloat162float(prior.x6) + v6), + __float2bfloat16(__bfloat162float(prior.x7) + v7)}; + break; + } + } + } + if (u + 1 < nacc) { + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + } + } + if (epi != EPI_BIAS) warpgroup::sync(1); + if (epi_hold >= 0 && warpgroup::laneid() == 0) + arrive(drained[epi_hold], CONSUMER_WARPS); + } + (void)my_tiles; + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr, + const unsigned char *native_predictor_active_mask = nullptr + MK2_PROGRESS_PARAMETER(native_progress_mask)) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&p_s)[STAGES][SLOTS] = alloc.allocate(); + out_tile_f (&o_sf) = alloc.allocate(); + out_tile &o_s = *reinterpret_cast(&o_sf); + bf16 *const route_stage = reinterpret_cast(&o_sf); + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_SCRATCH_STRIDE]; + static_assert(STAGES * SLOTS * sizeof(weight_tile) >= sizeof(sampler_storage)); + sampler_storage &sort_storage = *reinterpret_cast(&__shm[0]); + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS, 0); + arrive(drained[s], CONSUMER_WARPS); + } + } + int gemm_iter = 0; + int gemm_tile = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + const float *fold_src = nullptr, *fold_routes = nullptr; + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + fold_src = nullptr; + fold_routes = nullptr; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + grid_barrier(g.barrier); +#endif + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); +#endif + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, BATCH, + route_stage); + } + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + } + grid_barrier(g.barrier); + r6_sparse_map(bb->moe_w,bb->r6_map,bb->r6_counts); + grid_barrier(g.barrier); + r6_sparse_gather(bb->normed,bb->r6_packed,bb->r6_map,bb->r6_counts, + bb->r6_lut); + grid_barrier(g.barrier); + gemm(L.exp_up, bb->r6_views[0], BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0,0,nullptr,nullptr,DECODE_TILES,DECODE_ROWS, + nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->r6_views[1], BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + DECODE_TILES, DECODE_ROWS, nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + if (bb->layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_selected(bb->moe_out, bb->moe_w, idx); + grid_barrier(g.barrier); + } else { + fold_src = bb->moe_out; + fold_routes = bb->moe_w; + } + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int n_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = n_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + bb->pf_h[idx] = bb->prompt[idx]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = bb->pf_n_seq > 0 ? bb->pf_n_seq : 1; + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_gated, + T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, T, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * (BB_EXPERTS + 1) * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->pf_normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->pf_moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = bb_moe_fold_deterministic(bb->pf_moe_out, idx); + grid_barrier(g.barrier); + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, bb->pf_final, ln_scratch, + n_total); + grid_barrier(g.barrier); + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, bb->phoneme_b); + grid_barrier(g.barrier); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + for (int row = blockIdx.x + warp * gridDim.x; row < BATCH; + row += gridDim.x * WARPS) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + MK2_PROGRESS_GUARD(native_progress_mask, row); + const int arg = phoneme_warp_argmax(lg); + if (lane == 0) bb->phoneme_out[row] = arg; + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + __shared__ int native_run_predictor; + if (native_predictor_active_mask && bb && g.native_codes) { + if (threadIdx.x == 0) { + int any_active = 0; + for (int row = 0; row < BATCH; ++row) + any_active |= int(native_predictor_active_mask[row] != 0 && + MK2_PROGRESS_ENABLED(native_progress_mask, row)); + native_run_predictor = any_active; + } + __syncthreads(); + } + const bool skip_native_predictor = native_predictor_active_mask && bb && + g.native_codes && native_run_predictor == 0; + if (skip_native_predictor) { + for (int index = blockIdx.x * THREADS + threadIdx.x; + index < BATCH * N_POS; index += gridDim.x * THREADS) { + MK2_PROGRESS_GUARD(native_progress_mask, index / N_POS); + g.sampled[index] = 0; + g.native_codes[index] = 0; + } + grid_barrier(g.barrier); + } else { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + constexpr int DPL = D_HEAD / 32; + float qv[DPL], kn[DPL], vn[DPL]; + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + qv[j] = __bfloat162float(q[i]); + kn[j] = __bfloat162float(k_new[i]); + vn[j] = __bfloat162float(v_new[i]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += qv[j] * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + { + float acc = 0.f; + #pragma unroll + for (int j = 0; j < DPL; ++j) acc += qv[j] * kn[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[pos] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + float ov[DPL] = {}; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *vt = g.vcache + cache_base + t * D_HEAD; + const float s = sc[t]; + #pragma unroll + for (int j = 0; j < DPL; ++j) + ov[j] += s * __bfloat162float(vt[lane + j * 32]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + ov[j] += sc[pos] * vn[j]; + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + + lane + j * 32] = __float2bfloat16(ov[j]); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, + g.sampled_pos, sort_storage, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif +#ifdef ABLATE_SAMPLE + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { + __syncthreads(); + } else { + grid_barrier(g.barrier); + } +#endif + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x == 0) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + g.sampled[row * N_POS + pos] = code; + if (g.native_codes && MK2_PROGRESS_ENABLED(native_progress_mask, row)) { +#if MK2_INFLIGHT + g.native_codes[row * N_POS + pos] = + native_predictor_active_mask && !native_predictor_active_mask[row] + ? 0 : (long long)code; +#else + g.native_codes[row * N_POS + pos] = (long long)code; +#endif + } + } + } +#ifdef ABLATE_SAMPLE + grid_barrier(g.barrier); +#else + if (pos + 1 == n_pos) grid_barrier(g.barrier); +#endif + } + } + if (g.native_seq_len && blockIdx.x == 0) + for (int row = threadIdx.x; row < BATCH; row += blockDim.x) { + MK2_PROGRESS_GUARD(native_progress_mask, row); + g.native_seq_len[row] += 1; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2_startup( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr, + const unsigned char *native_predictor_active_mask = nullptr + MK2_PROGRESS_PARAMETER(native_progress_mask)) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&p_s)[STAGES][SLOTS] = alloc.allocate(); + out_tile_f (&o_sf) = alloc.allocate(); + out_tile &o_s = *reinterpret_cast(&o_sf); + bf16 *const route_stage = reinterpret_cast(&o_sf); + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_SCRATCH_STRIDE]; + static_assert(STAGES * SLOTS * sizeof(weight_tile) >= sizeof(sampler_storage)); + sampler_storage &sort_storage = *reinterpret_cast(&__shm[0]); + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS, 0); + arrive(drained[s], CONSUMER_WARPS); + } + } + int gemm_iter = 0; + int gemm_tile = 0; + __syncthreads(); + constexpr bool prefilling = false; + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + const float *fold_src = nullptr, *fold_routes = nullptr; + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + fold_src = nullptr; + fold_routes = nullptr; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + grid_barrier(g.barrier); +#endif + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); +#endif + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, BATCH, + route_stage); + } + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + } + grid_barrier(g.barrier); + r6_sparse_map(bb->moe_w,bb->r6_map,bb->r6_counts); + grid_barrier(g.barrier); + r6_sparse_gather(bb->normed,bb->r6_packed,bb->r6_map,bb->r6_counts, + bb->r6_lut); + grid_barrier(g.barrier); + gemm(L.exp_up, bb->r6_views[0], BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0,0,nullptr,nullptr,DECODE_TILES,DECODE_ROWS, + nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->r6_views[1], BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + DECODE_TILES, DECODE_ROWS, nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + if (bb->layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_selected(bb->moe_out, bb->moe_w, idx); + grid_barrier(g.barrier); + } else { + fold_src = bb->moe_out; + fold_routes = bb->moe_w; + } + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if constexpr (false) { + const int n_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = n_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + bb->pf_h[idx] = bb->prompt[idx]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = bb->pf_n_seq > 0 ? bb->pf_n_seq : 1; + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_gated, + T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, T, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * (BB_EXPERTS + 1) * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->pf_normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->pf_moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = bb_moe_fold_deterministic(bb->pf_moe_out, idx); + grid_barrier(g.barrier); + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, bb->pf_final, ln_scratch, + n_total); + grid_barrier(g.barrier); + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, bb->phoneme_b); + grid_barrier(g.barrier); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + for (int row = blockIdx.x + warp * gridDim.x; row < BATCH; + row += gridDim.x * WARPS) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + MK2_PROGRESS_GUARD(native_progress_mask, row); + const int arg = phoneme_warp_argmax(lg); + if (lane == 0) bb->phoneme_out[row] = arg; + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + __shared__ int native_run_predictor; + if (native_predictor_active_mask && bb && g.native_codes) { + if (threadIdx.x == 0) { + int any_active = 0; + for (int row = 0; row < BATCH; ++row) + any_active |= int(native_predictor_active_mask[row] != 0 && + MK2_PROGRESS_ENABLED(native_progress_mask, row)); + native_run_predictor = any_active; + } + __syncthreads(); + } + const bool skip_native_predictor = native_predictor_active_mask && bb && + g.native_codes && native_run_predictor == 0; + if constexpr (true) { + for (int index = blockIdx.x * THREADS + threadIdx.x; + index < BATCH * N_POS; index += gridDim.x * THREADS) { + MK2_PROGRESS_GUARD(native_progress_mask, index / N_POS); + g.sampled[index] = 0; + g.native_codes[index] = 0; + } + grid_barrier(g.barrier); + } else { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + constexpr int DPL = D_HEAD / 32; + float qv[DPL], kn[DPL], vn[DPL]; + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + qv[j] = __bfloat162float(q[i]); + kn[j] = __bfloat162float(k_new[i]); + vn[j] = __bfloat162float(v_new[i]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += qv[j] * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + { + float acc = 0.f; + #pragma unroll + for (int j = 0; j < DPL; ++j) acc += qv[j] * kn[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[pos] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + float ov[DPL] = {}; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *vt = g.vcache + cache_base + t * D_HEAD; + const float s = sc[t]; + #pragma unroll + for (int j = 0; j < DPL; ++j) + ov[j] += s * __bfloat162float(vt[lane + j * 32]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + ov[j] += sc[pos] * vn[j]; + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + + lane + j * 32] = __float2bfloat16(ov[j]); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, + g.sampled_pos, sort_storage, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif +#ifdef ABLATE_SAMPLE + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { + __syncthreads(); + } else { + grid_barrier(g.barrier); + } +#endif + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x == 0) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + g.sampled[row * N_POS + pos] = code; + if (g.native_codes && MK2_PROGRESS_ENABLED(native_progress_mask, row)) { +#if MK2_INFLIGHT + g.native_codes[row * N_POS + pos] = + native_predictor_active_mask && !native_predictor_active_mask[row] + ? 0 : (long long)code; +#else + g.native_codes[row * N_POS + pos] = (long long)code; +#endif + } + } + } +#ifdef ABLATE_SAMPLE + grid_barrier(g.barrier); +#else + if (pos + 1 == n_pos) grid_barrier(g.barrier); +#endif + } + } + if (g.native_seq_len && blockIdx.x == 0) + for (int row = threadIdx.x; row < BATCH; row += blockDim.x) { + MK2_PROGRESS_GUARD(native_progress_mask, row); + g.native_seq_len[row] += 1; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 64 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * SLOTS * sizeof(weight_tile) + + sizeof(out_tile_f) + 2048; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; + const unsigned char *native_predictor_active_mask = nullptr; + MK2_PROGRESS_MEMBER + int maximum_grid = 0; + bool grid_reset_pending = false; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * SLOTS * sizeof(weight_tile) + + sizeof(out_tile_f) + 2048), 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaFuncSetAttribute(predictor_mk2_startup, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = std::min(out->sms * out->blocks_per_sm, 4096); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 64 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->maximum_grid = info.sms * info.blocks_per_sm; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + mk2_check(owned.alloc(&hb.r6_packed,(size_t)BB_EXPERTS*BATCH*D_MODEL*sizeof(bf16))); + mk2_check(owned.alloc(&hb.r6_map,(size_t)BB_EXPERTS*BATCH*sizeof(int))); + mk2_check(owned.alloc(&hb.r6_counts,BB_EXPERTS*sizeof(int))); + mk2_check(owned.alloc(&hb.r6_lut,(1+BB_EXPERTS*DECODE_TILES)*sizeof(int))); + act_gl sparse_views[]={ + act_gl{hb.r6_packed,nullptr,nullptr,(size_t)BB_EXPERTS*BATCH,(size_t)D_MODEL}, + act_gl{moe_h,nullptr,nullptr,(size_t)BB_EXPERTS*BATCH,(size_t)BB_MOE_INTER}}; + mk2_check(owned.alloc(&hb.r6_views,sizeof(sparse_views))); + mk2_check(cudaMemcpy(hb.r6_views,sparse_views,sizeof(sparse_views),cudaMemcpyHostToDevice)); + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.maximum_grid = pred->maximum_grid; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native, bool startup = false) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(s, &capture_status)); + if (st->pred.grid_reset_pending || capture_status != cudaStreamCaptureStatusNone) { + mk2_check(cudaMemsetAsync(st->pred.g.barrier, 0, 64 * sizeof(unsigned int), s)); + st->pred.grid_reset_pending = false; + } + if (startup) { + predictor_mk2_startup<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS, nullptr, + st->pred.native_predictor_active_mask + MK2_PROGRESS_ARGUMENT(st->pred.native_progress_mask)); + } else { + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS, nullptr, + native ? st->pred.native_predictor_active_mask : nullptr + MK2_PROGRESS_ARGUMENT(native ? st->pred.native_progress_mask : nullptr)); + } + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" int mk2_run_native_startup_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream) { + return mk2_enqueue_stage0(p, b, stream, true, true); +} +extern "C" int mk2_set_native_predictor_active_mask(MK2Predictor *p, + const unsigned char *mask) { + try { + if (!p || reinterpret_cast(p)->kind != 1) + throw std::runtime_error("startup mask requires a full Stage0 handle"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && st->pred.native_predictor_active_mask != mask) + throw std::runtime_error("startup mask pointer cannot change after first Stage0 enqueue"); + if (mask) + mk2_stage0_device_pointer(mask, st->pred.device, "native_predictor_active_mask"); + st->pred.native_predictor_active_mask = mask; + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown startup mask binding failure"; return 1; } +} +extern "C" int mk2_native_startup_skip_abi_version() { return 1; } +#include "../../../../../inflight/runtime.cuh" +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "native_prefill_impl.inc" +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *p, int grid) { + if (!p) { g_last_error = "grid setter requires full Stage0 handle"; return 1; } + auto *st = reinterpret_cast(p); + if (st->kind != 1 || grid < 1 || grid > st->maximum_grid) { + g_last_error = "requested Stage0 grid outside independently resident capacity"; + return 1; + } + if (st->grid != grid) st->grid_reset_pending = true; + st->grid = grid; + g_last_error.clear(); + return 0; +} +extern "C" int mk2_r8_bind_initial_state(MK2Predictor *p, const void *const *initial, const void *ready, const void *seq_len, int positions) { + try { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || st->bindings_initialized || !initial || !ready || !seq_len || positions <= 0) + throw std::runtime_error("invalid immutable initial-state binding"); + for (int i = 0; i < 14; ++i) { + mk2_stage0_device_pointer(initial[i], st->pred.device, "initial state"); + st->mb_host[i].r8_initial_ssm = static_cast(initial[i]); + st->mb_host[i].r8_initial_ready = static_cast(ready); + st->mb_host[i].r8_seq_len = static_cast(seq_len); + st->mb_host[i].r8_initial_positions = positions; + } + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" int mk2_r9_state_bind(void* p,const void*const* scales,const void*const* maps) {auto*st=(MK2Stage0State*)p;if(st->bindings_initialized)return 1;for(int i=0;i<14;++i){st->mb_host[i].r9_scales=(const float*)scales[i];st->mb_host[i].r9_map=(const int*)maps[i];}return 0;} +extern "C" int mk2_r9_state_bytes(){return sizeof(r9_state_t);} diff --git a/deploy/nemotron_native/kernels/b128/regions/stage0/stage0.cu b/deploy/nemotron_native/kernels/b128/regions/stage0/stage0.cu new file mode 100644 index 0000000..c32ea89 --- /dev/null +++ b/deploy/nemotron_native/kernels/b128/regions/stage0/stage0.cu @@ -0,0 +1,35 @@ +#define MK2_BATCH 128 +#define STAGES 3 +#define KCHUNK 128 +#define MK2_NO_MAIN 1 +#define MK2_DIRECT_INPUT 1 +#define MK2_B1_LOCAL_BARRIERS 1 +#define MK2_SKIP_HIDDEN_SNAPSHOT 1 +#define MK2_ACTIVE_EXPERTS 1 +#define MK2_SSM_SPLITS 1 +#define NMULT 2 +#define MK2_OUT_MULT 2 +#define SSM_PREFETCH_MAX 12 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#if MK2_INFLIGHT +#include +#endif +#include +#include +#include +#include +#include +#include +#include +#include +namespace b128_stage0 { +#include "native/stage0/predictor_mk2.cu" +} diff --git a/deploy/nemotron_native/kernels/b128/stage0.cu b/deploy/nemotron_native/kernels/b128/stage0.cu new file mode 100644 index 0000000..6c04cd7 --- /dev/null +++ b/deploy/nemotron_native/kernels/b128/stage0.cu @@ -0,0 +1 @@ +#include "regions/stage0/stage0.cu" diff --git a/deploy/nemotron_native/kernels/b2048/prefill.cu b/deploy/nemotron_native/kernels/b2048/prefill.cu new file mode 100644 index 0000000..6d21d7e --- /dev/null +++ b/deploy/nemotron_native/kernels/b2048/prefill.cu @@ -0,0 +1 @@ +#include "regions/prefill/prefill.cu" diff --git a/deploy/nemotron_native/kernels/b2048/regions/prefill/native/prefill/dedup_native_prefill_impl.inc b/deploy/nemotron_native/kernels/b2048/regions/prefill/native/prefill/dedup_native_prefill_impl.inc new file mode 100644 index 0000000..fe4c125 --- /dev/null +++ b/deploy/nemotron_native/kernels/b2048/regions/prefill/native/prefill/dedup_native_prefill_impl.inc @@ -0,0 +1,267 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 16 || BATCH == 32 || BATCH == 64 || BATCH == 128 || BATCH == 256 || BATCH == 512 || BATCH == 1024 || BATCH == 2048)) + throw std::runtime_error("native prefill requires a matching supported physical batch"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((std::uint64_t(c.prompt_positions) + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (rows > (INT_MAX - 132 * THREADS) / D_MODEL) + throw std::runtime_error("native prefill full-width prompt exceeds safe device loop bounds"); + if (padded > (INT_MAX - 132 * THREADS) / (BB_EXPERTS * BB_MOE_INTER)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((size_t(config->prompt_positions) + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + float *r9_workspace = nullptr; + mk2_check(owned.alloc(&r9_workspace,size_t(BATCH)*ssm_elements*sizeof(float))); + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + r9_state_t *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(r9_state_t))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(r9_state_t))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = r9_workspace; + st->mb_host[i].r9_private = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + mk2_check(owned.alloc(&st->mb_host[i].r8_canonical_ssm, ssm_elements * sizeof(float))); + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_dtda, 2 * MB_HEADS); + { + const size_t estride = padded; + const size_t grows = size_t(BB_EXPERTS) * estride; + mk2_check(owned.alloc(&hb.pf_moe_gath, grows * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_gath, 0, grows * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_hc, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_hc, 0, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_gidx, grows * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_gidx, 0, grows * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ecnt, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ecnt, 0, BB_EXPERTS * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ntiles, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ntiles, 0, BB_EXPERTS * sizeof(int))); + hb.pf_moe_estride = int(estride); + hb.pf_moe_gath_v = act_gl{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_v = act_gl{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + hb.pf_moe_gath_vn = act_gl_n{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_vn = act_gl_n{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + scratch_bytes += grows * (D_MODEL + BB_MOE_INTER) * sizeof(bf16); + } + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(owned.alloc(&hb.native_pf_cohort_check, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(hb.native_pf_cohort_check, 0, 2 * sizeof(unsigned int))); +#if PF_ACTIVE_TOKEN_EXPERTS + mk2_check(owned.alloc(&hb.native_pf_expert_masks, (padded/BATCH_PAD)*sizeof(unsigned int))); +#endif + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} +extern "C" int mk2_native_prefill_validate(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null compact prefill handle"); + auto *st = reinterpret_cast(p->base); + mk2_check(cudaStreamSynchronize(static_cast(stream))); + unsigned int status[2] = {0, 0}; + mk2_check(cudaMemcpy(status, st->host_bb.native_pf_cohort_check, + sizeof(status), cudaMemcpyDeviceToHost)); + if (status[0] != 0 || status[1] != 1) + throw std::runtime_error("compact prefill requires bit-identical prompts and identity slots"); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown compact prefill validation failure"; return 1; } +} +extern "C" int mk2_native_prefill_workspace_sequences() { return 1; } +extern "C" int mk2_native_prefill_effective_sequences(MK2NativePrefill *p, int *out) { + try { + if (!p || !out) throw std::runtime_error("null effective-sequence query"); + auto *st = reinterpret_cast(p->base); + mk2_check(cudaMemcpy(out, st->host_bb.native_pf_cohort_check + 1, + sizeof(int), cudaMemcpyDeviceToHost)); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" int mk2_r8_initial_state(MK2NativePrefill *p, void **initial, void **ready) { + if (!p || !initial || !ready) return 1; + auto *st = reinterpret_cast(p->base); + for (int i = 0; i < 14; ++i) initial[i] = st->mb_host[i].r8_canonical_ssm; + *ready = st->host_bb.native_pf_cohort_check + 1; + return 0; +} +extern "C" int mk2_r9_prefill_bind(MK2NativePrefill*p,const void*const*scales,const void*const*maps){if(!p)return 1;auto*st=(MK2Stage0State*)p->base;for(int i=0;i<14;++i){st->mb_host[i].r9_scales=(const float*)scales[i];st->mb_host[i].r9_map=(const int*)maps[i];}return (int)cudaMemcpy((void*)st->host_bb.mamba,st->mb_host.data(),14*sizeof(MambaLayer),cudaMemcpyHostToDevice);} diff --git a/deploy/nemotron_native/kernels/b2048/regions/prefill/native/prefill/dedup_predictor_mk2.cuh b/deploy/nemotron_native/kernels/b2048/regions/prefill/native/prefill/dedup_predictor_mk2.cuh new file mode 100644 index 0000000..f98dd53 --- /dev/null +++ b/deploy/nemotron_native/kernels/b2048/regions/prefill/native/prefill/dedup_predictor_mk2.cuh @@ -0,0 +1,4600 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef PF_ACTIVE_TOKEN_EXPERTS +#define PF_ACTIVE_TOKEN_EXPERTS 0 +#endif +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +#ifndef MK2_PREFILL_ONLY +#define MK2_PREFILL_ONLY 0 +#endif +#ifndef MK2_BATCH_TILE +#define MK2_BATCH_TILE (((MK2_BATCH + 15) / 16) * 16) +#endif +constexpr int BATCH_PAD = MK2_BATCH_TILE; +static_assert(BATCH >= 1 && (MK2_PREFILL_ONLY || BATCH <= BATCH_PAD)); +static_assert(BATCH_PAD == 16 || BATCH_PAD == 32 || BATCH_PAD == 64 || BATCH_PAD == 128, + "native prefill tile must be 16, 32, 64 or 128"); +static_assert(MK2_PREFILL_ONLY || BATCH_PAD == ((BATCH + 15) / 16) * 16, + "a wider token tile requires a prefill-only library"); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_PW_MULTI +#define PF_SSM_PW_MULTI 4 +#endif +#ifndef PF_SSM_PW_SINGLE +#define PF_SSM_PW_SINGLE 2 +#endif +static_assert(MB_HEAD_DIM % PF_SSM_PW_MULTI == 0 && + MB_HEAD_DIM % PF_SSM_PW_SINGLE == 0, + "head_dim must split evenly into scan channel groups"); +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef MK2_SMEM_SLACK +#define MK2_SMEM_SLACK 1024 +#endif +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +#ifndef MK2_THREADS +#define MK2_THREADS 256 +#endif +constexpr int THREADS = MK2_THREADS; +constexpr int NORM_THREADS = 256; +constexpr int NORM_WARPS = NORM_THREADS / 32; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +#ifndef MK2_GEMM_CONSUMERS +#define MK2_GEMM_CONSUMERS 1 +#endif +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 +#define MK2_HEAD_STAGE_BYTES 0 +#else +#define MK2_HEAD_STAGE_BYTES sizeof(out_tile_f) +#endif +constexpr int NCONS = MK2_GEMM_CONSUMERS; +static_assert(NCONS >= 1 && NCONS <= 3, "consumers occupy warpgroups 0, 2 and 3"); +static_assert(THREADS >= (NCONS + 1) * 128, "one producer plus NCONS consumers"); +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +constexpr int PAD_N = 32; +using act_tile_n = st_bf; +using out_tile_n = st_bf; +using act_gl_n = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_MOE_ACC_C, EPI_ACC_F32 }; +#include +constexpr int R9_STATE_MODE=1; +using r9_state_t=bf16; +__device__ __forceinline__ int r9_scale_index(int head,int channel,int n) {return (((head/8)*4+(head%8)/2)*4+channel/6)*128+n;} +__device__ __forceinline__ float r9_load(r9_state_t v,float scale) { +return __bfloat162float(v); +} +__device__ __forceinline__ r9_state_t r9_store(float x,float scale) { +return __float2bfloat16(x); +} +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; + float *r8_canonical_ssm = nullptr; + r9_state_t *r9_private = nullptr; + const float *r9_scales = nullptr; + const int *r9_map = nullptr; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + unsigned int *native_pf_expert_masks = nullptr; + float *pf_dtda = nullptr; + bf16 *pf_moe_gath = nullptr; + bf16 *pf_moe_hc = nullptr; + int *pf_moe_gidx = nullptr; + int *pf_moe_ecnt = nullptr; + int *pf_moe_ntiles = nullptr; + int pf_moe_estride = 0; + act_gl pf_moe_gath_v; + act_gl pf_moe_hc_v; + act_gl_n pf_moe_gath_vn; + act_gl_n pf_moe_hc_vn; + unsigned int *native_pf_cohort_check = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + volatile unsigned int *generation = state + 1; + const unsigned int seen = *generation; + if (atomicAdd(state, 1u) == gridDim.x - 1) { + atomicExch(state, 0u); + atomicAdd(state + 1, 1u); + } else { + while (*generation == seen) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const bf16 *in = src + (size_t)row * D_MODEL; + float v[PER_THREAD]; + float sum = 0.f, sq = 0.f; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (threadIdx.x < NORM_THREADS) { + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + v[j] = __bfloat162float(in[threadIdx.x + j * NORM_THREADS]); + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + if (lane == 0) { scratch[warp] = sum; scratch[NORM_WARPS + warp] = sq; } + } + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f, q = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) { s += scratch[w]; q += scratch[NORM_WARPS + w]; } + scratch[2 * NORM_WARPS] = s / D_MODEL; + scratch[2 * NORM_WARPS + 1] = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + } + __syncthreads(); + const float mean = scratch[2 * NORM_WARPS], inv = scratch[2 * NORM_WARPS + 1]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + if (threadIdx.x >= NORM_THREADS) break; + const int i = threadIdx.x + j * NORM_THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * weight[i]); + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, + int idx, unsigned mask); +__device__ __forceinline__ void norm_sub_sync(int sub) { + if (THREADS == NORM_THREADS) { __syncthreads(); return; } + asm volatile("bar.sync %0, %1;" :: "r"(6 + sub), "n"(NORM_THREADS) : "memory"); +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, const float *fold_src = nullptr, + const float *fold_routes = nullptr) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + constexpr int NSUB = THREADS / NORM_THREADS; + if (threadIdx.x >= NSUB * NORM_THREADS) return; + const int sub = threadIdx.x / NORM_THREADS; + const int tid = threadIdx.x - sub * NORM_THREADS; + float *const sc = scratch + sub * (NORM_WARPS + 1); + for (int row = blockIdx.x * NSUB + sub; row < n_rows; + row += gridDim.x * NSUB) { + const size_t base = (size_t)row * D_MODEL; + unsigned fold_mask = (1u << BB_EXPERTS) - 1u; + if (fold_src && fold_routes) { + fold_mask = 0u; + const float *rw = fold_routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) + if (rw[e] != 0.f) fold_mask |= 1u << e; + } + float v[PER_THREAD]; + float sq = 0.f; + const int warp = tid >> 5, lane = tid & 31; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + bf16 hv; + if (fold_src) { + hv = bb_moe_fold_deterministic(fold_src, (int)(base + i), fold_mask); + h[base + i] = hv; + } else { + hv = h[base + i]; + } + v[j] = __bfloat162float(hv) + __bfloat162float(r[base + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + if (lane == 0) sc[warp] = sq; + norm_sub_sync(sub); + if (tid == 0) { + float s = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) s += sc[w]; + sc[NORM_WARPS] = rsqrtf(s / D_MODEL + RMS_EPS); + } + norm_sub_sync(sub); + const float inv = sc[NORM_WARPS]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + out[base + i] = __float2bfloat16( + __bfloat162float(weight[i]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + } + norm_sub_sync(sub); + } +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + const int c = idx - row * MB_CONV_DIM; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + const float raw = __bfloat162float( + zxbcdt[(size_t)row * MB_IN_PROJ + MB_INNER + c]); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + xbc_act[(size_t)row * MB_CONV_DIM + c] = __float2bfloat16(acc); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + const int pairs = BATCH * MB_HEADS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_HEADS * split); + const int head = (pair / split) % MB_HEADS; + const int first_channel = (pair % split) * channels; + const int group = head / MB_HEADS_PER_GROUP; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + L.dt_bias[head]); + const float dA = __expf(L.A[head] * dt); + const float Dh = L.D[head]; + float *st = L.ssm_state + (size_t)slot_map[row] * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + for (int p = first_channel; p < first_channel + channels; ++p) { + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + const size_t at = (size_t)p * MB_STATE + n; + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + static_assert(MB_CONV_K == 4, "packed conv taps assume a 4-tap kernel"); + constexpr int CPT = 4; + static_assert(MB_CONV_DIM % CPT == 0, "channel vector width"); + constexpr int CGROUPS = MB_CONV_DIM / CPT; + const int total = n_tokens * CGROUPS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / CGROUPS; + const int c = (idx - t * CGROUPS) * CPT; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + float wv[CPT][MB_CONV_K]; + { + const float4 *wsrc = reinterpret_cast( + L.conv_w + (size_t)c * MB_CONV_K); + #pragma unroll + for (int j = 0; j < CPT / 2; ++j) { + const float4 w4 = wsrc[j]; + const unsigned uu0 = __float_as_uint(w4.x); + const unsigned uu1 = __float_as_uint(w4.y); + const unsigned uu2 = __float_as_uint(w4.z); + const unsigned uu3 = __float_as_uint(w4.w); + wv[2 * j][0] = __uint_as_float(uu0 << 16); + wv[2 * j][1] = __uint_as_float(uu0 & 0xffff0000u); + wv[2 * j][2] = __uint_as_float(uu1 << 16); + wv[2 * j][3] = __uint_as_float(uu1 & 0xffff0000u); + wv[2 * j + 1][0] = __uint_as_float(uu2 << 16); + wv[2 * j + 1][1] = __uint_as_float(uu2 & 0xffff0000u); + wv[2 * j + 1][2] = __uint_as_float(uu3 << 16); + wv[2 * j + 1][3] = __uint_as_float(uu3 & 0xffff0000u); + } + } + float acc[CPT]; + { + const uint2 b2 = *reinterpret_cast(L.conv_b + c); + acc[0] = __uint_as_float(b2.x << 16); + acc[1] = __uint_as_float(b2.x & 0xffff0000u); + acc[2] = __uint_as_float(b2.y << 16); + acc[3] = __uint_as_float(b2.y & 0xffff0000u); + } + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v[CPT]; + if (src < 0) { + #pragma unroll + for (int q = 0; q < CPT; ++q) + v[q] = __bfloat162float( + state[(size_t)q * (MB_CONV_K - 1) + pos + k]); + } else { + const uint2 p = *reinterpret_cast( + zxbcdt + (size_t)src * MB_IN_PROJ + MB_INNER + c); + v[0] = __uint_as_float(p.x << 16); + v[1] = __uint_as_float(p.x & 0xffff0000u); + v[2] = __uint_as_float(p.y << 16); + v[3] = __uint_as_float(p.y & 0xffff0000u); + } + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] += wv[q][k] * v[q]; + } + uint2 out; + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] = acc[q] / (1.f + __expf(-acc[q])); + out.x = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[0])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[1])) << 16); + out.y = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[2])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[3])) << 16); + *reinterpret_cast(xbc_act + (size_t)t * MB_CONV_DIM + c) = out; + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_dtda_prefill( + const MambaLayer &L, const bf16 *zxbcdt, float *dtda, int n_tokens) { + const int total = n_tokens * MB_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_HEADS; + const int head = idx - t * MB_HEADS; + const float dt = mamba_softplus( + __bfloat162float(zxbcdt[(size_t)t * MB_IN_PROJ + MB_INNER + + MB_CONV_DIM + head]) + L.dt_bias[head]); + dtda[(size_t)idx * 2] = dt; + dtda[(size_t)idx * 2 + 1] = __expf(L.A[head] * dt); + } +} +template +__device__ __forceinline__ void mamba_ssm_prefill_impl( + const MambaLayer &L, const bf16 *xbc_act, const float *dtda, bf16 *gated, + int n_pos, int n_seq, int slot, const int *seq_slot) { + constexpr int NS = MB_STATE / 32; + constexpr int PG = MB_HEAD_DIM / PW; + constexpr bool RS = (PW % 4) == 0; + constexpr int H1 = RS ? PW / 2 : 1; + constexpr int H2 = RS ? PW / 4 : 1; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const bool keep_hi = (lane & 16) != 0; + const bool keep_hi2 = (lane & 8) != 0; + const int warps_total = gridDim.x * WARPS; + const int units = n_seq * MB_HEADS * PG; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int pg = unit % PG; + const int rest = unit / PG; + const int head = rest % MB_HEADS; + const int seq = rest / MB_HEADS; + const int p0 = pg * PW; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Dh = L.D[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p0 * MB_STATE; + float sreg[PW][NS]; + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + sreg[q][j] = st[(size_t)q * MB_STATE + lane + j * 32]; + } + auto fetch = [&](int p, float (&oB)[NS], float (&oC)[NS], + float (&oX)[PW], float2 &od) { + const size_t tt = (size_t)p * n_seq + seq; + const bf16 *xa2 = xbc_act + tt * MB_CONV_DIM; + const bf16 *Bs = xa2 + MB_INNER + group * MB_STATE; + const bf16 *Cs = xa2 + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + od = *reinterpret_cast( + dtda + (tt * MB_HEADS + head) * 2); + #pragma unroll + for (int j = 0; j < NS; ++j) { + oB[j] = __bfloat162float(Bs[lane + j * 32]); + oC[j] = __bfloat162float(Cs[lane + j * 32]); + } + if constexpr (RS) { + const uint2 *xp = reinterpret_cast( + xa2 + head * MB_HEAD_DIM + p0); + #pragma unroll + for (int q = 0; q < PW; q += 4) { + const uint2 w = xp[q >> 2]; + oX[q] = __uint_as_float(w.x << 16); + oX[q + 1] = __uint_as_float(w.x & 0xffff0000u); + oX[q + 2] = __uint_as_float(w.y << 16); + oX[q + 3] = __uint_as_float(w.y & 0xffff0000u); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) + oX[q] = __bfloat162float(xa2[head * MB_HEAD_DIM + p0 + q]); + } + }; + float pB[NS], pC[NS], pX[PW]; + float2 pd; + fetch(0, pB, pC, pX, pd); + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const float dt = pd.x; + const float dA = pd.y; + float Bv[NS], Cv[NS]; + #pragma unroll + for (int j = 0; j < NS; ++j) { + Bv[j] = pB[j] * dt; + Cv[j] = pC[j]; + } + float acc[PW], xvs[PW]; + #pragma unroll + for (int q = 0; q < PW; ++q) xvs[q] = pX[q]; + if (pos + 1 < n_pos) fetch(pos + 1, pB, pC, pX, pd); + #pragma unroll + for (int q = 0; q < PW; ++q) { + const float xv = xvs[q]; + float a = 0.f; + #pragma unroll + for (int j = 0; j < NS; ++j) { + sreg[q][j] = sreg[q][j] * dA + Bv[j] * xv; + a += sreg[q][j] * Cv[j]; + } + acc[q] = a; + } + if constexpr (RS) { + float m1[H1], t1[H1], x1[H1]; + #pragma unroll + for (int q = 0; q < H1; ++q) { + m1[q] = keep_hi ? acc[q + H1] : acc[q]; + t1[q] = keep_hi ? acc[q] : acc[q + H1]; + x1[q] = keep_hi ? xvs[q + H1] : xvs[q]; + } + #pragma unroll + for (int q = 0; q < H1; ++q) + m1[q] += __shfl_xor_sync(0xffffffff, t1[q], 16); + float m2[H2], t2[H2], x2[H2]; + #pragma unroll + for (int q = 0; q < H2; ++q) { + m2[q] = keep_hi2 ? m1[q + H2] : m1[q]; + t2[q] = keep_hi2 ? m1[q] : m1[q + H2]; + x2[q] = keep_hi2 ? x1[q + H2] : x1[q]; + } + #pragma unroll + for (int q = 0; q < H2; ++q) + m2[q] += __shfl_xor_sync(0xffffffff, t2[q], 8); + #pragma unroll + for (int q = 0; q < H2; ++q) { + #pragma unroll + for (int off = 4; off > 0; off >>= 1) + m2[q] += __shfl_xor_sync(0xffffffff, m2[q], off); + } + if ((lane & 7) == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + + p0 + (keep_hi ? H1 : 0) + (keep_hi2 ? H2 : 0); + #pragma unroll + for (int q = 0; q < H2; ++q) + dst[q] = __float2bfloat16(m2[q] + x2[q] * Dh); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc[q] += __shfl_xor_sync(0xffffffff, acc[q], off); + } + if (lane == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + p0; + #pragma unroll + for (int q = 0; q < PW; ++q) + dst[q] = __float2bfloat16(acc[q] + xvs[q] * Dh); + } + } + } + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + st[(size_t)q * MB_STATE + lane + j * 32] = sreg[q][j]; + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, + const float *dtda, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + (void)zxbcdt; + if (n_seq > 1) + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens / n_seq, n_seq, slot, seq_slot); + else + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens, 1, slot, nullptr); +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16( + held[j] * inv * __bfloat162float(L.gnorm_w[idx])); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = BATCH * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / BB_ATT_HEADS; + const int head = pair % BB_ATT_HEADS; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const int kv = head / kv_per_q; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *q = base + head * BB_HEAD_DIM; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (head % kv_per_q == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + for (int tpos = 0; tpos < T; ++tpos) { +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new + : kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#endif + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m, s); + const float rescale = __expf(m - m_new); + const float w = __expf(s - m_new); + denom = denom * rescale + w; +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *vt = kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#else + const bf16 *vt = cur_token + ? v_new + : kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + w * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + constexpr int CHUNKS = BB_HEAD_DIM / 8; + const int total = n_tokens * BB_KV_HEADS * CHUNKS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int c = idx % CHUNKS; + const int rest = idx / CHUNKS; + const int t = rest / BB_KV_HEADS; + const int kv = rest - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM + c * 8; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM + c * 8; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv) + c * 8; + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv) + c * 8; + *reinterpret_cast(kc) = *reinterpret_cast(k_new); + *reinterpret_cast(vc) = *reinterpret_cast(v_new); + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float qv[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qv[j] = __bfloat162float(q[lane + j * 32]); + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + const int *const krow = + L.block_table + ((size_t)seq * 2 + 0) * L.max_blocks_per_seq; + const int *const vrow = + L.block_table + ((size_t)seq * 2 + 1) * L.max_blocks_per_seq; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + const int head_bias = kv * L.tokens_per_block * BB_HEAD_DIM; + int page = -1; + const bf16 *kpage = nullptr, *vpage = nullptr; + float kn[PER_LANE], vn[PER_LANE]; + { + page = 0; + kpage = L.pool_base + (size_t)(krow[0] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + (size_t)(vrow[0] & 0x7FFFFFFF) * elems_per_block; + const int at = head_bias + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + for (int s = 0; s <= pos; ++s) { + float kc[PER_LANE], vc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { kc[j] = kn[j]; vc[j] = vn[j]; } + if (s < pos) { + const int np = (s + 1) / L.tokens_per_block; + if (np != page) { + page = np; + kpage = L.pool_base + + (size_t)(krow[np] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + + (size_t)(vrow[np] & 0x7FFFFFFF) * elems_per_block; + } + const int at = head_bias + + ((s + 1) % L.tokens_per_block) * BB_HEAD_DIM + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) dot += qv[j] * kc[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[j] = acc[j] * rescale + wt * vc[j]; + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH, unsigned int *tile_masks = nullptr, + bf16 *gath = nullptr, int *gidx = nullptr, int *ecnt = nullptr, + int estride = 0, bf16 *row_stage = nullptr) { + __shared__ float score[BB_EXPERTS]; + __shared__ int s_pick[BB_TOP_K], s_slot[BB_TOP_K]; + constexpr int CH = D_MODEL / 8; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + { + const float4 *gsrc = + reinterpret_cast(normed + (size_t)row * D_MODEL); + float4 *gdst = reinterpret_cast(row_stage); + for (int i = threadIdx.x; i < CH; i += THREADS) gdst[i] = gsrc[i]; + } + if ((int)threadIdx.x < BB_EXPERTS) + moe_w[(size_t)row * BB_EXPERTS + threadIdx.x] = 0.f; + __syncthreads(); + for (int e = warp; e < BB_EXPERTS; e += WARPS) { + const bf16 *wr = gate_w + (size_t)e * D_MODEL; + float acc = 0.f; + for (int i = lane; i < D_MODEL; i += 32) + acc += __bfloat162float(row_stage[i]) * + __bfloat162float(wr[i]); + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) score[e] = 1.f / (1.f + __expf(-acc)); + } + __syncthreads(); + if (threadIdx.x == 0) { + float remaining[BB_EXPERTS]; + int pick[BB_TOP_K]; + for (int e = 0; e < BB_EXPERTS; ++e) + remaining[e] = score[e] + __bfloat162float(gate_bias[e]); + float total = 0.f; + for (int k = 0; k < BB_TOP_K; ++k) { + int best = 0; + for (int e = 1; e < BB_EXPERTS; ++e) + if (remaining[e] > remaining[best]) best = e; + pick[k] = best; + total += score[best]; + remaining[best] = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + unsigned int bits = 0u; + for (int k = 0; k < BB_TOP_K; ++k) { + moe_w[row * BB_EXPERTS + pick[k]] = score[pick[k]] * norm; + bits |= 1u << pick[k]; + s_pick[k] = pick[k]; + } + if (tile_masks) atomicOr(&tile_masks[row / BATCH_PAD], bits); + } + __syncthreads(); + if (gath) { + if ((int)threadIdx.x < BB_TOP_K) { + const int k = threadIdx.x; + const int slot = atomicAdd(&ecnt[s_pick[k]], 1); + s_slot[k] = slot; + gidx[(size_t)s_pick[k] * estride + slot] = row; + } + __syncthreads(); + const float4 *src = reinterpret_cast(row_stage); + for (int k = 0; k < BB_TOP_K; ++k) { + float4 *dst = reinterpret_cast( + gath + ((size_t)s_pick[k] * estride + s_slot[k]) * D_MODEL); + for (int i = threadIdx.x; i < CH; i += THREADS) dst[i] = src[i]; + } + } + __syncthreads(); + } +} +__device__ __forceinline__ unsigned moe_row_mask(const float *routes, int row) { + unsigned m = 0u; + const float *rw = routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) if (rw[e] != 0.f) m |= 1u << e; + return m; +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx, + unsigned mask) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + for (unsigned m = mask; m; m &= m - 1u) + value = __fadd_rn(value, parts[base + (size_t)(__ffs(m) - 1) * D_MODEL]); + return __float2bfloat16( + __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL])); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + uint64_t *keys, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int n = threadIdx.x; n < SORT_N; n += THREADS) { + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + keys[n] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + __syncthreads(); + for (int k = 2; k <= SORT_N; k <<= 1) { + if (k >= 64) __syncthreads(); + for (int j = k >> 1; j > 0; j >>= 1) { + for (int i = threadIdx.x; i < SORT_N; i += THREADS) { + const int partner = i ^ j; + if (partner > i) { + const bool descending = ((i & k) == 0); + if ((keys[i] < keys[partner]) == descending) { + const uint64_t t = keys[i]; + keys[i] = keys[partner]; + keys[partner] = t; + } + } + } + if (j >= 32) __syncthreads(); + else __syncwarp(); + } + } + __syncthreads(); + if (threadIdx.x == 0) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float v; + int n; + if (!eos_used && eos_v > av) { + v = eos_v; + n = AUDIO_EOS; + eos_used = true; + } else { + v = av; + n = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = + v + gumbel[((size_t)row * N_POS + pos) * TOP_K + j]; + if (noised > best_noised) { best_noised = noised; best_code = n; } + } + out_codes[row] = best_code; + } + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void prefill_expert_masks(const float *routes, unsigned int *masks, int rows) { + const int lane=threadIdx.x&31, warp=threadIdx.x>>5; + const int nt=(rows+BATCH_PAD-1)/BATCH_PAD; + for (int t=blockIdx.x*WARPS+warp;t(gath + g * D_MODEL + c * 8); + if (local < ecnt[e]) + *dst = *reinterpret_cast( + normed + (size_t)gidx[g] * D_MODEL + c * 8); + else + *dst = make_float4(0.f, 0.f, 0.f, 0.f); + } +} +__device__ __forceinline__ int gemm_eff(int tiles, int reps) { + if (NCONS <= 1) return 1; + const int g = (int)gridDim.x; + const int rounds1 = (tiles * reps + g - 1) / g; + if constexpr (NCONS == 3) { + const int max_eff = reps <= BB_EXPERTS ? 2 : 3; + const int rounds2 = (((tiles + 1) / 2) * reps + g - 1) / g; + if (max_eff == 2) return rounds2 < rounds1 ? 2 : 1; + const int rounds3 = (((tiles + 2) / 3) * reps + g - 1) / g; + if (rounds3 < rounds2 && rounds3 < rounds1) return 3; + return rounds2 < rounds1 ? 2 : 1; + } else { + const int roundsN = (((tiles + NCONS - 1) / NCONS) * reps + g - 1) / g; + return roundsN < rounds1 ? NCONS : 1; + } +} +__device__ __forceinline__ int gemm_groups(int tiles, int reps) { + const int eff = gemm_eff(tiles, reps); + return (tiles + eff - 1) / eff; +} +__device__ __forceinline__ int gemm_consumer_id() { + const int g = warpgroup::groupid(); + if (g == 1) return -1; + const int c = g == 0 ? 0 : g - 1; + return c < NCONS ? c : -1; +} +__device__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][BATCH_PAD]; + __shared__ int moe_gidx_stage[NCONS][BATCH_PAD]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + a_s[s], a, {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + a_s[s], a, {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / BATCH_PAD) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, BATCH_PAD> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], a_s[s]); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * BATCH_PAD; + if (epi == EPI_BIAS) { + warpgroup::store(o_sf, accum); + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + dest[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(o_sf[int2{col, row}] + + __bfloat162float(bias[base + col])); + } + warpgroup::sync(1 + cons); + continue; + } + warpgroup::store(o_s[cons], accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_s[cons][int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__device__ void gemm_narrow( + const weight_gl &w, const act_gl_n &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = PAD_N, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][PAD_N]; + __shared__ int moe_gidx_stage[NCONS][PAD_N]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + PAD_N - 1) / PAD_N; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + PAD_N - 1) / PAD_N; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / PAD_N) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, PAD_N> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], + *reinterpret_cast(&a_s[s])); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * PAD_N; + out_tile_n &o_sn = *reinterpret_cast(&o_s[cons]); + warpgroup::store(o_sn, accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_sn[int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&w_s)[STAGES * NCONS] = + alloc.allocate(); + act_tile (&a_s)[STAGES] = alloc.allocate(); + out_tile (&o_s)[NCONS] = alloc.allocate(); +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 + static_assert(NCONS * sizeof(out_tile) >= sizeof(out_tile_f), + "aliased head staging must fit the output tiles"); + out_tile_f &o_sf = *reinterpret_cast(&o_s[0]); +#else + out_tile_f (&o_sf) = alloc.allocate(); +#endif + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_WARPS + 2]; +#if MK2_PREFILL_ONLY + bf16 *const route_stage = reinterpret_cast(&o_s[0]); + static_assert(sizeof(o_s) >= D_MODEL * sizeof(bf16), + "router row stage must fit the output tiles"); +#else + __shared__ __align__(16) uint64_t sort_keys[SORT_N]; + bf16 *const route_stage = reinterpret_cast(&sort_keys[0]); + static_assert(sizeof(sort_keys) >= D_MODEL * sizeof(bf16), + "router row stage must fit the sampler workspace"); +#endif + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS * NCONS, 0); + for (int c = 0; c < CONSUMER_WARPS * NCONS; ++c) arrive(drained[s]); + } + } + int gemm_iter = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map); + grid_barrier(g.barrier); + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map); + grid_barrier(g.barrier); + mamba_gated_norm(L, bb->zxbcdt, bb->gated); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, nullptr, 1, BATCH, + bb->moe_w, BB_MOE_INTER / OUT_ATOM); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + 1, BATCH, bb->moe_w, D_MODEL / OUT_ATOM); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_deterministic(bb->moe_out, idx, + (1u << BB_EXPERTS) - 1u); + grid_barrier(g.barrier); + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int original_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + bool identical_cohort = false; + if constexpr (BATCH > 1) { + if (bb->native_pf_cohort_check && bb->native_prefill_reset && + bb->pf_n_seq == BATCH && !bb->pf_fuse_decode && + !bb->pf_layer_out && !bb->moe_w_inject && !bb->moe_w_audit) { + if (blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[0] = 0; + grid_barrier(g.barrier); + const auto *bits = reinterpret_cast(bb->prompt); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < original_total * D_MODEL; idx += gridDim.x * THREADS) { + const int canonical = (idx / (BATCH * D_MODEL)) * BATCH * D_MODEL + idx % D_MODEL; + if (bits[idx] != bits[canonical]) atomicOr(bb->native_pf_cohort_check, 1u); + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + if (bb->pf_seq_slot[row] != row || bb->slot_map[row] != row) + atomicOr(bb->native_pf_cohort_check, 1u); + grid_barrier(g.barrier); + identical_cohort = bb->native_pf_cohort_check[0] == 0; + } + } + if (!identical_cohort) { + if (bb->native_pf_cohort_check && blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[1] = 0; + grid_barrier(g.barrier); + return; + } + const int n_total = original_total / BATCH; + if (bb->native_pf_cohort_check && blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[1] = identical_cohort ? 1 : BATCH; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int reset_slots = identical_cohort ? 1 : BATCH; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * reset_slots * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % reset_slots; + const MambaLayer &L = bb->mamba[row_layer / reset_slots]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = original_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + const int src = identical_cohort ? (idx / D_MODEL) * BATCH * D_MODEL + idx % D_MODEL : idx; + bb->pf_h[idx] = bb->prompt[src]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = identical_cohort ? 1 : (bb->pf_n_seq > 0 ? bb->pf_n_seq : 1); + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T, + (!bb->pf_layer_out && layer > 0 && + bb->moe_of_layer[layer - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + if (bb->native_pf_expert_masks) + for (int i = blockIdx.x * THREADS + threadIdx.x; i < n_tok; + i += gridDim.x * THREADS) + bb->native_pf_expert_masks[i] = 0u; + if (bb->pf_moe_ecnt) + for (int e = blockIdx.x * THREADS + threadIdx.x; e < BB_EXPERTS; + e += gridDim.x * THREADS) + bb->pf_moe_ecnt[e] = 0; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + for(size_t j=blockIdx.x*THREADS+threadIdx.x;jpf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + mamba_dtda_prefill(L, bb->pf_zxbcdt, bb->pf_dtda, T); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_dtda, + bb->pf_gated, T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + for(size_t j=blockIdx.x*THREADS+threadIdx.x;jpf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, + T, bb->native_pf_expert_masks, bb->pf_moe_gath, + bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, route_stage); + } + grid_barrier(g.barrier); + int max_ecnt = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + max_ecnt = max(max_ecnt, bb->pf_moe_ecnt[e]); + const bool narrow = BATCH_PAD > PAD_N && max_ecnt <= PAD_N; + const int etile = narrow ? PAD_N : BATCH_PAD; + int nt_sum = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_sum += (bb->pf_moe_ecnt[e] + etile - 1) / etile; + const int up_steps = + gemm_groups(BB_MOE_INTER / OUT_ATOM, nt_sum) * + nt_sum * (D_MODEL / KCHUNK); + const int shup_steps = + gemm_groups(BB_SHARED_INTER / OUT_ATOM, n_tok) * + n_tok * (D_MODEL / KCHUNK); + const int down_steps = + gemm_groups(D_MODEL / OUT_ATOM, nt_sum) * + nt_sum * (BB_MOE_INTER / KCHUNK); + const int shdown_steps = + gemm_groups(D_MODEL / OUT_ATOM, n_tok) * + n_tok * (BB_SHARED_INTER / KCHUNK); + if (narrow) + gemm_narrow(L.exp_up, bb->pf_moe_gath_vn, + BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, PAD_N, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + else + gemm(L.exp_up, bb->pf_moe_gath_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, BATCH_PAD, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + up_steps, up_steps + shup_steps); + grid_barrier(g.barrier); + if (narrow) + gemm_narrow(L.exp_down, bb->pf_moe_hc_vn, + BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, PAD_N, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + else + gemm(L.exp_down, bb->pf_moe_hc_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + down_steps, down_steps + shdown_steps); + grid_barrier(g.barrier); + if (bb->pf_layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * D_MODEL; idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = + bb_moe_fold_deterministic(bb->pf_moe_out, idx, + moe_row_mask(bb->pf_moe_w, idx / D_MODEL)); + grid_barrier(g.barrier); + } + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, + identical_cohort ? bb->pf_normed : bb->pf_final, ln_scratch, + n_total, + (!bb->pf_layer_out && n_backbone_layers > 0 && + bb->moe_of_layer[n_backbone_layers - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + grid_barrier(g.barrier); + if (identical_cohort) { + const int worker = blockIdx.x * THREADS + threadIdx.x; + const int workers = gridDim.x * THREADS; + for (int i = worker; i < original_total * D_MODEL; i += workers) + bb->pf_final[i] = bb->pf_normed[(i / (BATCH * D_MODEL)) * D_MODEL + i % D_MODEL]; + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = worker; i < 14 * (BATCH - 1) * CV; i += workers) { + const int cell = i % CV, pair = i / CV; + const int row = pair % (BATCH - 1) + 1; + const MambaLayer &L = bb->mamba[pair / (BATCH - 1)]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = L.conv_state[cell]; + } + const int kv_columns = BB_KV_HEADS * BB_HEAD_DIM; + for (size_t i = worker; i < size_t(3) * 2 * (BATCH - 1) * n_total * kv_columns; i += workers) { + size_t q = i / kv_columns; + const int col = i % kv_columns; + const int pos = q % n_total; q /= n_total; + const int row = q % (BATCH - 1) + 1; q /= BATCH - 1; + const int kind = q % 2, layer = q / 2; + const AttnLayer &L = bb->attn[layer]; + const int offset = kv_offset(L, pos, col / BB_HEAD_DIM) + col % BB_HEAD_DIM; + kv_slot(L, row, kind, pos)[offset] = kv_slot(L, 0, kind, pos)[offset]; + } + grid_barrier(g.barrier); + } + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } +#if MK2_PREFILL_ONLY + return; +#else + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, bb->phoneme_b); + grid_barrier(g.barrier); + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + if (threadIdx.x == 0) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = 0; n < BB_PHONEME; ++n) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + bb->phoneme_out[row] = arg; + } + __syncthreads(); + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + for (int i = lane; i < D_HEAD; i += 32) { + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + for (int t = 0; t <= pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + for (int i = lane; i < D_HEAD; i += 32) { + float acc = 0.f; + for (int t = 0; t <= pos; ++t) + acc += sc[t] * __bfloat162float( + g.vcache[cache_base + t * D_HEAD + i]); + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + i] = + __float2bfloat16(acc); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, g.sampled_pos, + sort_keys, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { +#ifdef ABLATE_SAMPLE + __syncthreads(); +#endif + } else { + grid_barrier(g.barrier); + } + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x < BATCH && blockIdx.x == 0) { + const int code = g.sampled_pos[threadIdx.x]; + g.sampled[threadIdx.x * N_POS + pos] = code; + if (g.native_codes) + g.native_codes[threadIdx.x * N_POS + pos] = (long long)code; + } + grid_barrier(g.barrier); + } + if (g.native_seq_len && blockIdx.x == 0 && threadIdx.x < BATCH) + g.native_seq_len[threadIdx.x] += 1; + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +#endif +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 2 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK), + 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = MK2_PREFILL_ONLY + ? out->sms * out->blocks_per_sm + : std::min(out->sms * out->blocks_per_sm, D_FFN / OUT_ATOM); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + 0, + nullptr, + nullptr, + nullptr, nullptr, nullptr, nullptr, nullptr, 0, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "dedup_native_prefill_impl.inc" diff --git a/deploy/nemotron_native/kernels/b2048/regions/prefill/native/prefill/mk2_abi.cuh b/deploy/nemotron_native/kernels/b2048/regions/prefill/native/prefill/mk2_abi.cuh new file mode 100644 index 0000000..95a253b --- /dev/null +++ b/deploy/nemotron_native/kernels/b2048/regions/prefill/native/prefill/mk2_abi.cuh @@ -0,0 +1,82 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} diff --git a/deploy/nemotron_native/kernels/b2048/regions/prefill/native/prefill/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/b2048/regions/prefill/native/prefill/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/b2048/regions/prefill/native/prefill/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/b2048/regions/prefill/native/prefill/prefill.cu b/deploy/nemotron_native/kernels/b2048/regions/prefill/native/prefill/prefill.cu new file mode 100644 index 0000000..205c7f5 --- /dev/null +++ b/deploy/nemotron_native/kernels/b2048/regions/prefill/native/prefill/prefill.cu @@ -0,0 +1,18 @@ +#define MK2_BATCH 2048 +#define MK2_BATCH_TILE 64 +#define MK2_PREFILL_ONLY 1 +#define STAGES 3 +#define KCHUNK 128 +#define MK2_GEMM_CONSUMERS 3 +#define MK2_THREADS 512 +#define MK2_NO_MAIN 1 +#define MK2_DIRECT_INPUT 1 +#define MK2_B1_LOCAL_BARRIERS 1 +#define MK2_SKIP_HIDDEN_SNAPSHOT 1 +#define MK2_ACTIVE_EXPERTS 1 +#define MK2_SSM_SPLITS 1 +#define PF_SSM_PW_MULTI 12 +#define PF_ACTIVE_TOKEN_EXPERTS 1 +#include +#include +#include "dedup_predictor_mk2.cuh" diff --git a/deploy/nemotron_native/kernels/b2048/regions/prefill/prefill.cu b/deploy/nemotron_native/kernels/b2048/regions/prefill/prefill.cu new file mode 100644 index 0000000..f2e057b --- /dev/null +++ b/deploy/nemotron_native/kernels/b2048/regions/prefill/prefill.cu @@ -0,0 +1,20 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace b2048_prefill { +#include "native/prefill/prefill.cu" +} diff --git a/deploy/nemotron_native/kernels/b2048/regions/stage0/native/stage0/mk2_abi.cuh b/deploy/nemotron_native/kernels/b2048/regions/stage0/native/stage0/mk2_abi.cuh new file mode 100644 index 0000000..bf1e415 --- /dev/null +++ b/deploy/nemotron_native/kernels/b2048/regions/stage0/native/stage0/mk2_abi.cuh @@ -0,0 +1,85 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_set_native_predictor_active_mask(MK2Predictor *p, const unsigned char *mask); +int mk2_native_startup_skip_abi_version(); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *, int grid); diff --git a/deploy/nemotron_native/kernels/b2048/regions/stage0/native/stage0/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/b2048/regions/stage0/native/stage0/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/b2048/regions/stage0/native/stage0/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/b2048/regions/stage0/native/stage0/native_prefill_impl.inc b/deploy/nemotron_native/kernels/b2048/regions/stage0/native/stage0/native_prefill_impl.inc new file mode 100644 index 0000000..a254308 --- /dev/null +++ b/deploy/nemotron_native/kernels/b2048/regions/stage0/native/stage0/native_prefill_impl.inc @@ -0,0 +1,198 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 16 || BATCH == 32 || BATCH == 64)) + throw std::runtime_error("native prefill requires matching batch 1/16/32/64"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (padded > INT_MAX / ((BB_EXPERTS + 1) * D_MODEL)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + float *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(float))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(float))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} diff --git a/deploy/nemotron_native/kernels/b2048/regions/stage0/native/stage0/predictor_mk2.cu b/deploy/nemotron_native/kernels/b2048/regions/stage0/native/stage0/predictor_mk2.cu new file mode 100644 index 0000000..fe71f6f --- /dev/null +++ b/deploy/nemotron_native/kernels/b2048/regions/stage0/native/stage0/predictor_mk2.cu @@ -0,0 +1,6111 @@ +#include "../../../../../inflight/progress.cuh" +#if MK2_INFLIGHT +#include +#endif +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#define MK2_FUSE_CONV (MK2_SSM_SPLITS == 1) +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +constexpr int BATCH_PAD = 64; +constexpr int DECODE_TILES = BATCH / BATCH_PAD; +constexpr int DECODE_ROWS = BATCH_PAD; +static_assert(BATCH == 512 || BATCH == 1024 || BATCH == 2048); +static_assert(BATCH % BATCH_PAD == 0); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +#ifndef NMULT +#define NMULT 1 +#endif +#ifndef SSM_PREFETCH_MAX +#define SSM_PREFETCH_MAX 8 +#endif +constexpr int NM_MAX = (NMULT > 1 && DECODE_TILES % NMULT == 0) ? NMULT : 1; +#ifndef MK2_OUT_MULT +#define MK2_OUT_MULT 1 +#endif +constexpr int MT = MK2_OUT_MULT; +static_assert(MT >= 1 && MT <= NM_MAX, "an item holds at most NM_MAX accumulators"); +constexpr int SLOTS = 1 + NM_MAX; +static_assert(MT + 1 <= SLOTS); +constexpr int THREADS = 256; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +constexpr int NORM_SCRATCH_STRIDE = 2 * WARPS + 2; +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_ACC_F32 }; +#define EPI_ROW(idx) ((OUT_ATOM == 64 && BATCH_PAD == 64) \ + ? (((idx) & 15) | ((((idx) >> 6) & 3) << 4)) : ((idx) / (OUT_ATOM / 8))) +#define EPI_OCT(idx) ((OUT_ATOM == 64 && BATCH_PAD == 64) \ + ? ((((idx) >> 4) & 3) | ((((idx) >> 8) & 1) << 2)) : ((idx) % (OUT_ATOM / 8))) +struct alignas(16) bf16_oct { bf16 x0, x1, x2, x3, x4, x5, x6, x7; }; +static_assert(sizeof(bf16_oct) == 16); +#include +constexpr int R9_STATE_MODE=1; +using r9_state_t=bf16; +__device__ __forceinline__ int r9_scale_index(int head,int channel,int n) {return (((head/8)*4+(head%8)/2)*4+channel/6)*128+n;} +__device__ __forceinline__ float r9_load(r9_state_t v,float scale) { +return __bfloat162float(v); +} +__device__ __forceinline__ r9_state_t r9_store(float x,float scale) { +return __float2bfloat16(x); +} +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; + const float *r8_initial_ssm = nullptr; + const unsigned int *r8_initial_ready = nullptr; + const int *r8_seq_len = nullptr; + int r8_initial_positions = 0; + const float *r9_scales = nullptr; + const int *r9_map = nullptr; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + bf16 *r6_packed = nullptr; + int *r6_map = nullptr,*r6_counts = nullptr; + act_gl *r6_views = nullptr; + int *r6_lut = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +#define BARRIER_GEN_WORD 32 +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + constexpr unsigned long long GRID_BLOCKS = 132ull; + auto *wide_state = reinterpret_cast(state); + const unsigned long long arrival = atomicAdd(wide_state, 1ull); + const unsigned int rank = (unsigned int)(arrival % GRID_BLOCKS); + if (rank != GRID_BLOCKS - 1u) { + const unsigned long long target = arrival - rank + GRID_BLOCKS; + volatile unsigned long long *counter = wide_state; + while (*counter < target) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / THREADS; + float wv[PER_THREAD]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) wv[j] = weight[threadIdx.x + j * THREADS]; + float v[PER_THREAD], v_next[PER_THREAD]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + v[j] = __bfloat162float( + src[(size_t)blockIdx.x * D_MODEL + threadIdx.x + j * THREADS]); + for (int row = blockIdx.x, row_iter = 0; row < BATCH; + row += gridDim.x, ++row_iter) { + float *row_scratch = scratch + (row_iter & 1) * NORM_SCRATCH_STRIDE; + const int next_row = row + (int)gridDim.x; + if (next_row < BATCH) { + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + v_next[j] = __bfloat162float( + src[(size_t)next_row * D_MODEL + threadIdx.x + j * THREADS]); + } + float sum = 0.f, sq = 0.f; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) { + row_scratch[warp] = sum; + row_scratch[WARPS + warp] = sq; + } + __syncthreads(); + float s = 0.f, q = 0.f; + #pragma unroll + for (int w = 0; w < WARPS; ++w) { + s += row_scratch[w]; + q += row_scratch[WARPS + w]; + } + const float mean = s / D_MODEL; + const float inv = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * wv[j]); + } + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) v[j] = v_next[j]; + } +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, bf16 *snapshot = nullptr, int snapshot_layer = -1, + float *snapshot_fp32 = nullptr, float *out_fp32 = nullptr, + const float *fold_src = nullptr, const float *fold_routes = nullptr) { + constexpr int PER_THREAD = D_MODEL / THREADS; + bf16 wv[PER_THREAD]; + extern __shared__ int __shm[]; + bf16 *const row_stage = reinterpret_cast(__shm); + constexpr int ROW_SLOT = 2 * D_MODEL; + constexpr int ROW_VECTORS = D_MODEL / 8; + if ((int)blockIdx.x < n_rows && threadIdx.x < ROW_VECTORS) { + const int i = 8 * threadIdx.x; + const size_t at = (size_t)blockIdx.x * D_MODEL + i; + const uint32_t rs = static_cast( + __cvta_generic_to_shared(row_stage + D_MODEL + i)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;" + :: "r"(rs), "l"(r + at) : "memory"); + if (!fold_src) { + const uint32_t hs = static_cast( + __cvta_generic_to_shared(row_stage + i)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;" + :: "r"(hs), "l"(h + at) : "memory"); + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + wv[j] = weight[threadIdx.x + j * THREADS]; + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); + __syncthreads(); + for (int row = blockIdx.x, row_iter = 0; row < n_rows; + row += gridDim.x, ++row_iter) { + bf16 *const current = row_stage + (row_iter & 1) * ROW_SLOT; + const int next_row = row + (int)gridDim.x; + if (next_row < n_rows) { + bf16 *const next = row_stage + ((row_iter + 1) & 1) * ROW_SLOT; + if (threadIdx.x < ROW_VECTORS) { + const int i = 8 * threadIdx.x; + const size_t at = (size_t)next_row * D_MODEL + i; + const uint32_t rs = static_cast( + __cvta_generic_to_shared(next + D_MODEL + i)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;" + :: "r"(rs), "l"(r + at) : "memory"); + if (!fold_src) { + const uint32_t hs = static_cast( + __cvta_generic_to_shared(next + i)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;" + :: "r"(hs), "l"(h + at) : "memory"); + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } + float *row_scratch = scratch + (row_iter & 1) * NORM_SCRATCH_STRIDE; + const size_t base = (size_t)row * D_MODEL; + float fv[PER_THREAD]; + if (fold_src) { + unsigned mask = 0u; + const float *rw = fold_routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) + if (rw[e] != 0.f) mask |= 1u << e; + const float *pe = fold_src + + (size_t)row * (BB_EXPERTS + 1) * D_MODEL + threadIdx.x; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) fv[j] = 0.f; + #pragma unroll + for (int e = 0; e <= BB_EXPERTS; ++e) { + const bool take = e == BB_EXPERTS || ((mask >> e) & 1u); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + fv[j] = __fadd_rn(fv[j], + take ? pe[(size_t)e * D_MODEL + j * THREADS] : 0.f); + } + } + float v[PER_THREAD]; + float sq = 0.f; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + const bf16 hv = fold_src ? __float2bfloat16(fv[j]) : current[i]; + if (fold_src) h[base + i] = hv; + if (snapshot) + snapshot[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = hv; + if (snapshot_fp32) + snapshot_fp32[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = + __bfloat162float(hv); + v[j] = __bfloat162float(hv) + + __bfloat162float(current[D_MODEL + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) row_scratch[warp] = sq; + __syncthreads(); + float s = 0.f; + #pragma unroll + for (int w = 0; w < WARPS; ++w) s += row_scratch[w]; + const float inv = rsqrtf(s / D_MODEL + RMS_EPS); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + const bf16 normed = __float2bfloat16( + __bfloat162float(wv[j]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + out[base + i] = normed; + if (out_fp32) out_fp32[base + i] = __bfloat162float(normed); + } + if (next_row < n_rows) { + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); + __syncthreads(); + } + } +} +__device__ __forceinline__ bf16 mamba_conv_channel( + const MambaLayer &L, bf16 *state_slot, const bf16 *zrow, int c) { + const unsigned ch = (unsigned)c; + const bf16 *w = L.conv_w + ch * (unsigned)MB_CONV_K; + bf16 *state = state_slot + ch * (unsigned)(MB_CONV_K - 1); + const bf16 rawb = zrow[c]; + const float raw = __bfloat162float(rawb); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = rawb; + return __float2bfloat16(acc); +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map MK2_PROGRESS_PARAMETER(progress_mask)) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + MK2_PROGRESS_GUARD(progress_mask, row); + const int c = idx - row * MB_CONV_DIM; + xbc_act[(size_t)row * MB_CONV_DIM + c] = mamba_conv_channel( + L, L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride, + zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER, c); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map MK2_PROGRESS_PARAMETER(progress_mask)) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int sb4 = (lane >> 4) & 1, sb3 = (lane >> 3) & 1; + const int sb2 = (lane >> 2) & 1, sb1 = (lane >> 1) & 1; + const int ssm_own_ch = 6 * sb4 + 3 * sb3 + 2 * sb2 + sb1; + const bool ssm_owner = (2 * sb2 + sb1) < 3 && (lane & 1) == 0; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + constexpr bool FUSE_GNORM = (split == 1); + extern __shared__ int __shm[]; + bf16 (*const gn_stage)[MB_GROUP_SIZE] = + reinterpret_cast(&__shm[0]); + constexpr bool FUSE_CONV = (MK2_FUSE_CONV != 0); + bf16 (*const conv_x)[MB_GROUP_SIZE] = gn_stage + WARPS; + constexpr int SSM_NSTATE = MB_STATE / 32; + constexpr int SSM_PF = channels < SSM_PREFETCH_MAX ? channels : SSM_PREFETCH_MAX; + constexpr int SSM_GROUP_BYTES = SSM_PF * MB_STATE * (int)sizeof(r9_state_t); + constexpr bool SSM_SMEM_OK = R9_STATE_MODE == 1 && channels == 2 * SSM_PF && + SSM_GROUP_BYTES % 512 == 0; + constexpr int SSM_RING = 4; + constexpr int SSM_GROUPS_PER_HEAD = channels / SSM_PF; + constexpr int SSM_PAIRS = MB_HEADS_PER_GROUP * SSM_GROUPS_PER_HEAD; + static_assert(!SSM_SMEM_OK || (SSM_RING & (SSM_RING - 1)) == 0, + "the flat (head, channel group) ring index is masked"); + static_assert(!SSM_SMEM_OK || SSM_PAIRS >= SSM_RING, + "the prologue must be able to fill the ring"); + static_assert(!SSM_SMEM_OK || + 2 * WARPS * MB_GROUP_SIZE * sizeof(bf16) + + WARPS * SSM_RING * SSM_GROUP_BYTES <= STAGES * SLOTS * sizeof(weight_tile)); + constexpr int SSM_F32_RING = 2; + constexpr int SSM_F32_GROUP_BYTES = SSM_PF * MB_STATE * (int)sizeof(float); + constexpr bool SSM_F32_OK = SSM_SMEM_OK && SSM_F32_GROUP_BYTES % 512 == 0 && + (SSM_F32_RING & (SSM_F32_RING - 1)) == 0 && SSM_PAIRS >= SSM_F32_RING && + SSM_F32_RING * SSM_F32_GROUP_BYTES <= SSM_RING * SSM_GROUP_BYTES; + char *const ssm_stage = reinterpret_cast(conv_x + WARPS); + const int pairs = BATCH * MB_GROUPS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_GROUPS * split); + MK2_PROGRESS_GUARD(progress_mask, row); + const int group = (pair / split) % MB_GROUPS; + const int first_channel = (pair % split) * channels; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const bool r8_initial = L.r8_initial_ssm && L.r8_initial_ready && + *L.r8_initial_ready == 1 && L.r8_seq_len[row] == L.r8_initial_positions; + r9_state_t *st = reinterpret_cast(L.ssm_state) + + (size_t)slot_map[row] * L.ssm_slot_stride; + char *const warp_ssm = + ssm_stage + (size_t)warp * SSM_RING * SSM_GROUP_BYTES; + auto issue_ssm = [&](int head_in_group, int channel0, int buffer) { + const int head = group * MB_HEADS_PER_GROUP + head_in_group; + const r9_state_t *src = st + + (size_t)(head * MB_HEAD_DIM + channel0) * MB_STATE; + const uint32_t dst = (uint32_t)__cvta_generic_to_shared( + warp_ssm + (size_t)buffer * SSM_GROUP_BYTES + lane * 16); + #pragma unroll + for (int chunk = 0; chunk < SSM_GROUP_BYTES / 512; ++chunk) + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: + "r"(dst + (uint32_t)chunk * 512), + "l"(reinterpret_cast(src) + + (size_t)chunk * 512 + lane * 16) : "memory"); + asm volatile("cp.async.commit_group;\n" ::: "memory"); + }; + auto issue_ssm_f32 = [&](int head_in_group, int channel0, int buffer) { + const int head = group * MB_HEADS_PER_GROUP + head_in_group; + const float *src = L.r8_initial_ssm + + (size_t)(head * MB_HEAD_DIM + channel0) * MB_STATE; + const uint32_t dst = (uint32_t)__cvta_generic_to_shared( + warp_ssm + (size_t)buffer * SSM_F32_GROUP_BYTES + lane * 16); + #pragma unroll + for (int chunk = 0; chunk < SSM_F32_GROUP_BYTES / 512; ++chunk) + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: + "r"(dst + (uint32_t)chunk * 512), + "l"(reinterpret_cast(src) + + (size_t)chunk * 512 + lane * 16) : "memory"); + asm volatile("cp.async.commit_group;\n" ::: "memory"); + }; + if constexpr (SSM_SMEM_OK) { + if (!r8_initial) { + #pragma unroll + for (int k = 0; k < SSM_RING; ++k) + issue_ssm(k / SSM_GROUPS_PER_HEAD, + (k % SSM_GROUPS_PER_HEAD) * SSM_PF, k); + } else if constexpr (SSM_F32_OK) { + #pragma unroll + for (int k = 0; k < SSM_F32_RING; ++k) + issue_ssm_f32(k / SSM_GROUPS_PER_HEAD, + (k % SSM_GROUPS_PER_HEAD) * SSM_PF, k); + } + } + float Braw[MB_STATE / 32]; + float Cv[MB_STATE / 32]; + const bf16 *xrow; + if constexpr (FUSE_CONV) { + bf16 *const cslot = + L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride; + const bf16 *const zrow = + zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER; + static_assert(MB_CONV_K == 4, "the state record is three 32-bit words"); + static_assert(MB_GROUP_SIZE % 64 == 0 && MB_STATE % 64 == 0, + "one paired round covers 64 channels of a warp"); + auto wide_lo = [](unsigned int v) { return __uint_as_float(v << 16); }; + auto wide_hi = [](unsigned int v) { return __uint_as_float(v & 0xffff0000u); }; + struct ConvCursor { + const uint2 *w; unsigned int *s; + const unsigned int *z; const unsigned int *b; + }; + auto conv_cursor = [&](int c) { + const unsigned int cu = (unsigned int)c; + return ConvCursor{ + reinterpret_cast( + L.conv_w + cu * (unsigned int)MB_CONV_K), + reinterpret_cast( + cslot + cu * (unsigned int)(MB_CONV_K - 1)), + reinterpret_cast(zrow + cu), + reinterpret_cast(L.conv_b + cu)}; + }; + auto conv_two = [&](uint2 w0, uint2 w1, unsigned int s0, + unsigned int s1, unsigned int s2, + unsigned int zp, unsigned int bp, float &a0, + float &a1, unsigned int &n0, unsigned int &n1, + unsigned int &n2) { + a0 = wide_lo(bp); + a0 += wide_lo(w0.x) * wide_lo(s0); + a0 += wide_hi(w0.x) * wide_hi(s0); + a0 += wide_lo(w0.y) * wide_lo(s1); + a0 += wide_hi(w0.y) * wide_lo(zp); + a0 = a0 / (1.f + __expf(-a0)); + a1 = wide_hi(bp); + a1 += wide_lo(w1.x) * wide_hi(s1); + a1 += wide_hi(w1.x) * wide_lo(s2); + a1 += wide_lo(w1.y) * wide_hi(s2); + a1 += wide_hi(w1.y) * wide_hi(zp); + a1 = a1 / (1.f + __expf(-a1)); + n0 = (s0 >> 16) | (s1 << 16); + n1 = (zp & 0xffffu) | (s2 << 16); + n2 = (s2 >> 16) | (zp & 0xffff0000u); + }; + struct ConvPairIn { + uint2 w0, w1; + unsigned int s0, s1, s2, zp, bp; + }; + auto conv_pair_fetch = [](const ConvCursor &q, int r) { + const uint2 *const wp = q.w + 64 * r; + const unsigned int *const sp = q.s + 96 * r; + return ConvPairIn{wp[0], wp[1], sp[0], sp[1], sp[2], + q.z[32 * r], q.b[32 * r]}; + }; + auto conv_pair_apply = [&](unsigned int *const sp, + const ConvPairIn &v, float &a0, + float &a1) { + unsigned int n0, n1, n2; + conv_two(v.w0, v.w1, v.s0, v.s1, v.s2, v.zp, v.bp, a0, a1, + n0, n1, n2); + sp[0] = n0; + sp[1] = n1; + sp[2] = n2; + }; + constexpr int QUAD_SPAN = 4 * 32; + static_assert(MB_INNER % 4 == 0 && MB_STATE % 4 == 0 && + MB_GROUP_SIZE % 4 == 0, + "every quad base is a four-channel boundary"); + static_assert(MB_STATE == QUAD_SPAN && + MB_GROUP_SIZE == QUAD_SPAN + 2 * 32, + "one quad round covers B and C; x adds a pair round"); + struct ConvQuad { + const uint4 *w; uint2 *s; const uint2 *z; const uint2 *b; + }; + auto conv_quad_cursor = [&](int c) { + const unsigned int cu = (unsigned int)c; + return ConvQuad{ + reinterpret_cast( + L.conv_w + cu * (unsigned int)MB_CONV_K), + reinterpret_cast( + cslot + cu * (unsigned int)(MB_CONV_K - 1)), + reinterpret_cast(zrow + cu), + reinterpret_cast(L.conv_b + cu)}; + }; + struct ConvQuadIn { + uint4 w01, w23; + uint2 sa, sb, sc, zp, bp; + }; + auto conv_quad_fetch = [](const ConvQuad &q) { + const uint2 *const sp = q.s; + return ConvQuadIn{q.w[0], q.w[1], sp[0], sp[1], sp[2], + q.z[0], q.b[0]}; + }; + auto conv_quad_apply = [&](uint2 *const sp, const ConvQuadIn &v, + float &a0, float &a1, float &a2, + float &a3) { + unsigned int n0, n1, n2, n3, n4, n5; + conv_two(make_uint2(v.w01.x, v.w01.y), + make_uint2(v.w01.z, v.w01.w), + v.sa.x, v.sa.y, v.sb.x, v.zp.x, v.bp.x, + a0, a1, n0, n1, n2); + conv_two(make_uint2(v.w23.x, v.w23.y), + make_uint2(v.w23.z, v.w23.w), + v.sb.y, v.sc.x, v.sc.y, v.zp.y, v.bp.y, + a2, a3, n3, n4, n5); + sp[0] = make_uint2(n0, n1); + sp[1] = make_uint2(n2, n3); + sp[2] = make_uint2(n4, n5); + }; + auto conv_pack = [](float a, float b) { + return (unsigned int)__bfloat16_as_ushort(__float2bfloat16(a)) | + ((unsigned int)__bfloat16_as_ushort(__float2bfloat16(b)) << 16); + }; + const ConvQuad qb = conv_quad_cursor( + MB_INNER + group * MB_STATE + 4 * lane); + const ConvQuad qc = conv_quad_cursor( + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE + 4 * lane); + static_assert(MB_GROUP_SIZE * (int)sizeof(bf16) % 8 == 0, + "each warp's staging row starts 8-byte aligned"); + const ConvQuad qx = conv_quad_cursor( + group * MB_GROUP_SIZE + 4 * lane); + const ConvCursor qt = conv_cursor( + group * MB_GROUP_SIZE + QUAD_SPAN + 2 * lane); + const ConvQuadIn vb = conv_quad_fetch(qb); + const ConvQuadIn vc = conv_quad_fetch(qc); + const ConvQuadIn vx = conv_quad_fetch(qx); + const ConvPairIn vt = conv_pair_fetch(qt, 0); + uint2 bpk, cpk; + { + float a0, a1, a2, a3; + conv_quad_apply(qb.s, vb, a0, a1, a2, a3); + bpk = make_uint2(conv_pack(a0, a1), conv_pack(a2, a3)); + conv_quad_apply(qc.s, vc, a0, a1, a2, a3); + cpk = make_uint2(conv_pack(a0, a1), conv_pack(a2, a3)); + } + { + float a0, a1, a2, a3; + conv_quad_apply(qx.s, vx, a0, a1, a2, a3); + *reinterpret_cast(&conv_x[warp][4 * lane]) = + make_uint2(conv_pack(a0, a1), conv_pack(a2, a3)); + } + { + float a0, a1; + conv_pair_apply(qt.s, vt, a0, a1); + conv_x[warp][QUAD_SPAN + 2 * lane] = __float2bfloat16(a0); + conv_x[warp][QUAD_SPAN + 2 * lane + 1] = __float2bfloat16(a1); + } + const bool qword = ((lane >> 1) & 1) != 0; + const bool qhalf = (lane & 1) != 0; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int src = (lane >> 2) + 8 * j; + const unsigned int bx = __shfl_sync(0xffffffffu, bpk.x, src); + const unsigned int by = __shfl_sync(0xffffffffu, bpk.y, src); + const unsigned int cx = __shfl_sync(0xffffffffu, cpk.x, src); + const unsigned int cy = __shfl_sync(0xffffffffu, cpk.y, src); + const unsigned int bg = qword ? by : bx; + const unsigned int cg = qword ? cy : cx; + Braw[j] = qhalf ? wide_hi(bg) : wide_lo(bg); + Cv[j] = qhalf ? wide_hi(cg) : wide_lo(cg); + } + __syncwarp(); + xrow = &conv_x[warp][0]; + } else { + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Braw[j] = __bfloat162float(Bsec[lane + j * 32]); + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + xrow = xa + (size_t)group * MB_GROUP_SIZE; + } + const int head_lane = group * MB_HEADS_PER_GROUP + + (lane < MB_HEADS_PER_GROUP ? lane : 0); + const float dt_lane = mamba_softplus( + __bfloat162float(dtrow[head_lane]) + L.dt_bias[head_lane]); + const float dA_lane = __expf(L.A[head_lane] * dt_lane); + const float D_lane = L.D[head_lane]; + #pragma unroll 1 + for (int head_in_group = 0; head_in_group < MB_HEADS_PER_GROUP; + ++head_in_group) { + const int head = group * MB_HEADS_PER_GROUP + head_in_group; + const float dt = __shfl_sync(0xffffffffu, dt_lane, head_in_group); + const float dA = __shfl_sync(0xffffffffu, dA_lane, head_in_group); + const float Dh = __shfl_sync(0xffffffffu, D_lane, head_in_group); + const float *st_in = L.r8_initial_ssm + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float Bv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) + Bv[j] = Braw[j] * dt; + constexpr int NSTATE = SSM_NSTATE; + constexpr int SSM_PREFETCH = SSM_PF; + static_assert(channels % SSM_PREFETCH == 0); + for (int g0 = 0; g0 < channels; g0 += SSM_PREFETCH) { + float sv[SSM_PREFETCH][NSTATE]; + float xvr[SSM_PREFETCH]; + if (r8_initial && SSM_F32_OK) { + asm volatile("cp.async.wait_group 1;\n" ::: "memory"); + __syncwarp(); + const int pair = head_in_group * SSM_GROUPS_PER_HEAD + + g0 / SSM_PREFETCH; + const int buffer = pair & (SSM_F32_RING - 1); + const float *const smf = reinterpret_cast( + warp_ssm + (size_t)buffer * SSM_F32_GROUP_BYTES); + static_assert(SSM_PREFETCH % 2 == 0 && MB_HEAD_DIM % 2 == 0, + "x staging vectorises in element pairs"); + const __nv_bfloat162 *const xq = + reinterpret_cast( + xrow + head_in_group * MB_HEAD_DIM + + first_channel + g0); + #pragma unroll + for (int gq = 0; gq < SSM_PREFETCH / 2; ++gq) { + const float2 xp = __bfloat1622float2(xq[gq]); + xvr[2 * gq] = xp.x; + xvr[2 * gq + 1] = xp.y; + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = smf[(size_t)gi * MB_STATE + lane + j * 32]; + } + __syncwarp(); + const int ahead = pair + SSM_F32_RING; + if (ahead < SSM_PAIRS) + issue_ssm_f32(ahead / SSM_GROUPS_PER_HEAD, + (ahead % SSM_GROUPS_PER_HEAD) * SSM_PF, + buffer); + else + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } else if (r8_initial) { + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const size_t prow = (size_t)(first_channel + g0 + gi) * MB_STATE; + xvr[gi] = __bfloat162float( + xrow[head_in_group * MB_HEAD_DIM + + first_channel + g0 + gi]); + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = st_in[prow + lane + j * 32]; + } + } else if constexpr (SSM_SMEM_OK) { + asm volatile("cp.async.wait_group 3;\n" ::: "memory"); + __syncwarp(); + const int pair = head_in_group * SSM_GROUPS_PER_HEAD + + g0 / SSM_PREFETCH; + const int buffer = pair & (SSM_RING - 1); + const r9_state_t *sm = reinterpret_cast( + warp_ssm + (size_t)buffer * SSM_GROUP_BYTES); + static_assert(SSM_PREFETCH % 2 == 0 && MB_HEAD_DIM % 2 == 0, + "x staging vectorises in element pairs"); + const __nv_bfloat162 *const xq = + reinterpret_cast( + xrow + head_in_group * MB_HEAD_DIM + + first_channel + g0); + #pragma unroll + for (int gq = 0; gq < SSM_PREFETCH / 2; ++gq) { + const float2 xp = __bfloat1622float2(xq[gq]); + xvr[2 * gq] = xp.x; + xvr[2 * gq + 1] = xp.y; + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = r9_load( + sm[(size_t)gi * MB_STATE + lane + j * 32], 1.f); + } + __syncwarp(); + const int ahead = pair + SSM_RING; + if (ahead < SSM_PAIRS) + issue_ssm(ahead / SSM_GROUPS_PER_HEAD, + (ahead % SSM_GROUPS_PER_HEAD) * SSM_PF, buffer); + else + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } else { + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + xvr[gi] = __bfloat162float( + xrow[head_in_group * MB_HEAD_DIM + + first_channel + g0 + gi]); + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + { + int p = first_channel+g0+gi; + int mapped = R9_STATE_MODE==4 ? L.r9_map[head*24+p] : head*24+p; + int n = lane+j*32; + float scale = R9_STATE_MODE==1 ? 1.f : L.r9_scales[r9_scale_index(mapped/24,mapped%24,n)]; + sv[gi][j] = r9_load(st[(size_t)mapped*128+n],scale); + } + } + } + if constexpr (SSM_PREFETCH == 12) { + float v6[6]; + #pragma unroll + for (int k = 0; k < 6; ++k) { + float pair_acc[2]; + #pragma unroll + for (int half = 0; half < 2; ++half) { + const int gi = k + 6 * half; + const int p = first_channel + g0 + gi; + const float xv = xvr[gi]; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < NSTATE; ++j) { + const int n = lane + j * 32; + const float s = sv[gi][j] * dA + Bv[j] * xv; + const int mapped = R9_STATE_MODE == 4 + ? L.r9_map[head * 24 + p] : head * 24 + p; + const float scale = R9_STATE_MODE == 1 ? 1.f + : L.r9_scales[r9_scale_index( + mapped / 24, mapped % 24, n)]; + st[(size_t)mapped * 128 + n] = r9_store(s, scale); + acc += s * Cv[j]; + } + pair_acc[half] = acc; + } + const float mine = sb4 ? pair_acc[1] : pair_acc[0]; + const float away = sb4 ? pair_acc[0] : pair_acc[1]; + v6[k] = mine + __shfl_xor_sync(0xffffffff, away, 16); + } + float v4[4]; + #pragma unroll + for (int k = 0; k < 3; ++k) { + const float mine = sb3 ? v6[3 + k] : v6[k]; + const float away = sb3 ? v6[k] : v6[3 + k]; + v4[k] = mine + __shfl_xor_sync(0xffffffff, away, 8); + } + v4[3] = 0.f; + float v2[2]; + #pragma unroll + for (int k = 0; k < 2; ++k) { + const float mine = sb2 ? v4[2 + k] : v4[k]; + const float away = sb2 ? v4[k] : v4[2 + k]; + v2[k] = mine + __shfl_xor_sync(0xffffffff, away, 4); + } + float acc = (sb1 ? v2[1] : v2[0]) + + __shfl_xor_sync(0xffffffff, sb1 ? v2[0] : v2[1], 2); + acc += __shfl_xor_sync(0xffffffff, acc, 1); + if (ssm_owner) { + const int p = first_channel + g0 + ssm_own_ch; + const float xv = __bfloat162float( + xrow[head_in_group * MB_HEAD_DIM + p]); + const bf16 outv = __float2bfloat16(acc + xv * Dh); + if constexpr (FUSE_GNORM) + gn_stage[warp][head_in_group * MB_HEAD_DIM + p] = outv; + else + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = outv; + } + } else { + float cacc[SSM_PREFETCH]; + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const int p = first_channel + g0 + gi; + const float xv = xvr[gi]; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < NSTATE; ++j) { + const int n = lane + j * 32; + const float s = sv[gi][j] * dA + Bv[j] * xv; + const int mapped = R9_STATE_MODE == 4 + ? L.r9_map[head * 24 + p] : head * 24 + p; + const float scale = R9_STATE_MODE == 1 ? 1.f + : L.r9_scales[r9_scale_index( + mapped / 24, mapped % 24, n)]; + st[(size_t)mapped * 128 + n] = r9_store(s, scale); + acc += s * Cv[j]; + } + cacc[gi] = acc; + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const int p = first_channel + g0 + gi; + float acc = cacc[gi]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) { + const bf16 outv = __float2bfloat16(acc + xvr[gi] * Dh); + if constexpr (FUSE_GNORM) + gn_stage[warp][head_in_group * MB_HEAD_DIM + p] = outv; + else + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = outv; + } + } + } + } + } + if constexpr (SSM_SMEM_OK) { + if (!r8_initial || SSM_F32_OK) { + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); + __syncwarp(); + } + } + if constexpr (FUSE_GNORM) { + __syncwarp(); + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32], wg[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float gv = __bfloat162float(gn_stage[warp][lane + j * 32]) * + (z / (1.f + __expf(-z))); + wg[j] = __bfloat162float(L.gnorm_w[idx]); + held[j] = gv; + sq += gv * gv; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16(held[j] * inv * wg[j]); + } + __syncwarp(); + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_CONV_DIM; + const int c = idx - t * MB_CONV_DIM; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v; +#ifdef PF_CONV_ZERO_STATE + if (src < 0) v = 0.f; +#else + if (src < 0) v = __bfloat162float(state[pos + k]); +#endif + else v = __bfloat162float( + zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]); + acc += __bfloat162float(w[k]) * v; + } +#ifndef PF_CONV_NO_SILU + acc = acc / (1.f + __expf(-acc)); +#endif + xbc_act[(size_t)t * MB_CONV_DIM + c] = __float2bfloat16(acc); + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + const int units = n_seq * MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int seq = unit / (MB_HEADS * MB_HEAD_DIM); + const int rem = unit - seq * (MB_HEADS * MB_HEAD_DIM); + const int head = rem / MB_HEAD_DIM; + const int p = rem - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p * MB_STATE; + float sreg[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) sreg[j] = st[lane + j * 32]; + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const bf16 *xa = xbc_act + t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + const bf16 *dtrow = zxbcdt + t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + sreg[j] = sreg[j] * dA + __bfloat162float(Bsec[n]) * dt * xv; + acc += sreg[j] * __bfloat162float(Csec[n]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) st[lane + j * 32] = sreg[j]; + } + return; + } +#if PF_SSM_SPLIT_P + const int units = MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; unit += warps_total) { + const int head = unit / MB_HEAD_DIM; + const int p_only = unit - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; +#else + for (int head = blockIdx.x * WARPS + warp; head < MB_HEADS; + head += warps_total) { + const int group = head / MB_HEADS_PER_GROUP; +#endif + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *st = L.ssm_state + (size_t)slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE +#if PF_SSM_SPLIT_P + + (size_t)p_only * MB_STATE +#endif + ; + for (int t = 0; t < n_tokens; ++t) { + const bf16 *xa = xbc_act + (size_t)t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = + zxbcdt + (size_t)t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } +#if PF_SSM_SPLIT_P + { + const int p = p_only; +#else + for (int p = 0; p < MB_HEAD_DIM; ++p) { +#endif + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { +#if PF_SSM_SPLIT_P + const size_t at = (size_t)lane + j * 32; +#else + const size_t at = (size_t)p * MB_STATE + lane + j * 32; +#endif + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } + } +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32], wg[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + wg[j] = __bfloat162float(L.gnorm_w[idx]); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16(held[j] * inv * wg[j]); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len MK2_PROGRESS_PARAMETER(progress_mask)) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + constexpr int Q_PER_WARP = 3; + constexpr int GROUPS_PER_KV = 1; + const int pairs = BATCH * BB_KV_HEADS * GROUPS_PER_KV; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x + warp * gridDim.x; pair < pairs; + pair += gridDim.x * WARPS) { + const int row = pair / (BB_KV_HEADS * GROUPS_PER_KV); + MK2_PROGRESS_GUARD(progress_mask, row); + const int within_row = pair % (BB_KV_HEADS * GROUPS_PER_KV); + const int kv = within_row / GROUPS_PER_KV; + const int q_group = within_row % GROUPS_PER_KV; + const int q_first = q_group * Q_PER_WARP; + const int q_count = q_group == 0 ? Q_PER_WARP : 1; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (q_group == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[h][j] = 0.f; + } + float m[Q_PER_WARP], denom[Q_PER_WARP]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + m[h] = -CUDART_INF_F; + denom[h] = 0.f; + } + constexpr int PF = 8; + const int tpb = L.tokens_per_block; + const size_t elems_per_block = + (size_t)tpb * BB_KV_HEADS * BB_HEAD_DIM; + const int *bt_k = L.block_table + (size_t)(row * 2) * L.max_blocks_per_seq; + const int *bt_v = bt_k + L.max_blocks_per_seq; + const int head_base = kv * tpb * BB_HEAD_DIM; + float qr[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const bf16 *q = base + (kv * kv_per_q + q_first + h) * BB_HEAD_DIM; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qr[h][j] = __bfloat162float(q[lane + j * 32]); + } + for (int page0 = 0; page0 < T; page0 += tpb) { + const int k_page = bt_k[page0 / tpb] & 0x7FFFFFFF; + const int v_page = bt_v[page0 / tpb] & 0x7FFFFFFF; + const bf16 *kbase = L.pool_base + + (size_t)k_page * elems_per_block + + head_base; + const bf16 *vbase = L.pool_base + + (size_t)v_page * elems_per_block + + head_base; + const int tend = page0 + tpb < T ? page0 + tpb : T; + for (int t0 = page0; t0 < tend; t0 += PF) { + const int navail = tend - t0 < PF ? tend - t0 : PF; + float kr[PF][PER_LANE], vr[PF][PER_LANE]; + #pragma unroll + for (int u = 0; u < PF; ++u) { + const int tpos = t0 + (u < navail ? u : navail - 1); +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = vbase + (tpos - page0) * BB_HEAD_DIM; +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new : kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = cur_token + ? v_new : vbase + (tpos - page0) * BB_HEAD_DIM; +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + kr[u][j] = __bfloat162float(kt[i]); + vr[u][j] = __bfloat162float(vt[i]); + } + } + #pragma unroll + for (int u = 0; u < PF; ++u) { + if (u >= navail) break; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + dot += qr[h][j] * kr[u][j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m[h], s); + const float rescale = __expf(m[h] - m_new); + const float w = __expf(s - m_new); + denom[h] = denom[h] * rescale + w; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[h][j] = acc[h][j] * rescale + w * vr[u][j]; + m[h] = m_new; + } + } + } + } + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const float inv = 1.f / denom[h]; + const int head = kv * kv_per_q + q_first + h; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[h][j] * inv); + } + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * BB_KV_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / BB_KV_HEADS; + const int kv = idx - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv); + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv); + for (int i = 0; i < BB_HEAD_DIM; ++i) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + for (int s = 0; s <= pos; ++s) { + const bf16 *kt = kv_slot(L, seq, 0, s) + kv_offset(L, s, kv); + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + const bf16 *vt = kv_slot(L, seq, 1, s) + kv_offset(L, s, kv); + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + wt * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH, bf16 *row_stage = nullptr) { + __shared__ float score[BB_EXPERTS]; + constexpr int CH = D_MODEL / 8; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const char *gsrc = reinterpret_cast( + normed + (size_t)row * D_MODEL); + char *gdst = reinterpret_cast(row_stage); + const uint32_t act_s = (uint32_t)__cvta_generic_to_shared(gdst); + for (int i = threadIdx.x; i < CH; i += THREADS) { + const uint32_t dst = act_s + (uint32_t)i * 16; + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: + "r"(dst), "l"(gsrc + (size_t)i * 16) : "memory"); + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + if ((int)threadIdx.x < BB_EXPERTS) + moe_w[(size_t)row * BB_EXPERTS + threadIdx.x] = 0.f; + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); + __syncthreads(); + float acc0 = 0.f, acc1 = 0.f, acc2 = 0.f; + const bf16 *w0 = gate_w + (size_t)warp * D_MODEL; + const bf16 *w1 = gate_w + (size_t)(warp + WARPS) * D_MODEL; + const bf16 *w2 = gate_w + (size_t)(warp + 2 * WARPS) * D_MODEL; + for (int i = lane; i < D_MODEL; i += 128) { + bf16 sx0, sx1; + move::lds(sx0, act_s + (uint32_t)i * sizeof(bf16)); + move::lds(sx1, act_s + (uint32_t)(i + 32) * sizeof(bf16)); + float x0 = __bfloat162float(sx0); + float x1 = __bfloat162float(sx1); + acc0 += x0 * __bfloat162float(w0[i]); + acc1 += x0 * __bfloat162float(w1[i]); + acc2 += x0 * __bfloat162float(w2[i]); + acc0 += x1 * __bfloat162float(w0[i + 32]); + acc1 += x1 * __bfloat162float(w1[i + 32]); + acc2 += x1 * __bfloat162float(w2[i + 32]); + move::lds(sx0, act_s + (uint32_t)(i + 64) * sizeof(bf16)); + move::lds(sx1, act_s + (uint32_t)(i + 96) * sizeof(bf16)); + x0 = __bfloat162float(sx0); + x1 = __bfloat162float(sx1); + acc0 += x0 * __bfloat162float(w0[i + 64]); + acc1 += x0 * __bfloat162float(w1[i + 64]); + acc2 += x0 * __bfloat162float(w2[i + 64]); + acc0 += x1 * __bfloat162float(w0[i + 96]); + acc1 += x1 * __bfloat162float(w1[i + 96]); + acc2 += x1 * __bfloat162float(w2[i + 96]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + acc0 += __shfl_xor_sync(0xffffffffu, acc0, off); + acc1 += __shfl_xor_sync(0xffffffffu, acc1, off); + acc2 += __shfl_xor_sync(0xffffffffu, acc2, off); + } + if (lane == 0) { + score[warp] = 1.f / (1.f + __expf(-acc0)); + score[warp + WARPS] = 1.f / (1.f + __expf(-acc1)); + score[warp + 2 * WARPS] = 1.f / (1.f + __expf(-acc2)); + } + __syncthreads(); + if (threadIdx.x < 32) { + float remaining = (lane < BB_EXPERTS) + ? score[lane] + __bfloat162float(gate_bias[lane]) + : -CUDART_INF_F; + int pick[BB_TOP_K]; + float total = 0.f; + #pragma unroll + for (int k = 0; k < BB_TOP_K; ++k) { + float bv = remaining; + int bi = lane; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffff, bv, off); + const int oi = __shfl_xor_sync(0xffffffff, bi, off); + if (ov > bv || (ov == bv && oi < bi)) { bv = ov; bi = oi; } + } + pick[k] = bi; + total += score[bi]; + if (lane == bi) remaining = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + #pragma unroll + for (int k = 0; k < BB_TOP_K; ++k) + if (lane == pick[k]) moe_w[row * BB_EXPERTS + lane] = score[lane] * norm; + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert <= BB_EXPERTS; ++expert) + value = __fadd_rn(value, parts[base + (size_t)expert * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ bf16 bb_moe_fold_selected( + const float *parts, const float *weights, int idx) { + const int row = idx / D_MODEL; + const size_t base = (size_t)row * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert < BB_EXPERTS; ++expert) { + const float part = weights[row * BB_EXPERTS + expert] != 0.f + ? parts[base + (size_t)expert * D_MODEL] : 0.f; + value = __fadd_rn(value, part); + } + value = __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ __forceinline__ void sample_argmax_merge( + float &value, int &rank, float other, int other_rank) { + if (other > value || (other == value && other_rank < rank)) { + value = other; + rank = other_rank; + } +} +__device__ __forceinline__ int phoneme_warp_argmax(const bf16 *lg) { + const int lane = threadIdx.x & 31; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = lane; n < BB_PHONEME; n += 32) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, best, off); + const int other_arg = __shfl_down_sync(0xffffffffu, arg, off); + sample_argmax_merge(best, arg, other, other_arg); + } + return arg; +} +__device__ __forceinline__ int sample_block_argmax(float value, int rank) { + __shared__ float warp_value[WARPS]; + __shared__ int warp_rank[WARPS]; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + if (lane == 0) { warp_value[warp] = value; warp_rank[warp] = rank; } + __syncthreads(); + if (warp == 0) { + value = lane < WARPS ? warp_value[lane] : -CUDART_INF_F; + rank = lane < WARPS ? warp_rank[lane] : INT_MAX; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + } + return rank; +} +__device__ __forceinline__ int sample_eos_insertion_rank( + const uint64_t *keys, float eos) { + int lo = 0, hi = TOP_K; + while (lo < hi) { + const int mid = (lo + hi) >> 1; + if (eos > ordered_to_float((uint32_t)(keys[mid] >> 32))) hi = mid; + else lo = mid + 1; + } + return lo; +} +__device__ __forceinline__ void sample_topk_serial( + const uint64_t *keys, float eos, const float *noise, int *out) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float value; + int code; + if (!eos_used && eos > av) { + value = eos; + code = AUDIO_EOS; + eos_used = true; + } else { + value = av; + code = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = value + noise[j]; + if (noised > best_noised) { + best_noised = noised; + best_code = code; + } + } + *out = best_code; +} +__device__ __forceinline__ void sample_topk_parallel( + const uint64_t *keys, float eos, const float *noise, int *out) { + const int j = threadIdx.x; + const bool sorted_nan = j < TOP_K && + isnan(ordered_to_float((uint32_t)(keys[j] >> 32))); + if (__syncthreads_or(sorted_nan)) { + if (threadIdx.x == 0) sample_topk_serial(keys, eos, noise, out); + return; + } + const int lane = threadIdx.x & 31; + int insertion = lane == 0 ? sample_eos_insertion_rank(keys, eos) : 0; + insertion = __shfl_sync(0xffffffffu, insertion, 0); + float value = -CUDART_INF_F; + int rank = INT_MAX; + if (j < TOP_K) { + const int a = j - (j > insertion); + const float v = j == insertion + ? eos : ordered_to_float((uint32_t)(keys[a] >> 32)); + const float noised = v + noise[j]; + if (noised > -CUDART_INF_F) { value = noised; rank = j; } + } + const int winner = sample_block_argmax(value, rank); + if (threadIdx.x == 0) { + if (winner == INT_MAX) *out = 0; + else if (winner == insertion) *out = AUDIO_EOS; + else { + const int a = winner - (winner > insertion); + *out = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + } + } +} +constexpr int SORT_ITEMS_PER_THREAD = SORT_N / THREADS; +struct sampler_storage { + uint64_t keys[2][SORT_N]; +}; +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + sampler_storage &storage, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + uint64_t thread_keys[SORT_ITEMS_PER_THREAD]; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + thread_keys[item] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + int cross_stage = 0; + for (int k = 2; k <= SORT_N; k <<= 1) { + for (int j = k >> 1; j > 0; j >>= 1) { + if (j >= 32) { + const int bank = cross_stage++ & 1; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + storage.keys[bank][n] = thread_keys[item]; + } + __syncthreads(); + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + const uint64_t other = storage.keys[bank][n ^ j]; + const bool take_hi = ((n & k) == 0) == ((n & j) == 0); + const uint64_t mine = thread_keys[item]; + thread_keys[item] = take_hi + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } else { + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + const uint64_t other = __shfl_xor_sync( + 0xffffffffu, thread_keys[item], j); + const bool take_hi = ((n & k) == 0) == ((n & j) == 0); + const uint64_t mine = thread_keys[item]; + thread_keys[item] = take_hi + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } + } + } + const int top_bank = cross_stage & 1; + if (threadIdx.x < TOP_K) + storage.keys[top_bank][threadIdx.x] = thread_keys[0]; + __syncthreads(); + sample_topk_parallel( + storage.keys[top_bank], eos_v, + gumbel + ((size_t)row * N_POS + pos) * TOP_K, + out_codes + row); + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void r6_sparse_map(const float* routes,int* map,int* counts + MK2_PROGRESS_PARAMETER(progress_mask)) { + __shared__ int r7_warp_counts[THREADS/32]; + const int lane=threadIdx.x%32, warp=threadIdx.x/32; + const unsigned lower=(1u<>1; + for(int e=0;e(packed)[(size_t)e*BATCH*row_vectors+i]= + row<0?make_uint4(0,0,0,0):reinterpret_cast(source)[(size_t)row*row_vectors+i%row_vectors]; + } + } +} +__device__ void r6_tile(int work,int output_tiles,const int* lut,int &tile,int &tok, + int mtg) { + const int groups=output_tiles/mtg; + tok=lut[1+work/groups]; + const int e=tok/DECODE_TILES; + tile=e*output_tiles+(work%groups)*mtg; +} +__device__ __forceinline__ void mk2_group_tile( + int work, int grouped, int &tile, int &tok, int &here) { + here = work < grouped ? NM_MAX : 1; + const int unit = work < grouped ? work * NM_MAX : grouped * (NM_MAX - 1) + work; + tile = unit / DECODE_TILES; + tok = unit - tile * DECODE_TILES; +} +__device__ __forceinline__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&p_s)[STAGES][SLOTS], out_tile &o_s, + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + int &tile_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = DECODE_TILES, + int rows_valid = DECODE_ROWS, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const int *r6_lut = nullptr,const int *r6_map = nullptr) { + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + constexpr int QSLOT = MT + NM_MAX; + constexpr bool WIDE = MT > 1 && NM_MAX > 1 && STAGES * SLOTS >= 2 * QSLOT; + constexpr int PSTAGES = WIDE ? 2 : STAGES; + constexpr int PSLOT = WIDE ? QSLOT : SLOTS; + constexpr int ACCN = WIDE ? MT * NM_MAX : NM_MAX; + constexpr int HEX_MT = 3, HEX_NM = 2; + constexpr int HEX_SLOT = HEX_MT + HEX_NM; + auto epi_publish = [](out_tile &st, const rt_fl<16, BATCH_PAD> &src) { + using acc_t = rt_fl<16, BATCH_PAD>; + static_assert(acc_t::height == 1, "one row block per warp"); + static_assert(out_tile::rows == BATCH_PAD && out_tile::cols == OUT_ATOM, + "staging tile must be square for an in-place transpose"); + const int t = ::kittens::laneid(); + const uint32_t base = + static_cast(__cvta_generic_to_shared(&st.data[0])); + const int row0 = 8 * (t >> 4) + (t & 7); + const int col0 = 16 * warpgroup::warpid() + 8 * ((t >> 3) & 1); + #pragma unroll + for (int j = 0; j < acc_t::width; ++j) { + bf16_2 q0 = base_types::convertor::convert( + src.tiles[0][j].data[0]); + bf16_2 q1 = base_types::convertor::convert( + src.tiles[0][j].data[1]); + bf16_2 q2 = base_types::convertor::convert( + src.tiles[0][j].data[2]); + bf16_2 q3 = base_types::convertor::convert( + src.tiles[0][j].data[3]); + move::stsm4t( + out_tile::idx(base, {16 * j + row0, col0}), q0, q1, q2, q3); + } + }; + weight_tile *const pf = &p_s[0][0]; + const bool arena10 = + reinterpret_cast(&o_sf) == + reinterpret_cast(pf + STAGES * SLOTS); + const bool two_d = !r6_lut && out_tiles_per_expert == 0 + && n_tok == DECODE_TILES && DECODE_TILES % NM_MAX == 0; + int n_tiles = compact_tiles ? __popc(active_mask)*compact_tiles : n_out/OUT_ATOM; + constexpr int QTOK = DECODE_TILES / NM_MAX; + const int out_tiles = n_tiles; + const bool tall3 = WIDE && two_d && out_tiles == 24; + const bool split72 = WIDE && two_d && DECODE_TILES == 16 && out_tiles == 72; + const bool split81 = WIDE && two_d && DECODE_TILES == 16 && out_tiles == 81; + const bool split17 = WIDE && two_d && DECODE_TILES == 16 && out_tiles == 17; + const bool mixed40 = WIDE && two_d && DECODE_TILES == 16 && out_tiles == 40; + const int n_tiles_hex_all = (out_tiles / HEX_MT) * (DECODE_TILES / HEX_NM); + const bool hex = WIDE && two_d && arena10 && DECODE_TILES == 16 && DECODE_TILES % HEX_NM == 0 && + (out_tiles == 96 || out_tiles == 81 || out_tiles == 72 || + out_tiles == 40) && + (epi == EPI_STORE || epi == EPI_GELU) && !r6_lut && + out_tiles_per_expert == 0; + const int hex_tok = DECODE_TILES / HEX_NM; + const int hex_wide = (out_tiles == 72 || out_tiles == 81) + ? 128 : n_tiles_hex_all; + const int pslot = hex ? HEX_SLOT : PSLOT; + int q_units = 0, q_rem_base = 0; + int grouped = 0; + if constexpr (WIDE) { + if (two_d) { + q_units = (out_tiles / MT) * QTOK; + q_rem_base = out_tiles - out_tiles % MT; + } + } else if (two_d) { + const int g = (int)gridDim.x; + const int units = n_tiles * DECODE_TILES; + const int rounds = (units + g - 1) / g; + const int cap = units / NM_MAX; + grouped = g * (rounds / NM_MAX); + if (grouped > cap) grouped = cap; + } + const int r6_output_tiles = r6_lut ? n_out/BB_EXPERTS/OUT_ATOM : 0; + const int r6_mtg = (MT > 1 && r6_lut && r6_output_tiles % MT == 0) ? MT : 1; + const int r6_groups = r6_lut ? r6_output_tiles/r6_mtg : 0; + const int r6_rows = r6_lut ? (r6_lut[0] & 0xFFF) : 0; + const int r6_pairs = (r6_lut && WIDE && r6_mtg == MT) ? (r6_lut[0] >> 12) : 0; + const int r6_pair_items = r6_pairs * r6_groups; + const int r6_single_mtg = (r6_lut && WIDE && r6_output_tiles % 3 == 0) + ? 3 : r6_mtg; + const int r6_single_groups = r6_lut ? r6_output_tiles/r6_single_mtg : 0; + const int r6_single_items = r6_lut + ? (r6_rows - 2 * r6_pairs) * r6_single_groups : 0; + if(r6_lut) { + n_tiles=r6_pair_items+r6_single_items; + n_tok=1; + } else if(two_d) { + if constexpr (WIDE) { + if (hex) n_tiles = (out_tiles == 72) ? 256 + : (out_tiles == 81) ? 264 + : (out_tiles == 40) + ? n_tiles_hex_all + hex_tok + : n_tiles_hex_all; + else if (mixed40) n_tiles = 244; + else if (split72) n_tiles = 360; + else if (split81) n_tiles = 384; + else if (split17) n_tiles = 132; + else if (tall3) n_tiles = (out_tiles / 3) * DECODE_TILES; + else n_tiles = q_units + (out_tiles % MT) * QTOK; + } + else n_tiles = grouped + (n_tiles*DECODE_TILES - grouped*NM_MAX); + n_tok=1; } + const int r6_extra = r6_lut ? n_tiles % (int)gridDim.x : 0; + const int r6_skew = r6_lut + ? (r6_single_items < r6_extra ? r6_single_items : r6_extra) : 0; + const int iters = k_total / KCHUNK; + int tile_start = (int)blockIdx.x - tile_io % (int)gridDim.x; + if (tile_start < 0) tile_start += (int)gridDim.x; + tile_io += n_tiles; + int tiles_here = 0; + for (int t = tile_start; t < n_tiles; t += gridDim.x) ++tiles_here; + const int iter_base = iter_io; + iter_io += tiles_here * iters * n_tok; + const int my_tiles = (n_tiles - tile_start + gridDim.x - 1) / gridDim.x; + if (warpgroup::groupid() == 1) { + if (hex) { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) { + if (out_tiles == 81 && work >= hex_wide) { + const int j = work - hex_wide; + const bool pair = j >= 128; + const int nt = pair ? 80 : 48 + 2 * (j / hex_tok); + const int ntok = HEX_NM * (j % hex_tok); + if (pair) { + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + wait(drained[s], (iter / PSTAGES) & 1); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], sizeof(weight_tile) + + HEX_NM * sizeof(act_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot], w, {nt, i}, filled[s]); + #pragma unroll + for (int m = 0; m < HEX_NM; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + 1 + m], a, + {ntok + m, i}, filled[s]); + } + } + } else { + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + wait(drained[s], (iter / PSTAGES) & 1); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], 2 * sizeof(weight_tile) + + HEX_NM * sizeof(act_tile)); + #pragma unroll + for (int c = 0; c < 2; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {nt + c, i}, filled[s]); + #pragma unroll + for (int m = 0; m < HEX_NM; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + 2 + m], a, + {ntok + m, i}, filled[s]); + } + } + } + continue; + } + if (out_tiles == 40 && work >= hex_wide) { + const int ntok = HEX_NM * (work - hex_wide); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + wait(drained[s], (iter / PSTAGES) & 1); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], sizeof(weight_tile) + + HEX_NM * sizeof(act_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot], w, {out_tiles - 1, i}, filled[s]); + #pragma unroll + for (int m = 0; m < HEX_NM; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + 1 + m], a, + {ntok + m, i}, filled[s]); + } + } + continue; + } + if (work >= hex_wide) { + const int j = work - hex_wide; + const int nt = HEX_MT * ((hex_wide + (j >> 1)) / hex_tok); + const int ntok = HEX_NM * ((hex_wide + (j >> 1)) % hex_tok) + (j & 1); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + wait(drained[s], (iter / PSTAGES) & 1); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], HEX_MT * sizeof(weight_tile) + + sizeof(act_tile)); + #pragma unroll + for (int c = 0; c < HEX_MT; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {nt + c, i}, filled[s]); + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + HEX_MT], a, {ntok, i}, filled[s]); + } + } + continue; + } + const int tile_id = HEX_MT * (work / hex_tok); + const int tok = HEX_NM * (work % hex_tok); + rt_fl<16, BATCH_PAD> pacc[HEX_NM]; + #pragma unroll + for (int u = 0; u < HEX_NM; ++u) warp::zero(pacc[u]); + int pend = -1, pend_ph = 0; + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(drained[s], phase); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], HEX_MT * sizeof(weight_tile) + + HEX_NM * sizeof(act_tile)); + #pragma unroll + for (int c = 0; c < HEX_MT; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {tile_id + c, i}, filled[s]); + #pragma unroll + for (int m = 0; m < HEX_NM; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + HEX_MT + m], a, {tok + m, i}, filled[s]); + } + if (pend >= 0) { { + wait(filled[pend], pend_ph); + warpgroup::mma_ABt(pacc[0], pf[pend * pslot + HEX_MT - 1], + pf[pend * pslot + HEX_MT]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + pacc[1], pf[pend * pslot + HEX_MT - 1], + pf[pend * pslot + HEX_MT + 1]); + warpgroup::mma_async_wait(); + } } + pend = s; pend_ph = phase; + } + { + wait(filled[pend], pend_ph); + warpgroup::mma_ABt(pacc[0], pf[pend * pslot + HEX_MT - 1], + pf[pend * pslot + HEX_MT]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + pacc[1], pf[pend * pslot + HEX_MT - 1], + pf[pend * pslot + HEX_MT + 1]); + warpgroup::mma_async_wait(); + } + const int es = iter % PSTAGES; + wait(drained[es], (iter / PSTAGES) & 1); + out_tile_f &q_sf = + *reinterpret_cast(&pf[es * pslot]); + out_tile &q_s = *reinterpret_cast(&q_sf); + const int qbase = (tile_id + HEX_MT - 1) * OUT_ATOM; + constexpr int OUT_OCTS = OUT_ATOM / 8; + #pragma unroll + for (int u = 0; u < HEX_NM; ++u) { + const int grow0 = (tok + u) * BATCH_PAD; + epi_publish(q_s, pacc[u]); + warpgroup::sync(3); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); + idx < OUT_OCTS * rows_valid; idx += 128) { + const int row = idx >> 3, col = (idx & 7) << 3; + const size_t at = + (size_t)(grow0 + row) * dest_stride + qbase + col; + const bf16_oct q = *reinterpret_cast( + &q_s[int2{row, col}]); + if (epi == EPI_STORE) { + *reinterpret_cast(dest + at) = q; + } else { + *reinterpret_cast(dest + at) = bf16_oct{ + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x0))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x1))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x2))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x3))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x4))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x5))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x6))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x7)))}; + } + } + warpgroup::sync(3); + } + } + } else if (warpgroup::laneid() == 0) { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) { + int tile_id = mk2_expert_tile(work, active_mask, compact_tiles), r6_tok=0; + int nm = 1, mt = 1; + if(r6_lut) { + int rwork = work - r6_skew; + if (rwork < 0) rwork += n_tiles; + const bool pair = rwork < r6_pair_items; + const int q = pair ? rwork / r6_groups + : (rwork-r6_pair_items) / r6_single_groups; + const int group = pair ? rwork-q*r6_groups + : rwork-r6_pair_items-q*r6_single_groups; + const int base_item = pair ? 2*q : 2*r6_pairs+q; + const int ltok = r6_lut[1 + base_item]; + r6_tok = ltok; + tile_id = (ltok / DECODE_TILES) * r6_output_tiles + + group * (pair ? r6_mtg : r6_single_mtg); + mt = pair ? r6_mtg : r6_single_mtg; + nm = pair ? NM_MAX : 1; + } + else if(two_d) { + if constexpr (WIDE) { + if (mixed40) { + const bool quad = work < 132; + const int item = quad ? work : work - 132; + if (quad && item < 84) { + tile_id = 2 * (item / 7); + r6_tok = 2 * (item % 7); + } else if (quad) { + const int tail = item - 84; + tile_id = 2 * (12 + tail / 6); + r6_tok = 2 * (tail % 6); + } else if (item < 48) { + tile_id = 2 * (item / 4) + (item % 4) / 2; + r6_tok = 14 + item % 2; + } else { + const int tail = item - 48; + tile_id = 2 * (12 + tail / 8) + (tail % 8) / 4; + r6_tok = 12 + tail % 4; + } + mt = quad ? MT : 1; + nm = quad ? NM_MAX : 1; + } else if (split72) { + if (work < 264) { + tile_id = 2 * (work / 8); + r6_tok = 2 * (work % 8); + mt = MT; nm = NM_MAX; + } else { + const int i = work - 264, c = i >> 2, q = i & 3; + tile_id = 2 * (33 + (c >> 3)) + (q >> 1); + r6_tok = 2 * (c & 7) + (q & 1); + mt = 1; nm = 1; + } + } else if (split81) { + if (work < 264) { + tile_id = 2 * (work / 8); + r6_tok = 2 * (work % 8); + mt = MT; nm = NM_MAX; + } else if (work < 376) { + const int i = work - 264; + tile_id = 2 * (33 + i / 16); + r6_tok = i % 16; + mt = MT; nm = 1; + } else { + tile_id = 80; + r6_tok = 2 * (work - 376); + mt = 1; nm = NM_MAX; + } + } else if (split17) { + if (work < 8) { + tile_id = 14; + r6_tok = work; + mt = 3; nm = 1; + } else if (work < 120) { + const int i = work - 8; + tile_id = i >> 3; + r6_tok = 2 * (i & 7); + mt = 1; nm = NM_MAX; + } else { + const int i = work - 120; + tile_id = 14 + (i >> 2); + r6_tok = 8 + 2 * (i & 3); + mt = 1; nm = NM_MAX; + } + } else if (tall3) { + tile_id = 3 * (work / DECODE_TILES); + r6_tok = work % DECODE_TILES; + mt = 3; + nm = 1; + } else { + const int j = work < q_units ? work : work - q_units; + tile_id = work < q_units ? MT * (j / QTOK) + : q_rem_base + j / QTOK; + r6_tok = NM_MAX * (j % QTOK); + mt = work < q_units ? MT : 1; + nm = NM_MAX; + } + } else mk2_group_tile(work,grouped,tile_id,r6_tok,nm); + } + for (int token_iteration = 0; token_iteration < n_tok; ++token_iteration) + for (int i = 0; i < iters; ++i, ++iter) { + const int tok = (r6_lut || two_d) ? r6_tok : token_iteration; + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(drained[s], phase); + tma::expect_bytes(filled[s], + mt * sizeof(weight_tile) + nm * sizeof(act_tile)); + if (mt == 1) { + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot], w, {tile_id, i}, filled[s]); + } else if (mt == 3) { + #pragma unroll + for (int c = 0; c < 3; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {tile_id + c, i}, filled[s]); + } else { + #pragma unroll + for (int c = 0; c < MT; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {tile_id + c, i}, filled[s]); + } + const int a_col = out_tiles_per_expert && !r6_lut + ? (tile_id / out_tiles_per_expert) * act_k_tiles_per_expert + i + : i; + if (NM_MAX == 1 || nm == NM_MAX) { + #pragma unroll + for (int m = 0; m < NM_MAX; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + mt + m], a, {tok + m, a_col}, filled[s]); + } else { + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + mt], a, {tok, a_col}, filled[s]); + } + } + } + } + } else { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) + for (int token_iteration = 0; token_iteration < n_tok; ++token_iteration) { + int tile_id=mk2_expert_tile(work, active_mask, compact_tiles),tok=token_iteration; + int nm = 1, mt = 1; + if(r6_lut) { + int rwork = work - r6_skew; + if (rwork < 0) rwork += n_tiles; + const bool pair = rwork < r6_pair_items; + const int q = pair ? rwork / r6_groups + : (rwork-r6_pair_items) / r6_single_groups; + const int group = pair ? rwork-q*r6_groups + : rwork-r6_pair_items-q*r6_single_groups; + const int base_item = pair ? 2*q : 2*r6_pairs+q; + const int ltok = r6_lut[1 + base_item]; + tok = ltok; + tile_id = (ltok / DECODE_TILES) * r6_output_tiles + + group * (pair ? r6_mtg : r6_single_mtg); + mt = pair ? r6_mtg : r6_single_mtg; + nm = pair ? NM_MAX : 1; + } + else if(two_d) { + if constexpr (WIDE) { + if (hex) { + if (out_tiles == 81 && work >= hex_wide) { + const int j = work - hex_wide; + const bool pair = j >= 128; + tile_id = pair ? 80 : 48 + 2 * (j / hex_tok); + tok = HEX_NM * (j % hex_tok); + mt = pair ? 1 : 2; + nm = HEX_NM; + } else if (out_tiles == 40 && work >= hex_wide) { + tile_id = out_tiles - 1; + tok = HEX_NM * (work - hex_wide); + mt = 1; + nm = HEX_NM; + } else { + const int hw = work < hex_wide + ? work : hex_wide + ((work - hex_wide) >> 1); + tile_id = HEX_MT * (hw / hex_tok); + tok = HEX_NM * (hw % hex_tok) + + (work < hex_wide ? 0 : ((work - hex_wide) & 1)); + mt = HEX_MT; nm = work < hex_wide ? HEX_NM : 1; + } + } else if (mixed40) { + const bool quad = work < 132; + const int item = quad ? work : work - 132; + if (quad && item < 84) { + tile_id = 2 * (item / 7); + tok = 2 * (item % 7); + } else if (quad) { + const int tail = item - 84; + tile_id = 2 * (12 + tail / 6); + tok = 2 * (tail % 6); + } else if (item < 48) { + tile_id = 2 * (item / 4) + (item % 4) / 2; + tok = 14 + item % 2; + } else { + const int tail = item - 48; + tile_id = 2 * (12 + tail / 8) + (tail % 8) / 4; + tok = 12 + tail % 4; + } + mt = quad ? MT : 1; + nm = quad ? NM_MAX : 1; + } else if (split72) { + if (work < 264) { + tile_id = 2 * (work / 8); + tok = 2 * (work % 8); + mt = MT; nm = NM_MAX; + } else { + const int i = work - 264, c = i >> 2, q = i & 3; + tile_id = 2 * (33 + (c >> 3)) + (q >> 1); + tok = 2 * (c & 7) + (q & 1); + mt = 1; nm = 1; + } + } else if (split81) { + if (work < 264) { + tile_id = 2 * (work / 8); + tok = 2 * (work % 8); + mt = MT; nm = NM_MAX; + } else if (work < 376) { + const int i = work - 264; + tile_id = 2 * (33 + i / 16); + tok = i % 16; + mt = MT; nm = 1; + } else { + tile_id = 80; + tok = 2 * (work - 376); + mt = 1; nm = NM_MAX; + } + } else if (split17) { + if (work < 8) { + tile_id = 14; + tok = work; + mt = 3; nm = 1; + } else if (work < 120) { + const int i = work - 8; + tile_id = i >> 3; + tok = 2 * (i & 7); + mt = 1; nm = NM_MAX; + } else { + const int i = work - 120; + tile_id = 14 + (i >> 2); + tok = 8 + 2 * (i & 3); + mt = 1; nm = NM_MAX; + } + } else if (tall3) { + tile_id = 3 * (work / DECODE_TILES); + tok = work % DECODE_TILES; + mt = 3; + nm = 1; + } else { + const int j = work < q_units ? work : work - q_units; + tile_id = work < q_units ? MT * (j / QTOK) + : q_rem_base + j / QTOK; + tok = NM_MAX * (j % QTOK); + mt = work < q_units ? MT : 1; + nm = NM_MAX; + } + } else mk2_group_tile(work,grouped,tile_id,tok,nm); + } + int epi_hold = -1; + rt_fl<16, BATCH_PAD> accum[ACCN]; + if (hex && nm == 1) { + #pragma unroll + for (int u = 0; u < 3; ++u) warp::zero(accum[u]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + HEX_MT]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[1], pf[s * pslot + 1], pf[s * pslot + HEX_MT]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[2], pf[s * pslot + 2], pf[s * pslot + HEX_MT]); + warpgroup::mma_async_wait(); + if (i + 1 < iters) { + if (warpgroup::laneid() == 0) + arrive(drained[s], CONSUMER_WARPS); + } else epi_hold = s; + } + } else if (hex && mt == HEX_MT) { + #pragma unroll + for (int u = 0; u < ACCN; ++u) warp::zero(accum[u]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], + pf[s * pslot + HEX_MT]); + #pragma unroll + for (int u = 1; u < ACCN; ++u) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[u], pf[s * pslot + u / HEX_NM], + pf[s * pslot + HEX_MT + u % HEX_NM]); + warpgroup::mma_async_wait(); + if (i + 1 < iters) { + if (warpgroup::laneid() == 0) + arrive(drained[s], CONSUMER_WARPS); + } else epi_hold = s; + } + } else if (WIDE && mt == 3) { + #pragma unroll + for (int u = 0; u < 3; ++u) warp::zero(accum[u]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + 3]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[1], pf[s * pslot + 1], pf[s * pslot + 3]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[2], pf[s * pslot + 2], pf[s * pslot + 3]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else if (WIDE && mt == MT && nm == NM_MAX) { + #pragma unroll + for (int u = 0; u < ACCN; ++u) warp::zero(accum[u]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + MT]); + #pragma unroll + for (int u = 1; u < ACCN; ++u) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[u], pf[s * pslot + u / NM_MAX], + pf[s * pslot + MT + u % NM_MAX]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else if (MT > 1 && mt == MT) { + #pragma unroll + for (int c = 0; c < MT; ++c) warp::zero(accum[c]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + MT]); + #pragma unroll + for (int c = 1; c < MT; ++c) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[c], pf[s * pslot + c], pf[s * pslot + MT]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else if (NM_MAX == 1 || nm == NM_MAX) { + #pragma unroll + for (int m = 0; m < NM_MAX; ++m) warp::zero(accum[m]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + 1]); + #pragma unroll + for (int m = 1; m < NM_MAX; ++m) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[m], pf[s * pslot], pf[s * pslot + 1 + m]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else { + warp::zero(accum[0]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + 1]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } + out_tile_f &e_sf = epi_hold >= 0 + ? *reinterpret_cast(&pf[epi_hold * pslot]) + : o_sf; + out_tile *const e_sp = reinterpret_cast(&e_sf); + const int nacc = mt * nm; + if (epi != EPI_BIAS) { + epi_publish(e_sp[0], accum[0]); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + } + #pragma unroll + for (int u = 0; u < ACCN; ++u) { + if (u >= mt * nm) break; + const int c = (mt > 1) ? ((nm > 1) ? (u / NM_MAX) : u) : 0; + const int m = (nm > 1) ? (u % NM_MAX) : 0; + const int base = (r6_lut ? (tile_id + c)%r6_output_tiles + : (tile_id + c)) * OUT_ATOM; + const int grow0 = (tok + m) * BATCH_PAD; + if (epi == EPI_BIAS) { + const int bw = 16 * warpgroup::warpid() + + (::kittens::laneid() >> 2); + const float b0 = __bfloat162float(bias[base + bw]); + const float b1 = __bfloat162float(bias[base + bw + 8]); + #pragma unroll + for (int j = 0; j < rt_fl<16, BATCH_PAD>::width; ++j) { + accum[u].tiles[0][j].data[0].x += b0; + accum[u].tiles[0][j].data[0].y += b0; + accum[u].tiles[0][j].data[1].x += b1; + accum[u].tiles[0][j].data[1].y += b1; + accum[u].tiles[0][j].data[2].x += b0; + accum[u].tiles[0][j].data[2].y += b0; + accum[u].tiles[0][j].data[3].x += b1; + accum[u].tiles[0][j].data[3].y += b1; + } + epi_publish(e_sp[0], accum[u]); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + constexpr int OUT_OCTS = OUT_ATOM / 8; + for (int idx = warpgroup::laneid(); idx < OUT_OCTS * rows_valid; + idx += 128) { + const int row = idx >> 3, col = (idx & 7) << 3; + const size_t at = (size_t)(grow0 + row) * dest_stride + base + col; + *reinterpret_cast(dest + at) = + *reinterpret_cast( + &e_sp[0][int2{row, col}]); + } + warpgroup::sync(1); + continue; + } + if (u + 1 < ACCN && u + 1 < nacc) + epi_publish(e_sp[(u + 1) & 1], + accum[u + 1 < ACCN ? u + 1 : 0]); + const out_tile &e_sr = e_sp[u & 1]; + constexpr int OUT_OCTS = OUT_ATOM / 8; + for (int idx = warpgroup::laneid(); idx < OUT_OCTS * rows_valid; + idx += 128) { + const int row = idx >> 3; + const int col = (idx & 7) << 3; + const int grow = grow0 + row; + const size_t at = + (size_t)grow * dest_stride + base + col; + const bf16_oct src = *reinterpret_cast( + &e_sr[int2{row, col}]); + if (epi == EPI_STORE) { + *reinterpret_cast(dest + at) = src; + continue; + } + const float v0 = __bfloat162float(src.x0); + const float v1 = __bfloat162float(src.x1); + const float v2 = __bfloat162float(src.x2); + const float v3 = __bfloat162float(src.x3); + const float v4 = __bfloat162float(src.x4); + const float v5 = __bfloat162float(src.x5); + const float v6 = __bfloat162float(src.x6); + const float v7 = __bfloat162float(src.x7); + switch (epi) { + case EPI_STORE: break; + case EPI_SILU: + *reinterpret_cast(dest + at) = bf16_oct{ + __float2bfloat16(v0 / (1.f + __expf(-v0))), + __float2bfloat16(v1 / (1.f + __expf(-v1))), + __float2bfloat16(v2 / (1.f + __expf(-v2))), + __float2bfloat16(v3 / (1.f + __expf(-v3))), + __float2bfloat16(v4 / (1.f + __expf(-v4))), + __float2bfloat16(v5 / (1.f + __expf(-v5))), + __float2bfloat16(v6 / (1.f + __expf(-v6))), + __float2bfloat16(v7 / (1.f + __expf(-v7)))}; + break; + case EPI_MOE_ACC: { + const int expert = (tile_id + c) / out_tiles_per_expert; + const int local = ((tile_id + c) % out_tiles_per_expert) * OUT_ATOM + + col; + const int original_row=r6_map?r6_map[grow]:grow; + if(original_row<0) break; + const float wt = moe_w[original_row * BB_EXPERTS + expert]; + if (wt != 0.f) { + const size_t dst = + ((size_t)original_row * (BB_EXPERTS + 1) + expert) * + D_MODEL + local; + *reinterpret_cast(moe_out + dst) = make_float4( + __fmul_rn(wt, v0), __fmul_rn(wt, v1), + __fmul_rn(wt, v2), __fmul_rn(wt, v3)); + *reinterpret_cast(moe_out + dst + 4) = make_float4( + __fmul_rn(wt, v4), __fmul_rn(wt, v5), + __fmul_rn(wt, v6), __fmul_rn(wt, v7)); + } + break; + } + case EPI_ACC_F32: + *reinterpret_cast(moe_out + + ((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * + D_MODEL + base + col) = make_float4(v0, v1, v2, v3); + *reinterpret_cast(moe_out + + ((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * + D_MODEL + base + col + 4) = make_float4(v4, v5, v6, v7); + break; + case EPI_BIAS: break; + case EPI_GELU: + *reinterpret_cast(dest + at) = bf16_oct{ + __float2bfloat16(gelu_tanh_fp32(v0)), + __float2bfloat16(gelu_tanh_fp32(v1)), + __float2bfloat16(gelu_tanh_fp32(v2)), + __float2bfloat16(gelu_tanh_fp32(v3)), + __float2bfloat16(gelu_tanh_fp32(v4)), + __float2bfloat16(gelu_tanh_fp32(v5)), + __float2bfloat16(gelu_tanh_fp32(v6)), + __float2bfloat16(gelu_tanh_fp32(v7))}; + break; + case EPI_RESIDUAL: { + const bf16_oct prior = + *reinterpret_cast(residual + at); + *reinterpret_cast(residual + at) = bf16_oct{ + __float2bfloat16(__bfloat162float(prior.x0) + v0), + __float2bfloat16(__bfloat162float(prior.x1) + v1), + __float2bfloat16(__bfloat162float(prior.x2) + v2), + __float2bfloat16(__bfloat162float(prior.x3) + v3), + __float2bfloat16(__bfloat162float(prior.x4) + v4), + __float2bfloat16(__bfloat162float(prior.x5) + v5), + __float2bfloat16(__bfloat162float(prior.x6) + v6), + __float2bfloat16(__bfloat162float(prior.x7) + v7)}; + break; + } + } + } + if (u + 1 < nacc) { + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + } + } + if (epi != EPI_BIAS) warpgroup::sync(1); + if (epi_hold >= 0 && warpgroup::laneid() == 0) + arrive(drained[epi_hold], CONSUMER_WARPS); + } + (void)my_tiles; + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr, + const unsigned char *native_predictor_active_mask = nullptr + MK2_PROGRESS_PARAMETER(native_progress_mask)) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&p_s)[STAGES][SLOTS] = alloc.allocate(); + out_tile_f (&o_sf) = alloc.allocate(); + out_tile &o_s = *reinterpret_cast(&o_sf); + bf16 *const route_stage = reinterpret_cast(&o_sf); + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_SCRATCH_STRIDE]; + static_assert(STAGES * SLOTS * sizeof(weight_tile) >= sizeof(sampler_storage)); + sampler_storage &sort_storage = *reinterpret_cast(&__shm[0]); + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS, 0); + arrive(drained[s], CONSUMER_WARPS); + } + } + int gemm_iter = 0; + int gemm_tile = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + const float *fold_src = nullptr, *fold_routes = nullptr; + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + fold_src = nullptr; + fold_routes = nullptr; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + grid_barrier(g.barrier); +#endif + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); +#endif + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, BATCH, + route_stage); + } + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + } + grid_barrier(g.barrier); + r6_sparse_map(bb->moe_w,bb->r6_map,bb->r6_counts + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + r6_sparse_gather(bb->normed,bb->r6_packed,bb->r6_map,bb->r6_counts, + bb->r6_lut); + grid_barrier(g.barrier); + gemm(L.exp_up, bb->r6_views[0], BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0,0,nullptr,nullptr,DECODE_TILES,DECODE_ROWS, + nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->r6_views[1], BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + DECODE_TILES, DECODE_ROWS, nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + if (bb->layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_selected(bb->moe_out, bb->moe_w, idx); + grid_barrier(g.barrier); + } else { + fold_src = bb->moe_out; + fold_routes = bb->moe_w; + } + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int n_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = n_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + bb->pf_h[idx] = bb->prompt[idx]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = bb->pf_n_seq > 0 ? bb->pf_n_seq : 1; + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_gated, + T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, T, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * (BB_EXPERTS + 1) * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->pf_normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->pf_moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = bb_moe_fold_deterministic(bb->pf_moe_out, idx); + grid_barrier(g.barrier); + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, bb->pf_final, ln_scratch, + n_total); + grid_barrier(g.barrier); + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, bb->phoneme_b); + grid_barrier(g.barrier); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + for (int row = blockIdx.x + warp * gridDim.x; row < BATCH; + row += gridDim.x * WARPS) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + MK2_PROGRESS_GUARD(native_progress_mask, row); + const int arg = phoneme_warp_argmax(lg); + if (lane == 0) bb->phoneme_out[row] = arg; + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + __shared__ int native_run_predictor; + if (native_predictor_active_mask && bb && g.native_codes) { + if (threadIdx.x == 0) { + int any_active = 0; + for (int row = 0; row < BATCH; ++row) + any_active |= int(native_predictor_active_mask[row] != 0 && + MK2_PROGRESS_ENABLED(native_progress_mask, row)); + native_run_predictor = any_active; + } + __syncthreads(); + } + const bool skip_native_predictor = native_predictor_active_mask && bb && + g.native_codes && native_run_predictor == 0; + if (skip_native_predictor) { + for (int index = blockIdx.x * THREADS + threadIdx.x; + index < BATCH * N_POS; index += gridDim.x * THREADS) { + MK2_PROGRESS_GUARD(native_progress_mask, index / N_POS); + g.sampled[index] = 0; + g.native_codes[index] = 0; + } + grid_barrier(g.barrier); + } else { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + constexpr int DPL = D_HEAD / 32; + float qv[DPL], kn[DPL], vn[DPL]; + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + qv[j] = __bfloat162float(q[i]); + kn[j] = __bfloat162float(k_new[i]); + vn[j] = __bfloat162float(v_new[i]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += qv[j] * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + { + float acc = 0.f; + #pragma unroll + for (int j = 0; j < DPL; ++j) acc += qv[j] * kn[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[pos] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + float ov[DPL] = {}; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *vt = g.vcache + cache_base + t * D_HEAD; + const float s = sc[t]; + #pragma unroll + for (int j = 0; j < DPL; ++j) + ov[j] += s * __bfloat162float(vt[lane + j * 32]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + ov[j] += sc[pos] * vn[j]; + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + + lane + j * 32] = __float2bfloat16(ov[j]); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, + g.sampled_pos, sort_storage, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif +#ifdef ABLATE_SAMPLE + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { + __syncthreads(); + } else { + grid_barrier(g.barrier); + } +#endif + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x == 0) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + g.sampled[row * N_POS + pos] = code; + if (g.native_codes && MK2_PROGRESS_ENABLED(native_progress_mask, row)) { +#if MK2_INFLIGHT + g.native_codes[row * N_POS + pos] = + native_predictor_active_mask && !native_predictor_active_mask[row] + ? 0 : (long long)code; +#else + g.native_codes[row * N_POS + pos] = (long long)code; +#endif + } + } + } +#ifdef ABLATE_SAMPLE + grid_barrier(g.barrier); +#else + if (pos + 1 == n_pos) grid_barrier(g.barrier); +#endif + } + } + if (g.native_seq_len && blockIdx.x == 0) + for (int row = threadIdx.x; row < BATCH; row += blockDim.x) { + MK2_PROGRESS_GUARD(native_progress_mask, row); + g.native_seq_len[row] += 1; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2_startup( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr, + const unsigned char *native_predictor_active_mask = nullptr + MK2_PROGRESS_PARAMETER(native_progress_mask)) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&p_s)[STAGES][SLOTS] = alloc.allocate(); + out_tile_f (&o_sf) = alloc.allocate(); + out_tile &o_s = *reinterpret_cast(&o_sf); + bf16 *const route_stage = reinterpret_cast(&o_sf); + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_SCRATCH_STRIDE]; + static_assert(STAGES * SLOTS * sizeof(weight_tile) >= sizeof(sampler_storage)); + sampler_storage &sort_storage = *reinterpret_cast(&__shm[0]); + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS, 0); + arrive(drained[s], CONSUMER_WARPS); + } + } + int gemm_iter = 0; + int gemm_tile = 0; + __syncthreads(); + constexpr bool prefilling = false; + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + const float *fold_src = nullptr, *fold_routes = nullptr; + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + fold_src = nullptr; + fold_routes = nullptr; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + grid_barrier(g.barrier); +#endif + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); +#endif + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, BATCH, + route_stage); + } + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + } + grid_barrier(g.barrier); + r6_sparse_map(bb->moe_w,bb->r6_map,bb->r6_counts + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + r6_sparse_gather(bb->normed,bb->r6_packed,bb->r6_map,bb->r6_counts, + bb->r6_lut); + grid_barrier(g.barrier); + gemm(L.exp_up, bb->r6_views[0], BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0,0,nullptr,nullptr,DECODE_TILES,DECODE_ROWS, + nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->r6_views[1], BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + DECODE_TILES, DECODE_ROWS, nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + if (bb->layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_selected(bb->moe_out, bb->moe_w, idx); + grid_barrier(g.barrier); + } else { + fold_src = bb->moe_out; + fold_routes = bb->moe_w; + } + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if constexpr (false) { + const int n_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = n_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + bb->pf_h[idx] = bb->prompt[idx]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = bb->pf_n_seq > 0 ? bb->pf_n_seq : 1; + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_gated, + T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, T, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * (BB_EXPERTS + 1) * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->pf_normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->pf_moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = bb_moe_fold_deterministic(bb->pf_moe_out, idx); + grid_barrier(g.barrier); + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, bb->pf_final, ln_scratch, + n_total); + grid_barrier(g.barrier); + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, bb->phoneme_b); + grid_barrier(g.barrier); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + for (int row = blockIdx.x + warp * gridDim.x; row < BATCH; + row += gridDim.x * WARPS) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + MK2_PROGRESS_GUARD(native_progress_mask, row); + const int arg = phoneme_warp_argmax(lg); + if (lane == 0) bb->phoneme_out[row] = arg; + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + __shared__ int native_run_predictor; + if (native_predictor_active_mask && bb && g.native_codes) { + if (threadIdx.x == 0) { + int any_active = 0; + for (int row = 0; row < BATCH; ++row) + any_active |= int(native_predictor_active_mask[row] != 0 && + MK2_PROGRESS_ENABLED(native_progress_mask, row)); + native_run_predictor = any_active; + } + __syncthreads(); + } + const bool skip_native_predictor = native_predictor_active_mask && bb && + g.native_codes && native_run_predictor == 0; + if constexpr (true) { + for (int index = blockIdx.x * THREADS + threadIdx.x; + index < BATCH * N_POS; index += gridDim.x * THREADS) { + MK2_PROGRESS_GUARD(native_progress_mask, index / N_POS); + g.sampled[index] = 0; + g.native_codes[index] = 0; + } + grid_barrier(g.barrier); + } else { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + constexpr int DPL = D_HEAD / 32; + float qv[DPL], kn[DPL], vn[DPL]; + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + qv[j] = __bfloat162float(q[i]); + kn[j] = __bfloat162float(k_new[i]); + vn[j] = __bfloat162float(v_new[i]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += qv[j] * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + { + float acc = 0.f; + #pragma unroll + for (int j = 0; j < DPL; ++j) acc += qv[j] * kn[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[pos] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + float ov[DPL] = {}; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *vt = g.vcache + cache_base + t * D_HEAD; + const float s = sc[t]; + #pragma unroll + for (int j = 0; j < DPL; ++j) + ov[j] += s * __bfloat162float(vt[lane + j * 32]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + ov[j] += sc[pos] * vn[j]; + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + + lane + j * 32] = __float2bfloat16(ov[j]); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, + g.sampled_pos, sort_storage, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif +#ifdef ABLATE_SAMPLE + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { + __syncthreads(); + } else { + grid_barrier(g.barrier); + } +#endif + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x == 0) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + g.sampled[row * N_POS + pos] = code; + if (g.native_codes && MK2_PROGRESS_ENABLED(native_progress_mask, row)) { +#if MK2_INFLIGHT + g.native_codes[row * N_POS + pos] = + native_predictor_active_mask && !native_predictor_active_mask[row] + ? 0 : (long long)code; +#else + g.native_codes[row * N_POS + pos] = (long long)code; +#endif + } + } + } +#ifdef ABLATE_SAMPLE + grid_barrier(g.barrier); +#else + if (pos + 1 == n_pos) grid_barrier(g.barrier); +#endif + } + } + if (g.native_seq_len && blockIdx.x == 0) + for (int row = threadIdx.x; row < BATCH; row += blockDim.x) { + MK2_PROGRESS_GUARD(native_progress_mask, row); + g.native_seq_len[row] += 1; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 64 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * SLOTS * sizeof(weight_tile) + + sizeof(out_tile_f) + 2048; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; + const unsigned char *native_predictor_active_mask = nullptr; + MK2_PROGRESS_MEMBER + int maximum_grid = 0; + bool grid_reset_pending = false; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * SLOTS * sizeof(weight_tile) + + sizeof(out_tile_f) + 2048), 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaFuncSetAttribute(predictor_mk2_startup, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = std::min(out->sms * out->blocks_per_sm, 4096); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 64 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->maximum_grid = info.sms * info.blocks_per_sm; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + mk2_check(owned.alloc(&hb.r6_packed,(size_t)BB_EXPERTS*BATCH*D_MODEL*sizeof(bf16))); + mk2_check(owned.alloc(&hb.r6_map,(size_t)BB_EXPERTS*BATCH*sizeof(int))); + mk2_check(owned.alloc(&hb.r6_counts,BB_EXPERTS*sizeof(int))); + mk2_check(owned.alloc(&hb.r6_lut,(1+BB_EXPERTS*DECODE_TILES)*sizeof(int))); + act_gl sparse_views[]={ + act_gl{hb.r6_packed,nullptr,nullptr,(size_t)BB_EXPERTS*BATCH,(size_t)D_MODEL}, + act_gl{moe_h,nullptr,nullptr,(size_t)BB_EXPERTS*BATCH,(size_t)BB_MOE_INTER}}; + mk2_check(owned.alloc(&hb.r6_views,sizeof(sparse_views))); + mk2_check(cudaMemcpy(hb.r6_views,sparse_views,sizeof(sparse_views),cudaMemcpyHostToDevice)); + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.maximum_grid = pred->maximum_grid; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native, bool startup = false) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(s, &capture_status)); + if (st->pred.grid_reset_pending || capture_status != cudaStreamCaptureStatusNone) { + mk2_check(cudaMemsetAsync(st->pred.g.barrier, 0, 64 * sizeof(unsigned int), s)); + st->pred.grid_reset_pending = false; + } + if (startup) { + predictor_mk2_startup<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS, nullptr, + st->pred.native_predictor_active_mask + MK2_PROGRESS_ARGUMENT(st->pred.native_progress_mask)); + } else { + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS, nullptr, + native ? st->pred.native_predictor_active_mask : nullptr + MK2_PROGRESS_ARGUMENT(native ? st->pred.native_progress_mask : nullptr)); + } + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" int mk2_run_native_startup_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream) { + return mk2_enqueue_stage0(p, b, stream, true, true); +} +extern "C" int mk2_set_native_predictor_active_mask(MK2Predictor *p, + const unsigned char *mask) { + try { + if (!p || reinterpret_cast(p)->kind != 1) + throw std::runtime_error("startup mask requires a full Stage0 handle"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && st->pred.native_predictor_active_mask != mask) + throw std::runtime_error("startup mask pointer cannot change after first Stage0 enqueue"); + if (mask) + mk2_stage0_device_pointer(mask, st->pred.device, "native_predictor_active_mask"); + st->pred.native_predictor_active_mask = mask; + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown startup mask binding failure"; return 1; } +} +extern "C" int mk2_native_startup_skip_abi_version() { return 1; } +#include "../../../../../inflight/runtime.cuh" +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "native_prefill_impl.inc" +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *p, int grid) { + if (!p) { g_last_error = "grid setter requires full Stage0 handle"; return 1; } + auto *st = reinterpret_cast(p); + if (st->kind != 1 || grid < 1 || grid > st->maximum_grid) { + g_last_error = "requested Stage0 grid outside independently resident capacity"; + return 1; + } + if (st->grid != grid) st->grid_reset_pending = true; + st->grid = grid; + g_last_error.clear(); + return 0; +} +extern "C" int mk2_r8_bind_initial_state(MK2Predictor *p, const void *const *initial, const void *ready, const void *seq_len, int positions) { + try { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || st->bindings_initialized || !initial || !ready || !seq_len || positions <= 0) + throw std::runtime_error("invalid immutable initial-state binding"); + for (int i = 0; i < 14; ++i) { + mk2_stage0_device_pointer(initial[i], st->pred.device, "initial state"); + st->mb_host[i].r8_initial_ssm = static_cast(initial[i]); + st->mb_host[i].r8_initial_ready = static_cast(ready); + st->mb_host[i].r8_seq_len = static_cast(seq_len); + st->mb_host[i].r8_initial_positions = positions; + } + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" int mk2_r9_state_bind(void* p,const void*const* scales,const void*const* maps) {auto*st=(MK2Stage0State*)p;if(st->bindings_initialized)return 1;for(int i=0;i<14;++i){st->mb_host[i].r9_scales=(const float*)scales[i];st->mb_host[i].r9_map=(const int*)maps[i];}return 0;} +extern "C" int mk2_r9_state_bytes(){return sizeof(r9_state_t);} diff --git a/deploy/nemotron_native/kernels/b2048/regions/stage0/stage0.cu b/deploy/nemotron_native/kernels/b2048/regions/stage0/stage0.cu new file mode 100644 index 0000000..ec50bd0 --- /dev/null +++ b/deploy/nemotron_native/kernels/b2048/regions/stage0/stage0.cu @@ -0,0 +1,35 @@ +#define MK2_BATCH 2048 +#define STAGES 3 +#define KCHUNK 128 +#define MK2_NO_MAIN 1 +#define MK2_DIRECT_INPUT 1 +#define MK2_B1_LOCAL_BARRIERS 1 +#define MK2_SKIP_HIDDEN_SNAPSHOT 1 +#define MK2_ACTIVE_EXPERTS 1 +#define MK2_SSM_SPLITS 1 +#define NMULT 2 +#define MK2_OUT_MULT 2 +#define SSM_PREFETCH_MAX 12 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#if MK2_INFLIGHT +#include +#endif +#include +#include +#include +#include +#include +#include +#include +#include +namespace b2048_stage0 { +#include "native/stage0/predictor_mk2.cu" +} diff --git a/deploy/nemotron_native/kernels/b2048/stage0.cu b/deploy/nemotron_native/kernels/b2048/stage0.cu new file mode 100644 index 0000000..6c04cd7 --- /dev/null +++ b/deploy/nemotron_native/kernels/b2048/stage0.cu @@ -0,0 +1 @@ +#include "regions/stage0/stage0.cu" diff --git a/deploy/nemotron_native/kernels/b256/prefill.cu b/deploy/nemotron_native/kernels/b256/prefill.cu new file mode 100644 index 0000000..6d21d7e --- /dev/null +++ b/deploy/nemotron_native/kernels/b256/prefill.cu @@ -0,0 +1 @@ +#include "regions/prefill/prefill.cu" diff --git a/deploy/nemotron_native/kernels/b256/regions/prefill/native/prefill/dedup_native_prefill_impl.inc b/deploy/nemotron_native/kernels/b256/regions/prefill/native/prefill/dedup_native_prefill_impl.inc new file mode 100644 index 0000000..fe4c125 --- /dev/null +++ b/deploy/nemotron_native/kernels/b256/regions/prefill/native/prefill/dedup_native_prefill_impl.inc @@ -0,0 +1,267 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 16 || BATCH == 32 || BATCH == 64 || BATCH == 128 || BATCH == 256 || BATCH == 512 || BATCH == 1024 || BATCH == 2048)) + throw std::runtime_error("native prefill requires a matching supported physical batch"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((std::uint64_t(c.prompt_positions) + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (rows > (INT_MAX - 132 * THREADS) / D_MODEL) + throw std::runtime_error("native prefill full-width prompt exceeds safe device loop bounds"); + if (padded > (INT_MAX - 132 * THREADS) / (BB_EXPERTS * BB_MOE_INTER)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((size_t(config->prompt_positions) + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + float *r9_workspace = nullptr; + mk2_check(owned.alloc(&r9_workspace,size_t(BATCH)*ssm_elements*sizeof(float))); + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + r9_state_t *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(r9_state_t))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(r9_state_t))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = r9_workspace; + st->mb_host[i].r9_private = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + mk2_check(owned.alloc(&st->mb_host[i].r8_canonical_ssm, ssm_elements * sizeof(float))); + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_dtda, 2 * MB_HEADS); + { + const size_t estride = padded; + const size_t grows = size_t(BB_EXPERTS) * estride; + mk2_check(owned.alloc(&hb.pf_moe_gath, grows * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_gath, 0, grows * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_hc, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_hc, 0, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_gidx, grows * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_gidx, 0, grows * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ecnt, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ecnt, 0, BB_EXPERTS * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ntiles, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ntiles, 0, BB_EXPERTS * sizeof(int))); + hb.pf_moe_estride = int(estride); + hb.pf_moe_gath_v = act_gl{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_v = act_gl{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + hb.pf_moe_gath_vn = act_gl_n{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_vn = act_gl_n{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + scratch_bytes += grows * (D_MODEL + BB_MOE_INTER) * sizeof(bf16); + } + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(owned.alloc(&hb.native_pf_cohort_check, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(hb.native_pf_cohort_check, 0, 2 * sizeof(unsigned int))); +#if PF_ACTIVE_TOKEN_EXPERTS + mk2_check(owned.alloc(&hb.native_pf_expert_masks, (padded/BATCH_PAD)*sizeof(unsigned int))); +#endif + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} +extern "C" int mk2_native_prefill_validate(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null compact prefill handle"); + auto *st = reinterpret_cast(p->base); + mk2_check(cudaStreamSynchronize(static_cast(stream))); + unsigned int status[2] = {0, 0}; + mk2_check(cudaMemcpy(status, st->host_bb.native_pf_cohort_check, + sizeof(status), cudaMemcpyDeviceToHost)); + if (status[0] != 0 || status[1] != 1) + throw std::runtime_error("compact prefill requires bit-identical prompts and identity slots"); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown compact prefill validation failure"; return 1; } +} +extern "C" int mk2_native_prefill_workspace_sequences() { return 1; } +extern "C" int mk2_native_prefill_effective_sequences(MK2NativePrefill *p, int *out) { + try { + if (!p || !out) throw std::runtime_error("null effective-sequence query"); + auto *st = reinterpret_cast(p->base); + mk2_check(cudaMemcpy(out, st->host_bb.native_pf_cohort_check + 1, + sizeof(int), cudaMemcpyDeviceToHost)); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" int mk2_r8_initial_state(MK2NativePrefill *p, void **initial, void **ready) { + if (!p || !initial || !ready) return 1; + auto *st = reinterpret_cast(p->base); + for (int i = 0; i < 14; ++i) initial[i] = st->mb_host[i].r8_canonical_ssm; + *ready = st->host_bb.native_pf_cohort_check + 1; + return 0; +} +extern "C" int mk2_r9_prefill_bind(MK2NativePrefill*p,const void*const*scales,const void*const*maps){if(!p)return 1;auto*st=(MK2Stage0State*)p->base;for(int i=0;i<14;++i){st->mb_host[i].r9_scales=(const float*)scales[i];st->mb_host[i].r9_map=(const int*)maps[i];}return (int)cudaMemcpy((void*)st->host_bb.mamba,st->mb_host.data(),14*sizeof(MambaLayer),cudaMemcpyHostToDevice);} diff --git a/deploy/nemotron_native/kernels/b256/regions/prefill/native/prefill/dedup_predictor_mk2.cuh b/deploy/nemotron_native/kernels/b256/regions/prefill/native/prefill/dedup_predictor_mk2.cuh new file mode 100644 index 0000000..f98dd53 --- /dev/null +++ b/deploy/nemotron_native/kernels/b256/regions/prefill/native/prefill/dedup_predictor_mk2.cuh @@ -0,0 +1,4600 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef PF_ACTIVE_TOKEN_EXPERTS +#define PF_ACTIVE_TOKEN_EXPERTS 0 +#endif +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +#ifndef MK2_PREFILL_ONLY +#define MK2_PREFILL_ONLY 0 +#endif +#ifndef MK2_BATCH_TILE +#define MK2_BATCH_TILE (((MK2_BATCH + 15) / 16) * 16) +#endif +constexpr int BATCH_PAD = MK2_BATCH_TILE; +static_assert(BATCH >= 1 && (MK2_PREFILL_ONLY || BATCH <= BATCH_PAD)); +static_assert(BATCH_PAD == 16 || BATCH_PAD == 32 || BATCH_PAD == 64 || BATCH_PAD == 128, + "native prefill tile must be 16, 32, 64 or 128"); +static_assert(MK2_PREFILL_ONLY || BATCH_PAD == ((BATCH + 15) / 16) * 16, + "a wider token tile requires a prefill-only library"); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_PW_MULTI +#define PF_SSM_PW_MULTI 4 +#endif +#ifndef PF_SSM_PW_SINGLE +#define PF_SSM_PW_SINGLE 2 +#endif +static_assert(MB_HEAD_DIM % PF_SSM_PW_MULTI == 0 && + MB_HEAD_DIM % PF_SSM_PW_SINGLE == 0, + "head_dim must split evenly into scan channel groups"); +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef MK2_SMEM_SLACK +#define MK2_SMEM_SLACK 1024 +#endif +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +#ifndef MK2_THREADS +#define MK2_THREADS 256 +#endif +constexpr int THREADS = MK2_THREADS; +constexpr int NORM_THREADS = 256; +constexpr int NORM_WARPS = NORM_THREADS / 32; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +#ifndef MK2_GEMM_CONSUMERS +#define MK2_GEMM_CONSUMERS 1 +#endif +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 +#define MK2_HEAD_STAGE_BYTES 0 +#else +#define MK2_HEAD_STAGE_BYTES sizeof(out_tile_f) +#endif +constexpr int NCONS = MK2_GEMM_CONSUMERS; +static_assert(NCONS >= 1 && NCONS <= 3, "consumers occupy warpgroups 0, 2 and 3"); +static_assert(THREADS >= (NCONS + 1) * 128, "one producer plus NCONS consumers"); +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +constexpr int PAD_N = 32; +using act_tile_n = st_bf; +using out_tile_n = st_bf; +using act_gl_n = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_MOE_ACC_C, EPI_ACC_F32 }; +#include +constexpr int R9_STATE_MODE=1; +using r9_state_t=bf16; +__device__ __forceinline__ int r9_scale_index(int head,int channel,int n) {return (((head/8)*4+(head%8)/2)*4+channel/6)*128+n;} +__device__ __forceinline__ float r9_load(r9_state_t v,float scale) { +return __bfloat162float(v); +} +__device__ __forceinline__ r9_state_t r9_store(float x,float scale) { +return __float2bfloat16(x); +} +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; + float *r8_canonical_ssm = nullptr; + r9_state_t *r9_private = nullptr; + const float *r9_scales = nullptr; + const int *r9_map = nullptr; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + unsigned int *native_pf_expert_masks = nullptr; + float *pf_dtda = nullptr; + bf16 *pf_moe_gath = nullptr; + bf16 *pf_moe_hc = nullptr; + int *pf_moe_gidx = nullptr; + int *pf_moe_ecnt = nullptr; + int *pf_moe_ntiles = nullptr; + int pf_moe_estride = 0; + act_gl pf_moe_gath_v; + act_gl pf_moe_hc_v; + act_gl_n pf_moe_gath_vn; + act_gl_n pf_moe_hc_vn; + unsigned int *native_pf_cohort_check = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + volatile unsigned int *generation = state + 1; + const unsigned int seen = *generation; + if (atomicAdd(state, 1u) == gridDim.x - 1) { + atomicExch(state, 0u); + atomicAdd(state + 1, 1u); + } else { + while (*generation == seen) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const bf16 *in = src + (size_t)row * D_MODEL; + float v[PER_THREAD]; + float sum = 0.f, sq = 0.f; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (threadIdx.x < NORM_THREADS) { + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + v[j] = __bfloat162float(in[threadIdx.x + j * NORM_THREADS]); + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + if (lane == 0) { scratch[warp] = sum; scratch[NORM_WARPS + warp] = sq; } + } + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f, q = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) { s += scratch[w]; q += scratch[NORM_WARPS + w]; } + scratch[2 * NORM_WARPS] = s / D_MODEL; + scratch[2 * NORM_WARPS + 1] = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + } + __syncthreads(); + const float mean = scratch[2 * NORM_WARPS], inv = scratch[2 * NORM_WARPS + 1]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + if (threadIdx.x >= NORM_THREADS) break; + const int i = threadIdx.x + j * NORM_THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * weight[i]); + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, + int idx, unsigned mask); +__device__ __forceinline__ void norm_sub_sync(int sub) { + if (THREADS == NORM_THREADS) { __syncthreads(); return; } + asm volatile("bar.sync %0, %1;" :: "r"(6 + sub), "n"(NORM_THREADS) : "memory"); +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, const float *fold_src = nullptr, + const float *fold_routes = nullptr) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + constexpr int NSUB = THREADS / NORM_THREADS; + if (threadIdx.x >= NSUB * NORM_THREADS) return; + const int sub = threadIdx.x / NORM_THREADS; + const int tid = threadIdx.x - sub * NORM_THREADS; + float *const sc = scratch + sub * (NORM_WARPS + 1); + for (int row = blockIdx.x * NSUB + sub; row < n_rows; + row += gridDim.x * NSUB) { + const size_t base = (size_t)row * D_MODEL; + unsigned fold_mask = (1u << BB_EXPERTS) - 1u; + if (fold_src && fold_routes) { + fold_mask = 0u; + const float *rw = fold_routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) + if (rw[e] != 0.f) fold_mask |= 1u << e; + } + float v[PER_THREAD]; + float sq = 0.f; + const int warp = tid >> 5, lane = tid & 31; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + bf16 hv; + if (fold_src) { + hv = bb_moe_fold_deterministic(fold_src, (int)(base + i), fold_mask); + h[base + i] = hv; + } else { + hv = h[base + i]; + } + v[j] = __bfloat162float(hv) + __bfloat162float(r[base + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + if (lane == 0) sc[warp] = sq; + norm_sub_sync(sub); + if (tid == 0) { + float s = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) s += sc[w]; + sc[NORM_WARPS] = rsqrtf(s / D_MODEL + RMS_EPS); + } + norm_sub_sync(sub); + const float inv = sc[NORM_WARPS]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + out[base + i] = __float2bfloat16( + __bfloat162float(weight[i]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + } + norm_sub_sync(sub); + } +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + const int c = idx - row * MB_CONV_DIM; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + const float raw = __bfloat162float( + zxbcdt[(size_t)row * MB_IN_PROJ + MB_INNER + c]); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + xbc_act[(size_t)row * MB_CONV_DIM + c] = __float2bfloat16(acc); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + const int pairs = BATCH * MB_HEADS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_HEADS * split); + const int head = (pair / split) % MB_HEADS; + const int first_channel = (pair % split) * channels; + const int group = head / MB_HEADS_PER_GROUP; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + L.dt_bias[head]); + const float dA = __expf(L.A[head] * dt); + const float Dh = L.D[head]; + float *st = L.ssm_state + (size_t)slot_map[row] * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + for (int p = first_channel; p < first_channel + channels; ++p) { + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + const size_t at = (size_t)p * MB_STATE + n; + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + static_assert(MB_CONV_K == 4, "packed conv taps assume a 4-tap kernel"); + constexpr int CPT = 4; + static_assert(MB_CONV_DIM % CPT == 0, "channel vector width"); + constexpr int CGROUPS = MB_CONV_DIM / CPT; + const int total = n_tokens * CGROUPS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / CGROUPS; + const int c = (idx - t * CGROUPS) * CPT; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + float wv[CPT][MB_CONV_K]; + { + const float4 *wsrc = reinterpret_cast( + L.conv_w + (size_t)c * MB_CONV_K); + #pragma unroll + for (int j = 0; j < CPT / 2; ++j) { + const float4 w4 = wsrc[j]; + const unsigned uu0 = __float_as_uint(w4.x); + const unsigned uu1 = __float_as_uint(w4.y); + const unsigned uu2 = __float_as_uint(w4.z); + const unsigned uu3 = __float_as_uint(w4.w); + wv[2 * j][0] = __uint_as_float(uu0 << 16); + wv[2 * j][1] = __uint_as_float(uu0 & 0xffff0000u); + wv[2 * j][2] = __uint_as_float(uu1 << 16); + wv[2 * j][3] = __uint_as_float(uu1 & 0xffff0000u); + wv[2 * j + 1][0] = __uint_as_float(uu2 << 16); + wv[2 * j + 1][1] = __uint_as_float(uu2 & 0xffff0000u); + wv[2 * j + 1][2] = __uint_as_float(uu3 << 16); + wv[2 * j + 1][3] = __uint_as_float(uu3 & 0xffff0000u); + } + } + float acc[CPT]; + { + const uint2 b2 = *reinterpret_cast(L.conv_b + c); + acc[0] = __uint_as_float(b2.x << 16); + acc[1] = __uint_as_float(b2.x & 0xffff0000u); + acc[2] = __uint_as_float(b2.y << 16); + acc[3] = __uint_as_float(b2.y & 0xffff0000u); + } + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v[CPT]; + if (src < 0) { + #pragma unroll + for (int q = 0; q < CPT; ++q) + v[q] = __bfloat162float( + state[(size_t)q * (MB_CONV_K - 1) + pos + k]); + } else { + const uint2 p = *reinterpret_cast( + zxbcdt + (size_t)src * MB_IN_PROJ + MB_INNER + c); + v[0] = __uint_as_float(p.x << 16); + v[1] = __uint_as_float(p.x & 0xffff0000u); + v[2] = __uint_as_float(p.y << 16); + v[3] = __uint_as_float(p.y & 0xffff0000u); + } + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] += wv[q][k] * v[q]; + } + uint2 out; + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] = acc[q] / (1.f + __expf(-acc[q])); + out.x = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[0])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[1])) << 16); + out.y = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[2])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[3])) << 16); + *reinterpret_cast(xbc_act + (size_t)t * MB_CONV_DIM + c) = out; + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_dtda_prefill( + const MambaLayer &L, const bf16 *zxbcdt, float *dtda, int n_tokens) { + const int total = n_tokens * MB_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_HEADS; + const int head = idx - t * MB_HEADS; + const float dt = mamba_softplus( + __bfloat162float(zxbcdt[(size_t)t * MB_IN_PROJ + MB_INNER + + MB_CONV_DIM + head]) + L.dt_bias[head]); + dtda[(size_t)idx * 2] = dt; + dtda[(size_t)idx * 2 + 1] = __expf(L.A[head] * dt); + } +} +template +__device__ __forceinline__ void mamba_ssm_prefill_impl( + const MambaLayer &L, const bf16 *xbc_act, const float *dtda, bf16 *gated, + int n_pos, int n_seq, int slot, const int *seq_slot) { + constexpr int NS = MB_STATE / 32; + constexpr int PG = MB_HEAD_DIM / PW; + constexpr bool RS = (PW % 4) == 0; + constexpr int H1 = RS ? PW / 2 : 1; + constexpr int H2 = RS ? PW / 4 : 1; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const bool keep_hi = (lane & 16) != 0; + const bool keep_hi2 = (lane & 8) != 0; + const int warps_total = gridDim.x * WARPS; + const int units = n_seq * MB_HEADS * PG; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int pg = unit % PG; + const int rest = unit / PG; + const int head = rest % MB_HEADS; + const int seq = rest / MB_HEADS; + const int p0 = pg * PW; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Dh = L.D[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p0 * MB_STATE; + float sreg[PW][NS]; + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + sreg[q][j] = st[(size_t)q * MB_STATE + lane + j * 32]; + } + auto fetch = [&](int p, float (&oB)[NS], float (&oC)[NS], + float (&oX)[PW], float2 &od) { + const size_t tt = (size_t)p * n_seq + seq; + const bf16 *xa2 = xbc_act + tt * MB_CONV_DIM; + const bf16 *Bs = xa2 + MB_INNER + group * MB_STATE; + const bf16 *Cs = xa2 + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + od = *reinterpret_cast( + dtda + (tt * MB_HEADS + head) * 2); + #pragma unroll + for (int j = 0; j < NS; ++j) { + oB[j] = __bfloat162float(Bs[lane + j * 32]); + oC[j] = __bfloat162float(Cs[lane + j * 32]); + } + if constexpr (RS) { + const uint2 *xp = reinterpret_cast( + xa2 + head * MB_HEAD_DIM + p0); + #pragma unroll + for (int q = 0; q < PW; q += 4) { + const uint2 w = xp[q >> 2]; + oX[q] = __uint_as_float(w.x << 16); + oX[q + 1] = __uint_as_float(w.x & 0xffff0000u); + oX[q + 2] = __uint_as_float(w.y << 16); + oX[q + 3] = __uint_as_float(w.y & 0xffff0000u); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) + oX[q] = __bfloat162float(xa2[head * MB_HEAD_DIM + p0 + q]); + } + }; + float pB[NS], pC[NS], pX[PW]; + float2 pd; + fetch(0, pB, pC, pX, pd); + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const float dt = pd.x; + const float dA = pd.y; + float Bv[NS], Cv[NS]; + #pragma unroll + for (int j = 0; j < NS; ++j) { + Bv[j] = pB[j] * dt; + Cv[j] = pC[j]; + } + float acc[PW], xvs[PW]; + #pragma unroll + for (int q = 0; q < PW; ++q) xvs[q] = pX[q]; + if (pos + 1 < n_pos) fetch(pos + 1, pB, pC, pX, pd); + #pragma unroll + for (int q = 0; q < PW; ++q) { + const float xv = xvs[q]; + float a = 0.f; + #pragma unroll + for (int j = 0; j < NS; ++j) { + sreg[q][j] = sreg[q][j] * dA + Bv[j] * xv; + a += sreg[q][j] * Cv[j]; + } + acc[q] = a; + } + if constexpr (RS) { + float m1[H1], t1[H1], x1[H1]; + #pragma unroll + for (int q = 0; q < H1; ++q) { + m1[q] = keep_hi ? acc[q + H1] : acc[q]; + t1[q] = keep_hi ? acc[q] : acc[q + H1]; + x1[q] = keep_hi ? xvs[q + H1] : xvs[q]; + } + #pragma unroll + for (int q = 0; q < H1; ++q) + m1[q] += __shfl_xor_sync(0xffffffff, t1[q], 16); + float m2[H2], t2[H2], x2[H2]; + #pragma unroll + for (int q = 0; q < H2; ++q) { + m2[q] = keep_hi2 ? m1[q + H2] : m1[q]; + t2[q] = keep_hi2 ? m1[q] : m1[q + H2]; + x2[q] = keep_hi2 ? x1[q + H2] : x1[q]; + } + #pragma unroll + for (int q = 0; q < H2; ++q) + m2[q] += __shfl_xor_sync(0xffffffff, t2[q], 8); + #pragma unroll + for (int q = 0; q < H2; ++q) { + #pragma unroll + for (int off = 4; off > 0; off >>= 1) + m2[q] += __shfl_xor_sync(0xffffffff, m2[q], off); + } + if ((lane & 7) == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + + p0 + (keep_hi ? H1 : 0) + (keep_hi2 ? H2 : 0); + #pragma unroll + for (int q = 0; q < H2; ++q) + dst[q] = __float2bfloat16(m2[q] + x2[q] * Dh); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc[q] += __shfl_xor_sync(0xffffffff, acc[q], off); + } + if (lane == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + p0; + #pragma unroll + for (int q = 0; q < PW; ++q) + dst[q] = __float2bfloat16(acc[q] + xvs[q] * Dh); + } + } + } + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + st[(size_t)q * MB_STATE + lane + j * 32] = sreg[q][j]; + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, + const float *dtda, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + (void)zxbcdt; + if (n_seq > 1) + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens / n_seq, n_seq, slot, seq_slot); + else + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens, 1, slot, nullptr); +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16( + held[j] * inv * __bfloat162float(L.gnorm_w[idx])); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = BATCH * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / BB_ATT_HEADS; + const int head = pair % BB_ATT_HEADS; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const int kv = head / kv_per_q; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *q = base + head * BB_HEAD_DIM; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (head % kv_per_q == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + for (int tpos = 0; tpos < T; ++tpos) { +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new + : kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#endif + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m, s); + const float rescale = __expf(m - m_new); + const float w = __expf(s - m_new); + denom = denom * rescale + w; +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *vt = kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#else + const bf16 *vt = cur_token + ? v_new + : kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + w * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + constexpr int CHUNKS = BB_HEAD_DIM / 8; + const int total = n_tokens * BB_KV_HEADS * CHUNKS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int c = idx % CHUNKS; + const int rest = idx / CHUNKS; + const int t = rest / BB_KV_HEADS; + const int kv = rest - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM + c * 8; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM + c * 8; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv) + c * 8; + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv) + c * 8; + *reinterpret_cast(kc) = *reinterpret_cast(k_new); + *reinterpret_cast(vc) = *reinterpret_cast(v_new); + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float qv[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qv[j] = __bfloat162float(q[lane + j * 32]); + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + const int *const krow = + L.block_table + ((size_t)seq * 2 + 0) * L.max_blocks_per_seq; + const int *const vrow = + L.block_table + ((size_t)seq * 2 + 1) * L.max_blocks_per_seq; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + const int head_bias = kv * L.tokens_per_block * BB_HEAD_DIM; + int page = -1; + const bf16 *kpage = nullptr, *vpage = nullptr; + float kn[PER_LANE], vn[PER_LANE]; + { + page = 0; + kpage = L.pool_base + (size_t)(krow[0] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + (size_t)(vrow[0] & 0x7FFFFFFF) * elems_per_block; + const int at = head_bias + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + for (int s = 0; s <= pos; ++s) { + float kc[PER_LANE], vc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { kc[j] = kn[j]; vc[j] = vn[j]; } + if (s < pos) { + const int np = (s + 1) / L.tokens_per_block; + if (np != page) { + page = np; + kpage = L.pool_base + + (size_t)(krow[np] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + + (size_t)(vrow[np] & 0x7FFFFFFF) * elems_per_block; + } + const int at = head_bias + + ((s + 1) % L.tokens_per_block) * BB_HEAD_DIM + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) dot += qv[j] * kc[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[j] = acc[j] * rescale + wt * vc[j]; + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH, unsigned int *tile_masks = nullptr, + bf16 *gath = nullptr, int *gidx = nullptr, int *ecnt = nullptr, + int estride = 0, bf16 *row_stage = nullptr) { + __shared__ float score[BB_EXPERTS]; + __shared__ int s_pick[BB_TOP_K], s_slot[BB_TOP_K]; + constexpr int CH = D_MODEL / 8; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + { + const float4 *gsrc = + reinterpret_cast(normed + (size_t)row * D_MODEL); + float4 *gdst = reinterpret_cast(row_stage); + for (int i = threadIdx.x; i < CH; i += THREADS) gdst[i] = gsrc[i]; + } + if ((int)threadIdx.x < BB_EXPERTS) + moe_w[(size_t)row * BB_EXPERTS + threadIdx.x] = 0.f; + __syncthreads(); + for (int e = warp; e < BB_EXPERTS; e += WARPS) { + const bf16 *wr = gate_w + (size_t)e * D_MODEL; + float acc = 0.f; + for (int i = lane; i < D_MODEL; i += 32) + acc += __bfloat162float(row_stage[i]) * + __bfloat162float(wr[i]); + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) score[e] = 1.f / (1.f + __expf(-acc)); + } + __syncthreads(); + if (threadIdx.x == 0) { + float remaining[BB_EXPERTS]; + int pick[BB_TOP_K]; + for (int e = 0; e < BB_EXPERTS; ++e) + remaining[e] = score[e] + __bfloat162float(gate_bias[e]); + float total = 0.f; + for (int k = 0; k < BB_TOP_K; ++k) { + int best = 0; + for (int e = 1; e < BB_EXPERTS; ++e) + if (remaining[e] > remaining[best]) best = e; + pick[k] = best; + total += score[best]; + remaining[best] = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + unsigned int bits = 0u; + for (int k = 0; k < BB_TOP_K; ++k) { + moe_w[row * BB_EXPERTS + pick[k]] = score[pick[k]] * norm; + bits |= 1u << pick[k]; + s_pick[k] = pick[k]; + } + if (tile_masks) atomicOr(&tile_masks[row / BATCH_PAD], bits); + } + __syncthreads(); + if (gath) { + if ((int)threadIdx.x < BB_TOP_K) { + const int k = threadIdx.x; + const int slot = atomicAdd(&ecnt[s_pick[k]], 1); + s_slot[k] = slot; + gidx[(size_t)s_pick[k] * estride + slot] = row; + } + __syncthreads(); + const float4 *src = reinterpret_cast(row_stage); + for (int k = 0; k < BB_TOP_K; ++k) { + float4 *dst = reinterpret_cast( + gath + ((size_t)s_pick[k] * estride + s_slot[k]) * D_MODEL); + for (int i = threadIdx.x; i < CH; i += THREADS) dst[i] = src[i]; + } + } + __syncthreads(); + } +} +__device__ __forceinline__ unsigned moe_row_mask(const float *routes, int row) { + unsigned m = 0u; + const float *rw = routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) if (rw[e] != 0.f) m |= 1u << e; + return m; +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx, + unsigned mask) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + for (unsigned m = mask; m; m &= m - 1u) + value = __fadd_rn(value, parts[base + (size_t)(__ffs(m) - 1) * D_MODEL]); + return __float2bfloat16( + __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL])); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + uint64_t *keys, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int n = threadIdx.x; n < SORT_N; n += THREADS) { + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + keys[n] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + __syncthreads(); + for (int k = 2; k <= SORT_N; k <<= 1) { + if (k >= 64) __syncthreads(); + for (int j = k >> 1; j > 0; j >>= 1) { + for (int i = threadIdx.x; i < SORT_N; i += THREADS) { + const int partner = i ^ j; + if (partner > i) { + const bool descending = ((i & k) == 0); + if ((keys[i] < keys[partner]) == descending) { + const uint64_t t = keys[i]; + keys[i] = keys[partner]; + keys[partner] = t; + } + } + } + if (j >= 32) __syncthreads(); + else __syncwarp(); + } + } + __syncthreads(); + if (threadIdx.x == 0) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float v; + int n; + if (!eos_used && eos_v > av) { + v = eos_v; + n = AUDIO_EOS; + eos_used = true; + } else { + v = av; + n = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = + v + gumbel[((size_t)row * N_POS + pos) * TOP_K + j]; + if (noised > best_noised) { best_noised = noised; best_code = n; } + } + out_codes[row] = best_code; + } + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void prefill_expert_masks(const float *routes, unsigned int *masks, int rows) { + const int lane=threadIdx.x&31, warp=threadIdx.x>>5; + const int nt=(rows+BATCH_PAD-1)/BATCH_PAD; + for (int t=blockIdx.x*WARPS+warp;t(gath + g * D_MODEL + c * 8); + if (local < ecnt[e]) + *dst = *reinterpret_cast( + normed + (size_t)gidx[g] * D_MODEL + c * 8); + else + *dst = make_float4(0.f, 0.f, 0.f, 0.f); + } +} +__device__ __forceinline__ int gemm_eff(int tiles, int reps) { + if (NCONS <= 1) return 1; + const int g = (int)gridDim.x; + const int rounds1 = (tiles * reps + g - 1) / g; + if constexpr (NCONS == 3) { + const int max_eff = reps <= BB_EXPERTS ? 2 : 3; + const int rounds2 = (((tiles + 1) / 2) * reps + g - 1) / g; + if (max_eff == 2) return rounds2 < rounds1 ? 2 : 1; + const int rounds3 = (((tiles + 2) / 3) * reps + g - 1) / g; + if (rounds3 < rounds2 && rounds3 < rounds1) return 3; + return rounds2 < rounds1 ? 2 : 1; + } else { + const int roundsN = (((tiles + NCONS - 1) / NCONS) * reps + g - 1) / g; + return roundsN < rounds1 ? NCONS : 1; + } +} +__device__ __forceinline__ int gemm_groups(int tiles, int reps) { + const int eff = gemm_eff(tiles, reps); + return (tiles + eff - 1) / eff; +} +__device__ __forceinline__ int gemm_consumer_id() { + const int g = warpgroup::groupid(); + if (g == 1) return -1; + const int c = g == 0 ? 0 : g - 1; + return c < NCONS ? c : -1; +} +__device__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][BATCH_PAD]; + __shared__ int moe_gidx_stage[NCONS][BATCH_PAD]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + a_s[s], a, {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + a_s[s], a, {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / BATCH_PAD) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, BATCH_PAD> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], a_s[s]); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * BATCH_PAD; + if (epi == EPI_BIAS) { + warpgroup::store(o_sf, accum); + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + dest[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(o_sf[int2{col, row}] + + __bfloat162float(bias[base + col])); + } + warpgroup::sync(1 + cons); + continue; + } + warpgroup::store(o_s[cons], accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_s[cons][int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__device__ void gemm_narrow( + const weight_gl &w, const act_gl_n &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = PAD_N, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][PAD_N]; + __shared__ int moe_gidx_stage[NCONS][PAD_N]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + PAD_N - 1) / PAD_N; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + PAD_N - 1) / PAD_N; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / PAD_N) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, PAD_N> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], + *reinterpret_cast(&a_s[s])); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * PAD_N; + out_tile_n &o_sn = *reinterpret_cast(&o_s[cons]); + warpgroup::store(o_sn, accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_sn[int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&w_s)[STAGES * NCONS] = + alloc.allocate(); + act_tile (&a_s)[STAGES] = alloc.allocate(); + out_tile (&o_s)[NCONS] = alloc.allocate(); +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 + static_assert(NCONS * sizeof(out_tile) >= sizeof(out_tile_f), + "aliased head staging must fit the output tiles"); + out_tile_f &o_sf = *reinterpret_cast(&o_s[0]); +#else + out_tile_f (&o_sf) = alloc.allocate(); +#endif + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_WARPS + 2]; +#if MK2_PREFILL_ONLY + bf16 *const route_stage = reinterpret_cast(&o_s[0]); + static_assert(sizeof(o_s) >= D_MODEL * sizeof(bf16), + "router row stage must fit the output tiles"); +#else + __shared__ __align__(16) uint64_t sort_keys[SORT_N]; + bf16 *const route_stage = reinterpret_cast(&sort_keys[0]); + static_assert(sizeof(sort_keys) >= D_MODEL * sizeof(bf16), + "router row stage must fit the sampler workspace"); +#endif + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS * NCONS, 0); + for (int c = 0; c < CONSUMER_WARPS * NCONS; ++c) arrive(drained[s]); + } + } + int gemm_iter = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map); + grid_barrier(g.barrier); + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map); + grid_barrier(g.barrier); + mamba_gated_norm(L, bb->zxbcdt, bb->gated); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, nullptr, 1, BATCH, + bb->moe_w, BB_MOE_INTER / OUT_ATOM); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + 1, BATCH, bb->moe_w, D_MODEL / OUT_ATOM); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_deterministic(bb->moe_out, idx, + (1u << BB_EXPERTS) - 1u); + grid_barrier(g.barrier); + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int original_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + bool identical_cohort = false; + if constexpr (BATCH > 1) { + if (bb->native_pf_cohort_check && bb->native_prefill_reset && + bb->pf_n_seq == BATCH && !bb->pf_fuse_decode && + !bb->pf_layer_out && !bb->moe_w_inject && !bb->moe_w_audit) { + if (blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[0] = 0; + grid_barrier(g.barrier); + const auto *bits = reinterpret_cast(bb->prompt); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < original_total * D_MODEL; idx += gridDim.x * THREADS) { + const int canonical = (idx / (BATCH * D_MODEL)) * BATCH * D_MODEL + idx % D_MODEL; + if (bits[idx] != bits[canonical]) atomicOr(bb->native_pf_cohort_check, 1u); + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + if (bb->pf_seq_slot[row] != row || bb->slot_map[row] != row) + atomicOr(bb->native_pf_cohort_check, 1u); + grid_barrier(g.barrier); + identical_cohort = bb->native_pf_cohort_check[0] == 0; + } + } + if (!identical_cohort) { + if (bb->native_pf_cohort_check && blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[1] = 0; + grid_barrier(g.barrier); + return; + } + const int n_total = original_total / BATCH; + if (bb->native_pf_cohort_check && blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[1] = identical_cohort ? 1 : BATCH; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int reset_slots = identical_cohort ? 1 : BATCH; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * reset_slots * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % reset_slots; + const MambaLayer &L = bb->mamba[row_layer / reset_slots]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = original_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + const int src = identical_cohort ? (idx / D_MODEL) * BATCH * D_MODEL + idx % D_MODEL : idx; + bb->pf_h[idx] = bb->prompt[src]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = identical_cohort ? 1 : (bb->pf_n_seq > 0 ? bb->pf_n_seq : 1); + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T, + (!bb->pf_layer_out && layer > 0 && + bb->moe_of_layer[layer - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + if (bb->native_pf_expert_masks) + for (int i = blockIdx.x * THREADS + threadIdx.x; i < n_tok; + i += gridDim.x * THREADS) + bb->native_pf_expert_masks[i] = 0u; + if (bb->pf_moe_ecnt) + for (int e = blockIdx.x * THREADS + threadIdx.x; e < BB_EXPERTS; + e += gridDim.x * THREADS) + bb->pf_moe_ecnt[e] = 0; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + for(size_t j=blockIdx.x*THREADS+threadIdx.x;jpf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + mamba_dtda_prefill(L, bb->pf_zxbcdt, bb->pf_dtda, T); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_dtda, + bb->pf_gated, T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + for(size_t j=blockIdx.x*THREADS+threadIdx.x;jpf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, + T, bb->native_pf_expert_masks, bb->pf_moe_gath, + bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, route_stage); + } + grid_barrier(g.barrier); + int max_ecnt = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + max_ecnt = max(max_ecnt, bb->pf_moe_ecnt[e]); + const bool narrow = BATCH_PAD > PAD_N && max_ecnt <= PAD_N; + const int etile = narrow ? PAD_N : BATCH_PAD; + int nt_sum = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_sum += (bb->pf_moe_ecnt[e] + etile - 1) / etile; + const int up_steps = + gemm_groups(BB_MOE_INTER / OUT_ATOM, nt_sum) * + nt_sum * (D_MODEL / KCHUNK); + const int shup_steps = + gemm_groups(BB_SHARED_INTER / OUT_ATOM, n_tok) * + n_tok * (D_MODEL / KCHUNK); + const int down_steps = + gemm_groups(D_MODEL / OUT_ATOM, nt_sum) * + nt_sum * (BB_MOE_INTER / KCHUNK); + const int shdown_steps = + gemm_groups(D_MODEL / OUT_ATOM, n_tok) * + n_tok * (BB_SHARED_INTER / KCHUNK); + if (narrow) + gemm_narrow(L.exp_up, bb->pf_moe_gath_vn, + BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, PAD_N, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + else + gemm(L.exp_up, bb->pf_moe_gath_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, BATCH_PAD, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + up_steps, up_steps + shup_steps); + grid_barrier(g.barrier); + if (narrow) + gemm_narrow(L.exp_down, bb->pf_moe_hc_vn, + BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, PAD_N, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + else + gemm(L.exp_down, bb->pf_moe_hc_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + down_steps, down_steps + shdown_steps); + grid_barrier(g.barrier); + if (bb->pf_layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * D_MODEL; idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = + bb_moe_fold_deterministic(bb->pf_moe_out, idx, + moe_row_mask(bb->pf_moe_w, idx / D_MODEL)); + grid_barrier(g.barrier); + } + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, + identical_cohort ? bb->pf_normed : bb->pf_final, ln_scratch, + n_total, + (!bb->pf_layer_out && n_backbone_layers > 0 && + bb->moe_of_layer[n_backbone_layers - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + grid_barrier(g.barrier); + if (identical_cohort) { + const int worker = blockIdx.x * THREADS + threadIdx.x; + const int workers = gridDim.x * THREADS; + for (int i = worker; i < original_total * D_MODEL; i += workers) + bb->pf_final[i] = bb->pf_normed[(i / (BATCH * D_MODEL)) * D_MODEL + i % D_MODEL]; + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = worker; i < 14 * (BATCH - 1) * CV; i += workers) { + const int cell = i % CV, pair = i / CV; + const int row = pair % (BATCH - 1) + 1; + const MambaLayer &L = bb->mamba[pair / (BATCH - 1)]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = L.conv_state[cell]; + } + const int kv_columns = BB_KV_HEADS * BB_HEAD_DIM; + for (size_t i = worker; i < size_t(3) * 2 * (BATCH - 1) * n_total * kv_columns; i += workers) { + size_t q = i / kv_columns; + const int col = i % kv_columns; + const int pos = q % n_total; q /= n_total; + const int row = q % (BATCH - 1) + 1; q /= BATCH - 1; + const int kind = q % 2, layer = q / 2; + const AttnLayer &L = bb->attn[layer]; + const int offset = kv_offset(L, pos, col / BB_HEAD_DIM) + col % BB_HEAD_DIM; + kv_slot(L, row, kind, pos)[offset] = kv_slot(L, 0, kind, pos)[offset]; + } + grid_barrier(g.barrier); + } + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } +#if MK2_PREFILL_ONLY + return; +#else + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, bb->phoneme_b); + grid_barrier(g.barrier); + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + if (threadIdx.x == 0) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = 0; n < BB_PHONEME; ++n) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + bb->phoneme_out[row] = arg; + } + __syncthreads(); + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + for (int i = lane; i < D_HEAD; i += 32) { + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + for (int t = 0; t <= pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + for (int i = lane; i < D_HEAD; i += 32) { + float acc = 0.f; + for (int t = 0; t <= pos; ++t) + acc += sc[t] * __bfloat162float( + g.vcache[cache_base + t * D_HEAD + i]); + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + i] = + __float2bfloat16(acc); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, g.sampled_pos, + sort_keys, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { +#ifdef ABLATE_SAMPLE + __syncthreads(); +#endif + } else { + grid_barrier(g.barrier); + } + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x < BATCH && blockIdx.x == 0) { + const int code = g.sampled_pos[threadIdx.x]; + g.sampled[threadIdx.x * N_POS + pos] = code; + if (g.native_codes) + g.native_codes[threadIdx.x * N_POS + pos] = (long long)code; + } + grid_barrier(g.barrier); + } + if (g.native_seq_len && blockIdx.x == 0 && threadIdx.x < BATCH) + g.native_seq_len[threadIdx.x] += 1; + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +#endif +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 2 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK), + 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = MK2_PREFILL_ONLY + ? out->sms * out->blocks_per_sm + : std::min(out->sms * out->blocks_per_sm, D_FFN / OUT_ATOM); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + 0, + nullptr, + nullptr, + nullptr, nullptr, nullptr, nullptr, nullptr, 0, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "dedup_native_prefill_impl.inc" diff --git a/deploy/nemotron_native/kernels/b256/regions/prefill/native/prefill/mk2_abi.cuh b/deploy/nemotron_native/kernels/b256/regions/prefill/native/prefill/mk2_abi.cuh new file mode 100644 index 0000000..95a253b --- /dev/null +++ b/deploy/nemotron_native/kernels/b256/regions/prefill/native/prefill/mk2_abi.cuh @@ -0,0 +1,82 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} diff --git a/deploy/nemotron_native/kernels/b256/regions/prefill/native/prefill/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/b256/regions/prefill/native/prefill/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/b256/regions/prefill/native/prefill/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/b256/regions/prefill/native/prefill/prefill.cu b/deploy/nemotron_native/kernels/b256/regions/prefill/native/prefill/prefill.cu new file mode 100644 index 0000000..e7091f5 --- /dev/null +++ b/deploy/nemotron_native/kernels/b256/regions/prefill/native/prefill/prefill.cu @@ -0,0 +1,18 @@ +#define MK2_BATCH 256 +#define MK2_BATCH_TILE 64 +#define MK2_PREFILL_ONLY 1 +#define STAGES 3 +#define KCHUNK 128 +#define MK2_GEMM_CONSUMERS 3 +#define MK2_THREADS 512 +#define MK2_NO_MAIN 1 +#define MK2_DIRECT_INPUT 1 +#define MK2_B1_LOCAL_BARRIERS 1 +#define MK2_SKIP_HIDDEN_SNAPSHOT 1 +#define MK2_ACTIVE_EXPERTS 1 +#define MK2_SSM_SPLITS 1 +#define PF_SSM_PW_MULTI 12 +#define PF_ACTIVE_TOKEN_EXPERTS 1 +#include +#include +#include "dedup_predictor_mk2.cuh" diff --git a/deploy/nemotron_native/kernels/b256/regions/prefill/prefill.cu b/deploy/nemotron_native/kernels/b256/regions/prefill/prefill.cu new file mode 100644 index 0000000..eda7dbd --- /dev/null +++ b/deploy/nemotron_native/kernels/b256/regions/prefill/prefill.cu @@ -0,0 +1,20 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace b256_prefill { +#include "native/prefill/prefill.cu" +} diff --git a/deploy/nemotron_native/kernels/b256/regions/stage0/native/stage0/mk2_abi.cuh b/deploy/nemotron_native/kernels/b256/regions/stage0/native/stage0/mk2_abi.cuh new file mode 100644 index 0000000..bf1e415 --- /dev/null +++ b/deploy/nemotron_native/kernels/b256/regions/stage0/native/stage0/mk2_abi.cuh @@ -0,0 +1,85 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_set_native_predictor_active_mask(MK2Predictor *p, const unsigned char *mask); +int mk2_native_startup_skip_abi_version(); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *, int grid); diff --git a/deploy/nemotron_native/kernels/b256/regions/stage0/native/stage0/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/b256/regions/stage0/native/stage0/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/b256/regions/stage0/native/stage0/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/b256/regions/stage0/native/stage0/native_prefill_impl.inc b/deploy/nemotron_native/kernels/b256/regions/stage0/native/stage0/native_prefill_impl.inc new file mode 100644 index 0000000..a254308 --- /dev/null +++ b/deploy/nemotron_native/kernels/b256/regions/stage0/native/stage0/native_prefill_impl.inc @@ -0,0 +1,198 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 16 || BATCH == 32 || BATCH == 64)) + throw std::runtime_error("native prefill requires matching batch 1/16/32/64"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (padded > INT_MAX / ((BB_EXPERTS + 1) * D_MODEL)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + float *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(float))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(float))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} diff --git a/deploy/nemotron_native/kernels/b256/regions/stage0/native/stage0/predictor_mk2.cu b/deploy/nemotron_native/kernels/b256/regions/stage0/native/stage0/predictor_mk2.cu new file mode 100644 index 0000000..ddf223f --- /dev/null +++ b/deploy/nemotron_native/kernels/b256/regions/stage0/native/stage0/predictor_mk2.cu @@ -0,0 +1,6107 @@ +#include "../../../../../inflight/progress.cuh" +#if MK2_INFLIGHT +#include +#endif +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#define MK2_FUSE_CONV (MK2_SSM_SPLITS == 1) +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +constexpr int BATCH_PAD = 64; +constexpr int DECODE_TILES = BATCH / BATCH_PAD; +constexpr int DECODE_ROWS = BATCH_PAD; +static_assert(BATCH == 128 || BATCH == 256 || BATCH == 512 || BATCH == 1024 || BATCH == 2048); +static_assert(BATCH % BATCH_PAD == 0); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +#ifndef NMULT +#define NMULT 1 +#endif +#ifndef SSM_PREFETCH_MAX +#define SSM_PREFETCH_MAX 8 +#endif +constexpr int NM_MAX = (NMULT > 1 && DECODE_TILES % NMULT == 0) ? NMULT : 1; +#ifndef MK2_OUT_MULT +#define MK2_OUT_MULT 1 +#endif +constexpr int MT = MK2_OUT_MULT; +static_assert(MT >= 1 && MT <= NM_MAX, "an item holds at most NM_MAX accumulators"); +constexpr int SLOTS = 1 + NM_MAX; +static_assert(MT + 1 <= SLOTS); +constexpr int THREADS = 256; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +constexpr int NORM_SCRATCH_STRIDE = 2 * WARPS + 2; +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_ACC_F32 }; +#define EPI_ROW(idx) ((OUT_ATOM == 64 && BATCH_PAD == 64) \ + ? (((idx) & 15) | ((((idx) >> 6) & 3) << 4)) : ((idx) / (OUT_ATOM / 8))) +#define EPI_OCT(idx) ((OUT_ATOM == 64 && BATCH_PAD == 64) \ + ? ((((idx) >> 4) & 3) | ((((idx) >> 8) & 1) << 2)) : ((idx) % (OUT_ATOM / 8))) +struct alignas(16) bf16_oct { bf16 x0, x1, x2, x3, x4, x5, x6, x7; }; +static_assert(sizeof(bf16_oct) == 16); +#include +constexpr int R9_STATE_MODE=1; +using r9_state_t=bf16; +__device__ __forceinline__ int r9_scale_index(int head,int channel,int n) {return (((head/8)*4+(head%8)/2)*4+channel/6)*128+n;} +__device__ __forceinline__ float r9_load(r9_state_t v,float scale) { +return __bfloat162float(v); +} +__device__ __forceinline__ r9_state_t r9_store(float x,float scale) { +return __float2bfloat16(x); +} +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; + const float *r8_initial_ssm = nullptr; + const unsigned int *r8_initial_ready = nullptr; + const int *r8_seq_len = nullptr; + int r8_initial_positions = 0; + const float *r9_scales = nullptr; + const int *r9_map = nullptr; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + bf16 *r6_packed = nullptr; + int *r6_map = nullptr,*r6_counts = nullptr; + act_gl *r6_views = nullptr; + int *r6_lut = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +#define BARRIER_GEN_WORD 32 +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + constexpr unsigned long long GRID_BLOCKS = 132ull; + auto *wide_state = reinterpret_cast(state); + const unsigned long long arrival = atomicAdd(wide_state, 1ull); + const unsigned int rank = (unsigned int)(arrival % GRID_BLOCKS); + if (rank != GRID_BLOCKS - 1u) { + const unsigned long long target = arrival - rank + GRID_BLOCKS; + volatile unsigned long long *counter = wide_state; + while (*counter < target) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / THREADS; + float wv[PER_THREAD]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) wv[j] = weight[threadIdx.x + j * THREADS]; + float v[PER_THREAD], v_next[PER_THREAD]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + v[j] = __bfloat162float( + src[(size_t)blockIdx.x * D_MODEL + threadIdx.x + j * THREADS]); + for (int row = blockIdx.x, row_iter = 0; row < BATCH; + row += gridDim.x, ++row_iter) { + float *row_scratch = scratch + (row_iter & 1) * NORM_SCRATCH_STRIDE; + const int next_row = row + (int)gridDim.x; + if (next_row < BATCH) { + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + v_next[j] = __bfloat162float( + src[(size_t)next_row * D_MODEL + threadIdx.x + j * THREADS]); + } + float sum = 0.f, sq = 0.f; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) { + row_scratch[warp] = sum; + row_scratch[WARPS + warp] = sq; + } + __syncthreads(); + float s = 0.f, q = 0.f; + #pragma unroll + for (int w = 0; w < WARPS; ++w) { + s += row_scratch[w]; + q += row_scratch[WARPS + w]; + } + const float mean = s / D_MODEL; + const float inv = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * wv[j]); + } + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) v[j] = v_next[j]; + } +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, bf16 *snapshot = nullptr, int snapshot_layer = -1, + float *snapshot_fp32 = nullptr, float *out_fp32 = nullptr, + const float *fold_src = nullptr, const float *fold_routes = nullptr) { + constexpr int PER_THREAD = D_MODEL / THREADS; + bf16 wv[PER_THREAD]; + extern __shared__ int __shm[]; + bf16 *const row_stage = reinterpret_cast(__shm); + constexpr int ROW_SLOT = 2 * D_MODEL; + constexpr int ROW_VECTORS = D_MODEL / 8; + if ((int)blockIdx.x < n_rows && threadIdx.x < ROW_VECTORS) { + const int i = 8 * threadIdx.x; + const size_t at = (size_t)blockIdx.x * D_MODEL + i; + const uint32_t rs = static_cast( + __cvta_generic_to_shared(row_stage + D_MODEL + i)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;" + :: "r"(rs), "l"(r + at) : "memory"); + if (!fold_src) { + const uint32_t hs = static_cast( + __cvta_generic_to_shared(row_stage + i)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;" + :: "r"(hs), "l"(h + at) : "memory"); + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + wv[j] = weight[threadIdx.x + j * THREADS]; + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); + __syncthreads(); + for (int row = blockIdx.x, row_iter = 0; row < n_rows; + row += gridDim.x, ++row_iter) { + bf16 *const current = row_stage + (row_iter & 1) * ROW_SLOT; + const int next_row = row + (int)gridDim.x; + if (next_row < n_rows) { + bf16 *const next = row_stage + ((row_iter + 1) & 1) * ROW_SLOT; + if (threadIdx.x < ROW_VECTORS) { + const int i = 8 * threadIdx.x; + const size_t at = (size_t)next_row * D_MODEL + i; + const uint32_t rs = static_cast( + __cvta_generic_to_shared(next + D_MODEL + i)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;" + :: "r"(rs), "l"(r + at) : "memory"); + if (!fold_src) { + const uint32_t hs = static_cast( + __cvta_generic_to_shared(next + i)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;" + :: "r"(hs), "l"(h + at) : "memory"); + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } + float *row_scratch = scratch + (row_iter & 1) * NORM_SCRATCH_STRIDE; + const size_t base = (size_t)row * D_MODEL; + float fv[PER_THREAD]; + if (fold_src) { + unsigned mask = 0u; + const float *rw = fold_routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) + if (rw[e] != 0.f) mask |= 1u << e; + const float *pe = fold_src + + (size_t)row * (BB_EXPERTS + 1) * D_MODEL + threadIdx.x; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) fv[j] = 0.f; + #pragma unroll + for (int e = 0; e <= BB_EXPERTS; ++e) { + const bool take = e == BB_EXPERTS || ((mask >> e) & 1u); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + fv[j] = __fadd_rn(fv[j], + take ? pe[(size_t)e * D_MODEL + j * THREADS] : 0.f); + } + } + float v[PER_THREAD]; + float sq = 0.f; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + const bf16 hv = fold_src ? __float2bfloat16(fv[j]) : current[i]; + if (fold_src) h[base + i] = hv; + if (snapshot) + snapshot[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = hv; + if (snapshot_fp32) + snapshot_fp32[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = + __bfloat162float(hv); + v[j] = __bfloat162float(hv) + + __bfloat162float(current[D_MODEL + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) row_scratch[warp] = sq; + __syncthreads(); + float s = 0.f; + #pragma unroll + for (int w = 0; w < WARPS; ++w) s += row_scratch[w]; + const float inv = rsqrtf(s / D_MODEL + RMS_EPS); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + const bf16 normed = __float2bfloat16( + __bfloat162float(wv[j]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + out[base + i] = normed; + if (out_fp32) out_fp32[base + i] = __bfloat162float(normed); + } + if (next_row < n_rows) { + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); + __syncthreads(); + } + } +} +__device__ __forceinline__ bf16 mamba_conv_channel( + const MambaLayer &L, bf16 *state_slot, const bf16 *zrow, int c) { + const unsigned ch = (unsigned)c; + const bf16 *w = L.conv_w + ch * (unsigned)MB_CONV_K; + bf16 *state = state_slot + ch * (unsigned)(MB_CONV_K - 1); + const bf16 rawb = zrow[c]; + const float raw = __bfloat162float(rawb); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = rawb; + return __float2bfloat16(acc); +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map MK2_PROGRESS_PARAMETER(progress_mask)) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + MK2_PROGRESS_GUARD(progress_mask, row); + const int c = idx - row * MB_CONV_DIM; + xbc_act[(size_t)row * MB_CONV_DIM + c] = mamba_conv_channel( + L, L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride, + zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER, c); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map MK2_PROGRESS_PARAMETER(progress_mask)) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int sb4 = (lane >> 4) & 1, sb3 = (lane >> 3) & 1; + const int sb2 = (lane >> 2) & 1, sb1 = (lane >> 1) & 1; + const int ssm_own_ch = 6 * sb4 + 3 * sb3 + 2 * sb2 + sb1; + const bool ssm_owner = (2 * sb2 + sb1) < 3 && (lane & 1) == 0; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + constexpr bool FUSE_GNORM = (split == 1); + extern __shared__ int __shm[]; + bf16 (*const gn_stage)[MB_GROUP_SIZE] = + reinterpret_cast(&__shm[0]); + constexpr bool FUSE_CONV = (MK2_FUSE_CONV != 0); + bf16 (*const conv_x)[MB_GROUP_SIZE] = gn_stage + WARPS; + constexpr int SSM_NSTATE = MB_STATE / 32; + constexpr int SSM_PF = channels < SSM_PREFETCH_MAX ? channels : SSM_PREFETCH_MAX; + constexpr int SSM_GROUP_BYTES = SSM_PF * MB_STATE * (int)sizeof(r9_state_t); + constexpr bool SSM_SMEM_OK = R9_STATE_MODE == 1 && channels == 2 * SSM_PF && + SSM_GROUP_BYTES % 512 == 0; + constexpr int SSM_RING = 4; + constexpr int SSM_GROUPS_PER_HEAD = channels / SSM_PF; + constexpr int SSM_PAIRS = MB_HEADS_PER_GROUP * SSM_GROUPS_PER_HEAD; + static_assert(!SSM_SMEM_OK || (SSM_RING & (SSM_RING - 1)) == 0, + "the flat (head, channel group) ring index is masked"); + static_assert(!SSM_SMEM_OK || SSM_PAIRS >= SSM_RING, + "the prologue must be able to fill the ring"); + static_assert(!SSM_SMEM_OK || + 2 * WARPS * MB_GROUP_SIZE * sizeof(bf16) + + WARPS * SSM_RING * SSM_GROUP_BYTES <= STAGES * SLOTS * sizeof(weight_tile)); + constexpr int SSM_F32_RING = 2; + constexpr int SSM_F32_GROUP_BYTES = SSM_PF * MB_STATE * (int)sizeof(float); + constexpr bool SSM_F32_OK = SSM_SMEM_OK && SSM_F32_GROUP_BYTES % 512 == 0 && + (SSM_F32_RING & (SSM_F32_RING - 1)) == 0 && SSM_PAIRS >= SSM_F32_RING && + SSM_F32_RING * SSM_F32_GROUP_BYTES <= SSM_RING * SSM_GROUP_BYTES; + char *const ssm_stage = reinterpret_cast(conv_x + WARPS); + const int pairs = BATCH * MB_GROUPS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_GROUPS * split); + MK2_PROGRESS_GUARD(progress_mask, row); + const int group = (pair / split) % MB_GROUPS; + const int first_channel = (pair % split) * channels; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const bool r8_initial = L.r8_initial_ssm && L.r8_initial_ready && + *L.r8_initial_ready == 1 && L.r8_seq_len[row] == L.r8_initial_positions; + r9_state_t *st = reinterpret_cast(L.ssm_state) + + (size_t)slot_map[row] * L.ssm_slot_stride; + char *const warp_ssm = + ssm_stage + (size_t)warp * SSM_RING * SSM_GROUP_BYTES; + auto issue_ssm = [&](int head_in_group, int channel0, int buffer) { + const int head = group * MB_HEADS_PER_GROUP + head_in_group; + const r9_state_t *src = st + + (size_t)(head * MB_HEAD_DIM + channel0) * MB_STATE; + const uint32_t dst = (uint32_t)__cvta_generic_to_shared( + warp_ssm + (size_t)buffer * SSM_GROUP_BYTES + lane * 16); + #pragma unroll + for (int chunk = 0; chunk < SSM_GROUP_BYTES / 512; ++chunk) + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: + "r"(dst + (uint32_t)chunk * 512), + "l"(reinterpret_cast(src) + + (size_t)chunk * 512 + lane * 16) : "memory"); + asm volatile("cp.async.commit_group;\n" ::: "memory"); + }; + auto issue_ssm_f32 = [&](int head_in_group, int channel0, int buffer) { + const int head = group * MB_HEADS_PER_GROUP + head_in_group; + const float *src = L.r8_initial_ssm + + (size_t)(head * MB_HEAD_DIM + channel0) * MB_STATE; + const uint32_t dst = (uint32_t)__cvta_generic_to_shared( + warp_ssm + (size_t)buffer * SSM_F32_GROUP_BYTES + lane * 16); + #pragma unroll + for (int chunk = 0; chunk < SSM_F32_GROUP_BYTES / 512; ++chunk) + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: + "r"(dst + (uint32_t)chunk * 512), + "l"(reinterpret_cast(src) + + (size_t)chunk * 512 + lane * 16) : "memory"); + asm volatile("cp.async.commit_group;\n" ::: "memory"); + }; + if constexpr (SSM_SMEM_OK) { + if (!r8_initial) { + #pragma unroll + for (int k = 0; k < SSM_RING; ++k) + issue_ssm(k / SSM_GROUPS_PER_HEAD, + (k % SSM_GROUPS_PER_HEAD) * SSM_PF, k); + } else if constexpr (SSM_F32_OK) { + #pragma unroll + for (int k = 0; k < SSM_F32_RING; ++k) + issue_ssm_f32(k / SSM_GROUPS_PER_HEAD, + (k % SSM_GROUPS_PER_HEAD) * SSM_PF, k); + } + } + float Braw[MB_STATE / 32]; + float Cv[MB_STATE / 32]; + const bf16 *xrow; + if constexpr (FUSE_CONV) { + bf16 *const cslot = + L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride; + const bf16 *const zrow = + zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER; + static_assert(MB_CONV_K == 4, "the state record is three 32-bit words"); + static_assert(MB_GROUP_SIZE % 64 == 0 && MB_STATE % 64 == 0, + "one paired round covers 64 channels of a warp"); + auto wide_lo = [](unsigned int v) { return __uint_as_float(v << 16); }; + auto wide_hi = [](unsigned int v) { return __uint_as_float(v & 0xffff0000u); }; + struct ConvCursor { + const uint2 *w; unsigned int *s; + const unsigned int *z; const unsigned int *b; + }; + auto conv_cursor = [&](int c) { + const unsigned int cu = (unsigned int)c; + return ConvCursor{ + reinterpret_cast( + L.conv_w + cu * (unsigned int)MB_CONV_K), + reinterpret_cast( + cslot + cu * (unsigned int)(MB_CONV_K - 1)), + reinterpret_cast(zrow + cu), + reinterpret_cast(L.conv_b + cu)}; + }; + auto conv_two = [&](uint2 w0, uint2 w1, unsigned int s0, + unsigned int s1, unsigned int s2, + unsigned int zp, unsigned int bp, float &a0, + float &a1, unsigned int &n0, unsigned int &n1, + unsigned int &n2) { + a0 = wide_lo(bp); + a0 += wide_lo(w0.x) * wide_lo(s0); + a0 += wide_hi(w0.x) * wide_hi(s0); + a0 += wide_lo(w0.y) * wide_lo(s1); + a0 += wide_hi(w0.y) * wide_lo(zp); + a0 = a0 / (1.f + __expf(-a0)); + a1 = wide_hi(bp); + a1 += wide_lo(w1.x) * wide_hi(s1); + a1 += wide_hi(w1.x) * wide_lo(s2); + a1 += wide_lo(w1.y) * wide_hi(s2); + a1 += wide_hi(w1.y) * wide_hi(zp); + a1 = a1 / (1.f + __expf(-a1)); + n0 = (s0 >> 16) | (s1 << 16); + n1 = (zp & 0xffffu) | (s2 << 16); + n2 = (s2 >> 16) | (zp & 0xffff0000u); + }; + struct ConvPairIn { + uint2 w0, w1; + unsigned int s0, s1, s2, zp, bp; + }; + auto conv_pair_fetch = [](const ConvCursor &q, int r) { + const uint2 *const wp = q.w + 64 * r; + const unsigned int *const sp = q.s + 96 * r; + return ConvPairIn{wp[0], wp[1], sp[0], sp[1], sp[2], + q.z[32 * r], q.b[32 * r]}; + }; + auto conv_pair_apply = [&](unsigned int *const sp, + const ConvPairIn &v, float &a0, + float &a1) { + unsigned int n0, n1, n2; + conv_two(v.w0, v.w1, v.s0, v.s1, v.s2, v.zp, v.bp, a0, a1, + n0, n1, n2); + sp[0] = n0; + sp[1] = n1; + sp[2] = n2; + }; + constexpr int QUAD_SPAN = 4 * 32; + static_assert(MB_INNER % 4 == 0 && MB_STATE % 4 == 0 && + MB_GROUP_SIZE % 4 == 0, + "every quad base is a four-channel boundary"); + static_assert(MB_STATE == QUAD_SPAN && + MB_GROUP_SIZE == QUAD_SPAN + 2 * 32, + "one quad round covers B and C; x adds a pair round"); + struct ConvQuad { + const uint4 *w; uint2 *s; const uint2 *z; const uint2 *b; + }; + auto conv_quad_cursor = [&](int c) { + const unsigned int cu = (unsigned int)c; + return ConvQuad{ + reinterpret_cast( + L.conv_w + cu * (unsigned int)MB_CONV_K), + reinterpret_cast( + cslot + cu * (unsigned int)(MB_CONV_K - 1)), + reinterpret_cast(zrow + cu), + reinterpret_cast(L.conv_b + cu)}; + }; + struct ConvQuadIn { + uint4 w01, w23; + uint2 sa, sb, sc, zp, bp; + }; + auto conv_quad_fetch = [](const ConvQuad &q) { + const uint2 *const sp = q.s; + return ConvQuadIn{q.w[0], q.w[1], sp[0], sp[1], sp[2], + q.z[0], q.b[0]}; + }; + auto conv_quad_apply = [&](uint2 *const sp, const ConvQuadIn &v, + float &a0, float &a1, float &a2, + float &a3) { + unsigned int n0, n1, n2, n3, n4, n5; + conv_two(make_uint2(v.w01.x, v.w01.y), + make_uint2(v.w01.z, v.w01.w), + v.sa.x, v.sa.y, v.sb.x, v.zp.x, v.bp.x, + a0, a1, n0, n1, n2); + conv_two(make_uint2(v.w23.x, v.w23.y), + make_uint2(v.w23.z, v.w23.w), + v.sb.y, v.sc.x, v.sc.y, v.zp.y, v.bp.y, + a2, a3, n3, n4, n5); + sp[0] = make_uint2(n0, n1); + sp[1] = make_uint2(n2, n3); + sp[2] = make_uint2(n4, n5); + }; + auto conv_pack = [](float a, float b) { + return (unsigned int)__bfloat16_as_ushort(__float2bfloat16(a)) | + ((unsigned int)__bfloat16_as_ushort(__float2bfloat16(b)) << 16); + }; + const ConvQuad qb = conv_quad_cursor( + MB_INNER + group * MB_STATE + 4 * lane); + const ConvQuad qc = conv_quad_cursor( + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE + 4 * lane); + static_assert(MB_GROUP_SIZE * (int)sizeof(bf16) % 8 == 0, + "each warp's staging row starts 8-byte aligned"); + const ConvQuad qx = conv_quad_cursor( + group * MB_GROUP_SIZE + 4 * lane); + const ConvCursor qt = conv_cursor( + group * MB_GROUP_SIZE + QUAD_SPAN + 2 * lane); + const ConvQuadIn vb = conv_quad_fetch(qb); + const ConvQuadIn vc = conv_quad_fetch(qc); + const ConvQuadIn vx = conv_quad_fetch(qx); + const ConvPairIn vt = conv_pair_fetch(qt, 0); + uint2 bpk, cpk; + { + float a0, a1, a2, a3; + conv_quad_apply(qb.s, vb, a0, a1, a2, a3); + bpk = make_uint2(conv_pack(a0, a1), conv_pack(a2, a3)); + conv_quad_apply(qc.s, vc, a0, a1, a2, a3); + cpk = make_uint2(conv_pack(a0, a1), conv_pack(a2, a3)); + } + { + float a0, a1, a2, a3; + conv_quad_apply(qx.s, vx, a0, a1, a2, a3); + *reinterpret_cast(&conv_x[warp][4 * lane]) = + make_uint2(conv_pack(a0, a1), conv_pack(a2, a3)); + } + { + float a0, a1; + conv_pair_apply(qt.s, vt, a0, a1); + conv_x[warp][QUAD_SPAN + 2 * lane] = __float2bfloat16(a0); + conv_x[warp][QUAD_SPAN + 2 * lane + 1] = __float2bfloat16(a1); + } + const bool qword = ((lane >> 1) & 1) != 0; + const bool qhalf = (lane & 1) != 0; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int src = (lane >> 2) + 8 * j; + const unsigned int bx = __shfl_sync(0xffffffffu, bpk.x, src); + const unsigned int by = __shfl_sync(0xffffffffu, bpk.y, src); + const unsigned int cx = __shfl_sync(0xffffffffu, cpk.x, src); + const unsigned int cy = __shfl_sync(0xffffffffu, cpk.y, src); + const unsigned int bg = qword ? by : bx; + const unsigned int cg = qword ? cy : cx; + Braw[j] = qhalf ? wide_hi(bg) : wide_lo(bg); + Cv[j] = qhalf ? wide_hi(cg) : wide_lo(cg); + } + __syncwarp(); + xrow = &conv_x[warp][0]; + } else { + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Braw[j] = __bfloat162float(Bsec[lane + j * 32]); + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + xrow = xa + (size_t)group * MB_GROUP_SIZE; + } + const int head_lane = group * MB_HEADS_PER_GROUP + + (lane < MB_HEADS_PER_GROUP ? lane : 0); + const float dt_lane = mamba_softplus( + __bfloat162float(dtrow[head_lane]) + L.dt_bias[head_lane]); + const float dA_lane = __expf(L.A[head_lane] * dt_lane); + const float D_lane = L.D[head_lane]; + #pragma unroll 1 + for (int head_in_group = 0; head_in_group < MB_HEADS_PER_GROUP; + ++head_in_group) { + const int head = group * MB_HEADS_PER_GROUP + head_in_group; + const float dt = __shfl_sync(0xffffffffu, dt_lane, head_in_group); + const float dA = __shfl_sync(0xffffffffu, dA_lane, head_in_group); + const float Dh = __shfl_sync(0xffffffffu, D_lane, head_in_group); + const float *st_in = L.r8_initial_ssm + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float Bv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) + Bv[j] = Braw[j] * dt; + constexpr int NSTATE = SSM_NSTATE; + constexpr int SSM_PREFETCH = SSM_PF; + static_assert(channels % SSM_PREFETCH == 0); + for (int g0 = 0; g0 < channels; g0 += SSM_PREFETCH) { + float sv[SSM_PREFETCH][NSTATE]; + float xvr[SSM_PREFETCH]; + if (r8_initial && SSM_F32_OK) { + asm volatile("cp.async.wait_group 1;\n" ::: "memory"); + __syncwarp(); + const int pair = head_in_group * SSM_GROUPS_PER_HEAD + + g0 / SSM_PREFETCH; + const int buffer = pair & (SSM_F32_RING - 1); + const float *const smf = reinterpret_cast( + warp_ssm + (size_t)buffer * SSM_F32_GROUP_BYTES); + static_assert(SSM_PREFETCH % 2 == 0 && MB_HEAD_DIM % 2 == 0, + "x staging vectorises in element pairs"); + const __nv_bfloat162 *const xq = + reinterpret_cast( + xrow + head_in_group * MB_HEAD_DIM + + first_channel + g0); + #pragma unroll + for (int gq = 0; gq < SSM_PREFETCH / 2; ++gq) { + const float2 xp = __bfloat1622float2(xq[gq]); + xvr[2 * gq] = xp.x; + xvr[2 * gq + 1] = xp.y; + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = smf[(size_t)gi * MB_STATE + lane + j * 32]; + } + __syncwarp(); + const int ahead = pair + SSM_F32_RING; + if (ahead < SSM_PAIRS) + issue_ssm_f32(ahead / SSM_GROUPS_PER_HEAD, + (ahead % SSM_GROUPS_PER_HEAD) * SSM_PF, + buffer); + else + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } else if (r8_initial) { + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const size_t prow = (size_t)(first_channel + g0 + gi) * MB_STATE; + xvr[gi] = __bfloat162float( + xrow[head_in_group * MB_HEAD_DIM + + first_channel + g0 + gi]); + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = st_in[prow + lane + j * 32]; + } + } else if constexpr (SSM_SMEM_OK) { + asm volatile("cp.async.wait_group 3;\n" ::: "memory"); + __syncwarp(); + const int pair = head_in_group * SSM_GROUPS_PER_HEAD + + g0 / SSM_PREFETCH; + const int buffer = pair & (SSM_RING - 1); + const r9_state_t *sm = reinterpret_cast( + warp_ssm + (size_t)buffer * SSM_GROUP_BYTES); + static_assert(SSM_PREFETCH % 2 == 0 && MB_HEAD_DIM % 2 == 0, + "x staging vectorises in element pairs"); + const __nv_bfloat162 *const xq = + reinterpret_cast( + xrow + head_in_group * MB_HEAD_DIM + + first_channel + g0); + #pragma unroll + for (int gq = 0; gq < SSM_PREFETCH / 2; ++gq) { + const float2 xp = __bfloat1622float2(xq[gq]); + xvr[2 * gq] = xp.x; + xvr[2 * gq + 1] = xp.y; + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = r9_load( + sm[(size_t)gi * MB_STATE + lane + j * 32], 1.f); + } + __syncwarp(); + const int ahead = pair + SSM_RING; + if (ahead < SSM_PAIRS) + issue_ssm(ahead / SSM_GROUPS_PER_HEAD, + (ahead % SSM_GROUPS_PER_HEAD) * SSM_PF, buffer); + else + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } else { + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + xvr[gi] = __bfloat162float( + xrow[head_in_group * MB_HEAD_DIM + + first_channel + g0 + gi]); + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + { + int p = first_channel+g0+gi; + int mapped = R9_STATE_MODE==4 ? L.r9_map[head*24+p] : head*24+p; + int n = lane+j*32; + float scale = R9_STATE_MODE==1 ? 1.f : L.r9_scales[r9_scale_index(mapped/24,mapped%24,n)]; + sv[gi][j] = r9_load(st[(size_t)mapped*128+n],scale); + } + } + } + if constexpr (SSM_PREFETCH == 12) { + float v6[6]; + #pragma unroll + for (int k = 0; k < 6; ++k) { + float pair_acc[2]; + #pragma unroll + for (int half = 0; half < 2; ++half) { + const int gi = k + 6 * half; + const int p = first_channel + g0 + gi; + const float xv = xvr[gi]; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < NSTATE; ++j) { + const int n = lane + j * 32; + const float s = sv[gi][j] * dA + Bv[j] * xv; + const int mapped = R9_STATE_MODE == 4 + ? L.r9_map[head * 24 + p] : head * 24 + p; + const float scale = R9_STATE_MODE == 1 ? 1.f + : L.r9_scales[r9_scale_index( + mapped / 24, mapped % 24, n)]; + st[(size_t)mapped * 128 + n] = r9_store(s, scale); + acc += s * Cv[j]; + } + pair_acc[half] = acc; + } + const float mine = sb4 ? pair_acc[1] : pair_acc[0]; + const float away = sb4 ? pair_acc[0] : pair_acc[1]; + v6[k] = mine + __shfl_xor_sync(0xffffffff, away, 16); + } + float v4[4]; + #pragma unroll + for (int k = 0; k < 3; ++k) { + const float mine = sb3 ? v6[3 + k] : v6[k]; + const float away = sb3 ? v6[k] : v6[3 + k]; + v4[k] = mine + __shfl_xor_sync(0xffffffff, away, 8); + } + v4[3] = 0.f; + float v2[2]; + #pragma unroll + for (int k = 0; k < 2; ++k) { + const float mine = sb2 ? v4[2 + k] : v4[k]; + const float away = sb2 ? v4[k] : v4[2 + k]; + v2[k] = mine + __shfl_xor_sync(0xffffffff, away, 4); + } + float acc = (sb1 ? v2[1] : v2[0]) + + __shfl_xor_sync(0xffffffff, sb1 ? v2[0] : v2[1], 2); + acc += __shfl_xor_sync(0xffffffff, acc, 1); + if (ssm_owner) { + const int p = first_channel + g0 + ssm_own_ch; + const float xv = __bfloat162float( + xrow[head_in_group * MB_HEAD_DIM + p]); + const bf16 outv = __float2bfloat16(acc + xv * Dh); + if constexpr (FUSE_GNORM) + gn_stage[warp][head_in_group * MB_HEAD_DIM + p] = outv; + else + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = outv; + } + } else { + float cacc[SSM_PREFETCH]; + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const int p = first_channel + g0 + gi; + const float xv = xvr[gi]; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < NSTATE; ++j) { + const int n = lane + j * 32; + const float s = sv[gi][j] * dA + Bv[j] * xv; + const int mapped = R9_STATE_MODE == 4 + ? L.r9_map[head * 24 + p] : head * 24 + p; + const float scale = R9_STATE_MODE == 1 ? 1.f + : L.r9_scales[r9_scale_index( + mapped / 24, mapped % 24, n)]; + st[(size_t)mapped * 128 + n] = r9_store(s, scale); + acc += s * Cv[j]; + } + cacc[gi] = acc; + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const int p = first_channel + g0 + gi; + float acc = cacc[gi]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) { + const bf16 outv = __float2bfloat16(acc + xvr[gi] * Dh); + if constexpr (FUSE_GNORM) + gn_stage[warp][head_in_group * MB_HEAD_DIM + p] = outv; + else + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = outv; + } + } + } + } + } + if constexpr (SSM_SMEM_OK) { + if (!r8_initial || SSM_F32_OK) { + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); + __syncwarp(); + } + } + if constexpr (FUSE_GNORM) { + __syncwarp(); + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32], wg[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float gv = __bfloat162float(gn_stage[warp][lane + j * 32]) * + (z / (1.f + __expf(-z))); + wg[j] = __bfloat162float(L.gnorm_w[idx]); + held[j] = gv; + sq += gv * gv; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16(held[j] * inv * wg[j]); + } + __syncwarp(); + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_CONV_DIM; + const int c = idx - t * MB_CONV_DIM; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v; +#ifdef PF_CONV_ZERO_STATE + if (src < 0) v = 0.f; +#else + if (src < 0) v = __bfloat162float(state[pos + k]); +#endif + else v = __bfloat162float( + zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]); + acc += __bfloat162float(w[k]) * v; + } +#ifndef PF_CONV_NO_SILU + acc = acc / (1.f + __expf(-acc)); +#endif + xbc_act[(size_t)t * MB_CONV_DIM + c] = __float2bfloat16(acc); + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + const int units = n_seq * MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int seq = unit / (MB_HEADS * MB_HEAD_DIM); + const int rem = unit - seq * (MB_HEADS * MB_HEAD_DIM); + const int head = rem / MB_HEAD_DIM; + const int p = rem - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p * MB_STATE; + float sreg[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) sreg[j] = st[lane + j * 32]; + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const bf16 *xa = xbc_act + t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + const bf16 *dtrow = zxbcdt + t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + sreg[j] = sreg[j] * dA + __bfloat162float(Bsec[n]) * dt * xv; + acc += sreg[j] * __bfloat162float(Csec[n]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) st[lane + j * 32] = sreg[j]; + } + return; + } +#if PF_SSM_SPLIT_P + const int units = MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; unit += warps_total) { + const int head = unit / MB_HEAD_DIM; + const int p_only = unit - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; +#else + for (int head = blockIdx.x * WARPS + warp; head < MB_HEADS; + head += warps_total) { + const int group = head / MB_HEADS_PER_GROUP; +#endif + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *st = L.ssm_state + (size_t)slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE +#if PF_SSM_SPLIT_P + + (size_t)p_only * MB_STATE +#endif + ; + for (int t = 0; t < n_tokens; ++t) { + const bf16 *xa = xbc_act + (size_t)t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = + zxbcdt + (size_t)t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } +#if PF_SSM_SPLIT_P + { + const int p = p_only; +#else + for (int p = 0; p < MB_HEAD_DIM; ++p) { +#endif + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { +#if PF_SSM_SPLIT_P + const size_t at = (size_t)lane + j * 32; +#else + const size_t at = (size_t)p * MB_STATE + lane + j * 32; +#endif + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } + } +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32], wg[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + wg[j] = __bfloat162float(L.gnorm_w[idx]); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16(held[j] * inv * wg[j]); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len MK2_PROGRESS_PARAMETER(progress_mask)) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + constexpr int Q_PER_WARP = 3; + constexpr int GROUPS_PER_KV = 1; + const int pairs = BATCH * BB_KV_HEADS * GROUPS_PER_KV; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x + warp * gridDim.x; pair < pairs; + pair += gridDim.x * WARPS) { + const int row = pair / (BB_KV_HEADS * GROUPS_PER_KV); + MK2_PROGRESS_GUARD(progress_mask, row); + const int within_row = pair % (BB_KV_HEADS * GROUPS_PER_KV); + const int kv = within_row / GROUPS_PER_KV; + const int q_group = within_row % GROUPS_PER_KV; + const int q_first = q_group * Q_PER_WARP; + const int q_count = q_group == 0 ? Q_PER_WARP : 1; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (q_group == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[h][j] = 0.f; + } + float m[Q_PER_WARP], denom[Q_PER_WARP]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + m[h] = -CUDART_INF_F; + denom[h] = 0.f; + } + constexpr int PF = 8; + const int tpb = L.tokens_per_block; + const size_t elems_per_block = + (size_t)tpb * BB_KV_HEADS * BB_HEAD_DIM; + const int *bt_k = L.block_table + (size_t)(row * 2) * L.max_blocks_per_seq; + const int *bt_v = bt_k + L.max_blocks_per_seq; + const int head_base = kv * tpb * BB_HEAD_DIM; + float qr[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const bf16 *q = base + (kv * kv_per_q + q_first + h) * BB_HEAD_DIM; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qr[h][j] = __bfloat162float(q[lane + j * 32]); + } + for (int page0 = 0; page0 < T; page0 += tpb) { + const int k_page = bt_k[page0 / tpb] & 0x7FFFFFFF; + const int v_page = bt_v[page0 / tpb] & 0x7FFFFFFF; + const bf16 *kbase = L.pool_base + + (size_t)k_page * elems_per_block + + head_base; + const bf16 *vbase = L.pool_base + + (size_t)v_page * elems_per_block + + head_base; + const int tend = page0 + tpb < T ? page0 + tpb : T; + for (int t0 = page0; t0 < tend; t0 += PF) { + const int navail = tend - t0 < PF ? tend - t0 : PF; + float kr[PF][PER_LANE], vr[PF][PER_LANE]; + #pragma unroll + for (int u = 0; u < PF; ++u) { + const int tpos = t0 + (u < navail ? u : navail - 1); +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = vbase + (tpos - page0) * BB_HEAD_DIM; +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new : kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = cur_token + ? v_new : vbase + (tpos - page0) * BB_HEAD_DIM; +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + kr[u][j] = __bfloat162float(kt[i]); + vr[u][j] = __bfloat162float(vt[i]); + } + } + #pragma unroll + for (int u = 0; u < PF; ++u) { + if (u >= navail) break; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + dot += qr[h][j] * kr[u][j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m[h], s); + const float rescale = __expf(m[h] - m_new); + const float w = __expf(s - m_new); + denom[h] = denom[h] * rescale + w; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[h][j] = acc[h][j] * rescale + w * vr[u][j]; + m[h] = m_new; + } + } + } + } + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const float inv = 1.f / denom[h]; + const int head = kv * kv_per_q + q_first + h; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[h][j] * inv); + } + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * BB_KV_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / BB_KV_HEADS; + const int kv = idx - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv); + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv); + for (int i = 0; i < BB_HEAD_DIM; ++i) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + for (int s = 0; s <= pos; ++s) { + const bf16 *kt = kv_slot(L, seq, 0, s) + kv_offset(L, s, kv); + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + const bf16 *vt = kv_slot(L, seq, 1, s) + kv_offset(L, s, kv); + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + wt * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH, bf16 *row_stage = nullptr) { + __shared__ float score[BB_EXPERTS]; + constexpr int CH = D_MODEL / 8; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const char *gsrc = reinterpret_cast( + normed + (size_t)row * D_MODEL); + char *gdst = reinterpret_cast(row_stage); + const uint32_t act_s = (uint32_t)__cvta_generic_to_shared(gdst); + for (int i = threadIdx.x; i < CH; i += THREADS) { + const uint32_t dst = act_s + (uint32_t)i * 16; + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: + "r"(dst), "l"(gsrc + (size_t)i * 16) : "memory"); + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + if ((int)threadIdx.x < BB_EXPERTS) + moe_w[(size_t)row * BB_EXPERTS + threadIdx.x] = 0.f; + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); + __syncthreads(); + float acc0 = 0.f, acc1 = 0.f, acc2 = 0.f; + const bf16 *w0 = gate_w + (size_t)warp * D_MODEL; + const bf16 *w1 = gate_w + (size_t)(warp + WARPS) * D_MODEL; + const bf16 *w2 = gate_w + (size_t)(warp + 2 * WARPS) * D_MODEL; + for (int i = lane; i < D_MODEL; i += 128) { + bf16 sx0, sx1; + move::lds(sx0, act_s + (uint32_t)i * sizeof(bf16)); + move::lds(sx1, act_s + (uint32_t)(i + 32) * sizeof(bf16)); + float x0 = __bfloat162float(sx0); + float x1 = __bfloat162float(sx1); + acc0 += x0 * __bfloat162float(w0[i]); + acc1 += x0 * __bfloat162float(w1[i]); + acc2 += x0 * __bfloat162float(w2[i]); + acc0 += x1 * __bfloat162float(w0[i + 32]); + acc1 += x1 * __bfloat162float(w1[i + 32]); + acc2 += x1 * __bfloat162float(w2[i + 32]); + move::lds(sx0, act_s + (uint32_t)(i + 64) * sizeof(bf16)); + move::lds(sx1, act_s + (uint32_t)(i + 96) * sizeof(bf16)); + x0 = __bfloat162float(sx0); + x1 = __bfloat162float(sx1); + acc0 += x0 * __bfloat162float(w0[i + 64]); + acc1 += x0 * __bfloat162float(w1[i + 64]); + acc2 += x0 * __bfloat162float(w2[i + 64]); + acc0 += x1 * __bfloat162float(w0[i + 96]); + acc1 += x1 * __bfloat162float(w1[i + 96]); + acc2 += x1 * __bfloat162float(w2[i + 96]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + acc0 += __shfl_xor_sync(0xffffffffu, acc0, off); + acc1 += __shfl_xor_sync(0xffffffffu, acc1, off); + acc2 += __shfl_xor_sync(0xffffffffu, acc2, off); + } + if (lane == 0) { + score[warp] = 1.f / (1.f + __expf(-acc0)); + score[warp + WARPS] = 1.f / (1.f + __expf(-acc1)); + score[warp + 2 * WARPS] = 1.f / (1.f + __expf(-acc2)); + } + __syncthreads(); + if (threadIdx.x < 32) { + float remaining = (lane < BB_EXPERTS) + ? score[lane] + __bfloat162float(gate_bias[lane]) + : -CUDART_INF_F; + int pick[BB_TOP_K]; + float total = 0.f; + #pragma unroll + for (int k = 0; k < BB_TOP_K; ++k) { + float bv = remaining; + int bi = lane; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffff, bv, off); + const int oi = __shfl_xor_sync(0xffffffff, bi, off); + if (ov > bv || (ov == bv && oi < bi)) { bv = ov; bi = oi; } + } + pick[k] = bi; + total += score[bi]; + if (lane == bi) remaining = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + #pragma unroll + for (int k = 0; k < BB_TOP_K; ++k) + if (lane == pick[k]) moe_w[row * BB_EXPERTS + lane] = score[lane] * norm; + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert <= BB_EXPERTS; ++expert) + value = __fadd_rn(value, parts[base + (size_t)expert * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ bf16 bb_moe_fold_selected( + const float *parts, const float *weights, int idx) { + const int row = idx / D_MODEL; + const size_t base = (size_t)row * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert < BB_EXPERTS; ++expert) { + const float part = weights[row * BB_EXPERTS + expert] != 0.f + ? parts[base + (size_t)expert * D_MODEL] : 0.f; + value = __fadd_rn(value, part); + } + value = __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ __forceinline__ void sample_argmax_merge( + float &value, int &rank, float other, int other_rank) { + if (other > value || (other == value && other_rank < rank)) { + value = other; + rank = other_rank; + } +} +__device__ __forceinline__ int phoneme_warp_argmax(const bf16 *lg) { + const int lane = threadIdx.x & 31; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = lane; n < BB_PHONEME; n += 32) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, best, off); + const int other_arg = __shfl_down_sync(0xffffffffu, arg, off); + sample_argmax_merge(best, arg, other, other_arg); + } + return arg; +} +__device__ __forceinline__ int sample_block_argmax(float value, int rank) { + __shared__ float warp_value[WARPS]; + __shared__ int warp_rank[WARPS]; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + if (lane == 0) { warp_value[warp] = value; warp_rank[warp] = rank; } + __syncthreads(); + if (warp == 0) { + value = lane < WARPS ? warp_value[lane] : -CUDART_INF_F; + rank = lane < WARPS ? warp_rank[lane] : INT_MAX; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + } + return rank; +} +__device__ __forceinline__ int sample_eos_insertion_rank( + const uint64_t *keys, float eos) { + int lo = 0, hi = TOP_K; + while (lo < hi) { + const int mid = (lo + hi) >> 1; + if (eos > ordered_to_float((uint32_t)(keys[mid] >> 32))) hi = mid; + else lo = mid + 1; + } + return lo; +} +__device__ __forceinline__ void sample_topk_serial( + const uint64_t *keys, float eos, const float *noise, int *out) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float value; + int code; + if (!eos_used && eos > av) { + value = eos; + code = AUDIO_EOS; + eos_used = true; + } else { + value = av; + code = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = value + noise[j]; + if (noised > best_noised) { + best_noised = noised; + best_code = code; + } + } + *out = best_code; +} +__device__ __forceinline__ void sample_topk_parallel( + const uint64_t *keys, float eos, const float *noise, int *out) { + const int j = threadIdx.x; + const bool sorted_nan = j < TOP_K && + isnan(ordered_to_float((uint32_t)(keys[j] >> 32))); + if (__syncthreads_or(sorted_nan)) { + if (threadIdx.x == 0) sample_topk_serial(keys, eos, noise, out); + return; + } + const int lane = threadIdx.x & 31; + int insertion = lane == 0 ? sample_eos_insertion_rank(keys, eos) : 0; + insertion = __shfl_sync(0xffffffffu, insertion, 0); + float value = -CUDART_INF_F; + int rank = INT_MAX; + if (j < TOP_K) { + const int a = j - (j > insertion); + const float v = j == insertion + ? eos : ordered_to_float((uint32_t)(keys[a] >> 32)); + const float noised = v + noise[j]; + if (noised > -CUDART_INF_F) { value = noised; rank = j; } + } + const int winner = sample_block_argmax(value, rank); + if (threadIdx.x == 0) { + if (winner == INT_MAX) *out = 0; + else if (winner == insertion) *out = AUDIO_EOS; + else { + const int a = winner - (winner > insertion); + *out = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + } + } +} +constexpr int SORT_ITEMS_PER_THREAD = SORT_N / THREADS; +struct sampler_storage { + uint64_t keys[2][SORT_N]; +}; +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + sampler_storage &storage, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + uint64_t thread_keys[SORT_ITEMS_PER_THREAD]; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + thread_keys[item] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + int cross_stage = 0; + for (int k = 2; k <= SORT_N; k <<= 1) { + for (int j = k >> 1; j > 0; j >>= 1) { + if (j >= 32) { + const int bank = cross_stage++ & 1; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + storage.keys[bank][n] = thread_keys[item]; + } + __syncthreads(); + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + const uint64_t other = storage.keys[bank][n ^ j]; + const bool take_hi = ((n & k) == 0) == ((n & j) == 0); + const uint64_t mine = thread_keys[item]; + thread_keys[item] = take_hi + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } else { + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + const uint64_t other = __shfl_xor_sync( + 0xffffffffu, thread_keys[item], j); + const bool take_hi = ((n & k) == 0) == ((n & j) == 0); + const uint64_t mine = thread_keys[item]; + thread_keys[item] = take_hi + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } + } + } + const int top_bank = cross_stage & 1; + if (threadIdx.x < TOP_K) + storage.keys[top_bank][threadIdx.x] = thread_keys[0]; + __syncthreads(); + sample_topk_parallel( + storage.keys[top_bank], eos_v, + gumbel + ((size_t)row * N_POS + pos) * TOP_K, + out_codes + row); + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void r6_sparse_map(const float* routes,int* map,int* counts) { + __shared__ int r7_warp_counts[THREADS/32]; + const int lane=threadIdx.x%32, warp=threadIdx.x/32; + const unsigned lower=(1u<>1; + for(int e=0;e(packed)[(size_t)e*BATCH*row_vectors+i]= + row<0?make_uint4(0,0,0,0):reinterpret_cast(source)[(size_t)row*row_vectors+i%row_vectors]; + } + } +} +__device__ void r6_tile(int work,int output_tiles,const int* lut,int &tile,int &tok, + int mtg) { + const int groups=output_tiles/mtg; + tok=lut[1+work/groups]; + const int e=tok/DECODE_TILES; + tile=e*output_tiles+(work%groups)*mtg; +} +__device__ __forceinline__ void mk2_group_tile( + int work, int grouped, int &tile, int &tok, int &here) { + here = work < grouped ? NM_MAX : 1; + const int unit = work < grouped ? work * NM_MAX : grouped * (NM_MAX - 1) + work; + tile = unit / DECODE_TILES; + tok = unit - tile * DECODE_TILES; +} +__device__ __forceinline__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&p_s)[STAGES][SLOTS], out_tile &o_s, + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + int &tile_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = DECODE_TILES, + int rows_valid = DECODE_ROWS, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const int *r6_lut = nullptr,const int *r6_map = nullptr) { + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + constexpr int QSLOT = MT + NM_MAX; + constexpr bool WIDE = MT > 1 && NM_MAX > 1 && STAGES * SLOTS >= 2 * QSLOT; + constexpr int PSTAGES = WIDE ? 2 : STAGES; + constexpr int PSLOT = WIDE ? QSLOT : SLOTS; + constexpr int ACCN = WIDE ? MT * NM_MAX : NM_MAX; + constexpr int HEX_MT = 3, HEX_NM = 2; + constexpr int HEX_SLOT = HEX_MT + HEX_NM; + auto epi_publish = [](out_tile &st, const rt_fl<16, BATCH_PAD> &src) { + using acc_t = rt_fl<16, BATCH_PAD>; + static_assert(acc_t::height == 1, "one row block per warp"); + static_assert(out_tile::rows == BATCH_PAD && out_tile::cols == OUT_ATOM, + "staging tile must be square for an in-place transpose"); + const int t = ::kittens::laneid(); + const uint32_t base = + static_cast(__cvta_generic_to_shared(&st.data[0])); + const int row0 = 8 * (t >> 4) + (t & 7); + const int col0 = 16 * warpgroup::warpid() + 8 * ((t >> 3) & 1); + #pragma unroll + for (int j = 0; j < acc_t::width; ++j) { + bf16_2 q0 = base_types::convertor::convert( + src.tiles[0][j].data[0]); + bf16_2 q1 = base_types::convertor::convert( + src.tiles[0][j].data[1]); + bf16_2 q2 = base_types::convertor::convert( + src.tiles[0][j].data[2]); + bf16_2 q3 = base_types::convertor::convert( + src.tiles[0][j].data[3]); + move::stsm4t( + out_tile::idx(base, {16 * j + row0, col0}), q0, q1, q2, q3); + } + }; + weight_tile *const pf = &p_s[0][0]; + const bool arena10 = + reinterpret_cast(&o_sf) == + reinterpret_cast(pf + STAGES * SLOTS); + const bool two_d = !r6_lut && out_tiles_per_expert == 0 + && n_tok == DECODE_TILES && DECODE_TILES % NM_MAX == 0; + int n_tiles = compact_tiles ? __popc(active_mask)*compact_tiles : n_out/OUT_ATOM; + constexpr int QTOK = DECODE_TILES / NM_MAX; + const int out_tiles = n_tiles; + const bool tall3 = WIDE && two_d && out_tiles == 24; + const bool split72 = WIDE && two_d && DECODE_TILES == 16 && out_tiles == 72; + const bool split81 = WIDE && two_d && DECODE_TILES == 16 && out_tiles == 81; + const bool split17 = WIDE && two_d && DECODE_TILES == 16 && out_tiles == 17; + const bool mixed40 = WIDE && two_d && DECODE_TILES == 16 && out_tiles == 40; + const int n_tiles_hex_all = (out_tiles / HEX_MT) * (DECODE_TILES / HEX_NM); + const bool hex = WIDE && two_d && arena10 && DECODE_TILES == 16 && DECODE_TILES % HEX_NM == 0 && + (out_tiles == 96 || out_tiles == 81 || out_tiles == 72 || + out_tiles == 40) && + (epi == EPI_STORE || epi == EPI_GELU) && !r6_lut && + out_tiles_per_expert == 0; + const int hex_tok = DECODE_TILES / HEX_NM; + const int hex_wide = (out_tiles == 72 || out_tiles == 81) + ? 128 : n_tiles_hex_all; + const int pslot = hex ? HEX_SLOT : PSLOT; + int q_units = 0, q_rem_base = 0; + int grouped = 0; + if constexpr (WIDE) { + if (two_d) { + q_units = (out_tiles / MT) * QTOK; + q_rem_base = out_tiles - out_tiles % MT; + } + } else if (two_d) { + const int g = (int)gridDim.x; + const int units = n_tiles * DECODE_TILES; + const int rounds = (units + g - 1) / g; + const int cap = units / NM_MAX; + grouped = g * (rounds / NM_MAX); + if (grouped > cap) grouped = cap; + } + const int r6_output_tiles = r6_lut ? n_out/BB_EXPERTS/OUT_ATOM : 0; + const int r6_mtg = (MT > 1 && r6_lut && r6_output_tiles % MT == 0) ? MT : 1; + const int r6_groups = r6_lut ? r6_output_tiles/r6_mtg : 0; + const int r6_rows = r6_lut ? (r6_lut[0] & 0xFFF) : 0; + const int r6_pairs = (r6_lut && WIDE && r6_mtg == MT) ? (r6_lut[0] >> 12) : 0; + const int r6_pair_items = r6_pairs * r6_groups; + const int r6_single_mtg = (r6_lut && WIDE && r6_output_tiles % 3 == 0) + ? 3 : r6_mtg; + const int r6_single_groups = r6_lut ? r6_output_tiles/r6_single_mtg : 0; + const int r6_single_items = r6_lut + ? (r6_rows - 2 * r6_pairs) * r6_single_groups : 0; + if(r6_lut) { + n_tiles=r6_pair_items+r6_single_items; + n_tok=1; + } else if(two_d) { + if constexpr (WIDE) { + if (hex) n_tiles = (out_tiles == 72) ? 256 + : (out_tiles == 81) ? 264 + : (out_tiles == 40) + ? n_tiles_hex_all + hex_tok + : n_tiles_hex_all; + else if (mixed40) n_tiles = 244; + else if (split72) n_tiles = 360; + else if (split81) n_tiles = 384; + else if (split17) n_tiles = 132; + else if (tall3) n_tiles = (out_tiles / 3) * DECODE_TILES; + else n_tiles = q_units + (out_tiles % MT) * QTOK; + } + else n_tiles = grouped + (n_tiles*DECODE_TILES - grouped*NM_MAX); + n_tok=1; } + const int r6_extra = r6_lut ? n_tiles % (int)gridDim.x : 0; + const int r6_skew = r6_lut + ? (r6_single_items < r6_extra ? r6_single_items : r6_extra) : 0; + const int iters = k_total / KCHUNK; + int tile_start = (int)blockIdx.x - tile_io % (int)gridDim.x; + if (tile_start < 0) tile_start += (int)gridDim.x; + tile_io += n_tiles; + int tiles_here = 0; + for (int t = tile_start; t < n_tiles; t += gridDim.x) ++tiles_here; + const int iter_base = iter_io; + iter_io += tiles_here * iters * n_tok; + const int my_tiles = (n_tiles - tile_start + gridDim.x - 1) / gridDim.x; + if (warpgroup::groupid() == 1) { + if (hex) { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) { + if (out_tiles == 81 && work >= hex_wide) { + const int j = work - hex_wide; + const bool pair = j >= 128; + const int nt = pair ? 80 : 48 + 2 * (j / hex_tok); + const int ntok = HEX_NM * (j % hex_tok); + if (pair) { + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + wait(drained[s], (iter / PSTAGES) & 1); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], sizeof(weight_tile) + + HEX_NM * sizeof(act_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot], w, {nt, i}, filled[s]); + #pragma unroll + for (int m = 0; m < HEX_NM; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + 1 + m], a, + {ntok + m, i}, filled[s]); + } + } + } else { + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + wait(drained[s], (iter / PSTAGES) & 1); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], 2 * sizeof(weight_tile) + + HEX_NM * sizeof(act_tile)); + #pragma unroll + for (int c = 0; c < 2; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {nt + c, i}, filled[s]); + #pragma unroll + for (int m = 0; m < HEX_NM; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + 2 + m], a, + {ntok + m, i}, filled[s]); + } + } + } + continue; + } + if (out_tiles == 40 && work >= hex_wide) { + const int ntok = HEX_NM * (work - hex_wide); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + wait(drained[s], (iter / PSTAGES) & 1); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], sizeof(weight_tile) + + HEX_NM * sizeof(act_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot], w, {out_tiles - 1, i}, filled[s]); + #pragma unroll + for (int m = 0; m < HEX_NM; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + 1 + m], a, + {ntok + m, i}, filled[s]); + } + } + continue; + } + if (work >= hex_wide) { + const int j = work - hex_wide; + const int nt = HEX_MT * ((hex_wide + (j >> 1)) / hex_tok); + const int ntok = HEX_NM * ((hex_wide + (j >> 1)) % hex_tok) + (j & 1); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + wait(drained[s], (iter / PSTAGES) & 1); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], HEX_MT * sizeof(weight_tile) + + sizeof(act_tile)); + #pragma unroll + for (int c = 0; c < HEX_MT; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {nt + c, i}, filled[s]); + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + HEX_MT], a, {ntok, i}, filled[s]); + } + } + continue; + } + const int tile_id = HEX_MT * (work / hex_tok); + const int tok = HEX_NM * (work % hex_tok); + rt_fl<16, BATCH_PAD> pacc[HEX_NM]; + #pragma unroll + for (int u = 0; u < HEX_NM; ++u) warp::zero(pacc[u]); + int pend = -1, pend_ph = 0; + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(drained[s], phase); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], HEX_MT * sizeof(weight_tile) + + HEX_NM * sizeof(act_tile)); + #pragma unroll + for (int c = 0; c < HEX_MT; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {tile_id + c, i}, filled[s]); + #pragma unroll + for (int m = 0; m < HEX_NM; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + HEX_MT + m], a, {tok + m, i}, filled[s]); + } + if (pend >= 0) { { + wait(filled[pend], pend_ph); + warpgroup::mma_ABt(pacc[0], pf[pend * pslot + HEX_MT - 1], + pf[pend * pslot + HEX_MT]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + pacc[1], pf[pend * pslot + HEX_MT - 1], + pf[pend * pslot + HEX_MT + 1]); + warpgroup::mma_async_wait(); + } } + pend = s; pend_ph = phase; + } + { + wait(filled[pend], pend_ph); + warpgroup::mma_ABt(pacc[0], pf[pend * pslot + HEX_MT - 1], + pf[pend * pslot + HEX_MT]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + pacc[1], pf[pend * pslot + HEX_MT - 1], + pf[pend * pslot + HEX_MT + 1]); + warpgroup::mma_async_wait(); + } + const int es = iter % PSTAGES; + wait(drained[es], (iter / PSTAGES) & 1); + out_tile_f &q_sf = + *reinterpret_cast(&pf[es * pslot]); + out_tile &q_s = *reinterpret_cast(&q_sf); + const int qbase = (tile_id + HEX_MT - 1) * OUT_ATOM; + constexpr int OUT_OCTS = OUT_ATOM / 8; + #pragma unroll + for (int u = 0; u < HEX_NM; ++u) { + const int grow0 = (tok + u) * BATCH_PAD; + epi_publish(q_s, pacc[u]); + warpgroup::sync(3); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); + idx < OUT_OCTS * rows_valid; idx += 128) { + const int row = idx >> 3, col = (idx & 7) << 3; + const size_t at = + (size_t)(grow0 + row) * dest_stride + qbase + col; + const bf16_oct q = *reinterpret_cast( + &q_s[int2{row, col}]); + if (epi == EPI_STORE) { + *reinterpret_cast(dest + at) = q; + } else { + *reinterpret_cast(dest + at) = bf16_oct{ + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x0))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x1))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x2))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x3))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x4))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x5))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x6))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x7)))}; + } + } + warpgroup::sync(3); + } + } + } else if (warpgroup::laneid() == 0) { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) { + int tile_id = mk2_expert_tile(work, active_mask, compact_tiles), r6_tok=0; + int nm = 1, mt = 1; + if(r6_lut) { + int rwork = work - r6_skew; + if (rwork < 0) rwork += n_tiles; + const bool pair = rwork < r6_pair_items; + const int q = pair ? rwork / r6_groups + : (rwork-r6_pair_items) / r6_single_groups; + const int group = pair ? rwork-q*r6_groups + : rwork-r6_pair_items-q*r6_single_groups; + const int base_item = pair ? 2*q : 2*r6_pairs+q; + const int ltok = r6_lut[1 + base_item]; + r6_tok = ltok; + tile_id = (ltok / DECODE_TILES) * r6_output_tiles + + group * (pair ? r6_mtg : r6_single_mtg); + mt = pair ? r6_mtg : r6_single_mtg; + nm = pair ? NM_MAX : 1; + } + else if(two_d) { + if constexpr (WIDE) { + if (mixed40) { + const bool quad = work < 132; + const int item = quad ? work : work - 132; + if (quad && item < 84) { + tile_id = 2 * (item / 7); + r6_tok = 2 * (item % 7); + } else if (quad) { + const int tail = item - 84; + tile_id = 2 * (12 + tail / 6); + r6_tok = 2 * (tail % 6); + } else if (item < 48) { + tile_id = 2 * (item / 4) + (item % 4) / 2; + r6_tok = 14 + item % 2; + } else { + const int tail = item - 48; + tile_id = 2 * (12 + tail / 8) + (tail % 8) / 4; + r6_tok = 12 + tail % 4; + } + mt = quad ? MT : 1; + nm = quad ? NM_MAX : 1; + } else if (split72) { + if (work < 264) { + tile_id = 2 * (work / 8); + r6_tok = 2 * (work % 8); + mt = MT; nm = NM_MAX; + } else { + const int i = work - 264, c = i >> 2, q = i & 3; + tile_id = 2 * (33 + (c >> 3)) + (q >> 1); + r6_tok = 2 * (c & 7) + (q & 1); + mt = 1; nm = 1; + } + } else if (split81) { + if (work < 264) { + tile_id = 2 * (work / 8); + r6_tok = 2 * (work % 8); + mt = MT; nm = NM_MAX; + } else if (work < 376) { + const int i = work - 264; + tile_id = 2 * (33 + i / 16); + r6_tok = i % 16; + mt = MT; nm = 1; + } else { + tile_id = 80; + r6_tok = 2 * (work - 376); + mt = 1; nm = NM_MAX; + } + } else if (split17) { + if (work < 8) { + tile_id = 14; + r6_tok = work; + mt = 3; nm = 1; + } else if (work < 120) { + const int i = work - 8; + tile_id = i >> 3; + r6_tok = 2 * (i & 7); + mt = 1; nm = NM_MAX; + } else { + const int i = work - 120; + tile_id = 14 + (i >> 2); + r6_tok = 8 + 2 * (i & 3); + mt = 1; nm = NM_MAX; + } + } else if (tall3) { + tile_id = 3 * (work / DECODE_TILES); + r6_tok = work % DECODE_TILES; + mt = 3; + nm = 1; + } else { + const int j = work < q_units ? work : work - q_units; + tile_id = work < q_units ? MT * (j / QTOK) + : q_rem_base + j / QTOK; + r6_tok = NM_MAX * (j % QTOK); + mt = work < q_units ? MT : 1; + nm = NM_MAX; + } + } else mk2_group_tile(work,grouped,tile_id,r6_tok,nm); + } + for (int token_iteration = 0; token_iteration < n_tok; ++token_iteration) + for (int i = 0; i < iters; ++i, ++iter) { + const int tok = (r6_lut || two_d) ? r6_tok : token_iteration; + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(drained[s], phase); + tma::expect_bytes(filled[s], + mt * sizeof(weight_tile) + nm * sizeof(act_tile)); + if (mt == 1) { + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot], w, {tile_id, i}, filled[s]); + } else if (mt == 3) { + #pragma unroll + for (int c = 0; c < 3; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {tile_id + c, i}, filled[s]); + } else { + #pragma unroll + for (int c = 0; c < MT; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {tile_id + c, i}, filled[s]); + } + const int a_col = out_tiles_per_expert && !r6_lut + ? (tile_id / out_tiles_per_expert) * act_k_tiles_per_expert + i + : i; + if (NM_MAX == 1 || nm == NM_MAX) { + #pragma unroll + for (int m = 0; m < NM_MAX; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + mt + m], a, {tok + m, a_col}, filled[s]); + } else { + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + mt], a, {tok, a_col}, filled[s]); + } + } + } + } + } else { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) + for (int token_iteration = 0; token_iteration < n_tok; ++token_iteration) { + int tile_id=mk2_expert_tile(work, active_mask, compact_tiles),tok=token_iteration; + int nm = 1, mt = 1; + if(r6_lut) { + int rwork = work - r6_skew; + if (rwork < 0) rwork += n_tiles; + const bool pair = rwork < r6_pair_items; + const int q = pair ? rwork / r6_groups + : (rwork-r6_pair_items) / r6_single_groups; + const int group = pair ? rwork-q*r6_groups + : rwork-r6_pair_items-q*r6_single_groups; + const int base_item = pair ? 2*q : 2*r6_pairs+q; + const int ltok = r6_lut[1 + base_item]; + tok = ltok; + tile_id = (ltok / DECODE_TILES) * r6_output_tiles + + group * (pair ? r6_mtg : r6_single_mtg); + mt = pair ? r6_mtg : r6_single_mtg; + nm = pair ? NM_MAX : 1; + } + else if(two_d) { + if constexpr (WIDE) { + if (hex) { + if (out_tiles == 81 && work >= hex_wide) { + const int j = work - hex_wide; + const bool pair = j >= 128; + tile_id = pair ? 80 : 48 + 2 * (j / hex_tok); + tok = HEX_NM * (j % hex_tok); + mt = pair ? 1 : 2; + nm = HEX_NM; + } else if (out_tiles == 40 && work >= hex_wide) { + tile_id = out_tiles - 1; + tok = HEX_NM * (work - hex_wide); + mt = 1; + nm = HEX_NM; + } else { + const int hw = work < hex_wide + ? work : hex_wide + ((work - hex_wide) >> 1); + tile_id = HEX_MT * (hw / hex_tok); + tok = HEX_NM * (hw % hex_tok) + + (work < hex_wide ? 0 : ((work - hex_wide) & 1)); + mt = HEX_MT; nm = work < hex_wide ? HEX_NM : 1; + } + } else if (mixed40) { + const bool quad = work < 132; + const int item = quad ? work : work - 132; + if (quad && item < 84) { + tile_id = 2 * (item / 7); + tok = 2 * (item % 7); + } else if (quad) { + const int tail = item - 84; + tile_id = 2 * (12 + tail / 6); + tok = 2 * (tail % 6); + } else if (item < 48) { + tile_id = 2 * (item / 4) + (item % 4) / 2; + tok = 14 + item % 2; + } else { + const int tail = item - 48; + tile_id = 2 * (12 + tail / 8) + (tail % 8) / 4; + tok = 12 + tail % 4; + } + mt = quad ? MT : 1; + nm = quad ? NM_MAX : 1; + } else if (split72) { + if (work < 264) { + tile_id = 2 * (work / 8); + tok = 2 * (work % 8); + mt = MT; nm = NM_MAX; + } else { + const int i = work - 264, c = i >> 2, q = i & 3; + tile_id = 2 * (33 + (c >> 3)) + (q >> 1); + tok = 2 * (c & 7) + (q & 1); + mt = 1; nm = 1; + } + } else if (split81) { + if (work < 264) { + tile_id = 2 * (work / 8); + tok = 2 * (work % 8); + mt = MT; nm = NM_MAX; + } else if (work < 376) { + const int i = work - 264; + tile_id = 2 * (33 + i / 16); + tok = i % 16; + mt = MT; nm = 1; + } else { + tile_id = 80; + tok = 2 * (work - 376); + mt = 1; nm = NM_MAX; + } + } else if (split17) { + if (work < 8) { + tile_id = 14; + tok = work; + mt = 3; nm = 1; + } else if (work < 120) { + const int i = work - 8; + tile_id = i >> 3; + tok = 2 * (i & 7); + mt = 1; nm = NM_MAX; + } else { + const int i = work - 120; + tile_id = 14 + (i >> 2); + tok = 8 + 2 * (i & 3); + mt = 1; nm = NM_MAX; + } + } else if (tall3) { + tile_id = 3 * (work / DECODE_TILES); + tok = work % DECODE_TILES; + mt = 3; + nm = 1; + } else { + const int j = work < q_units ? work : work - q_units; + tile_id = work < q_units ? MT * (j / QTOK) + : q_rem_base + j / QTOK; + tok = NM_MAX * (j % QTOK); + mt = work < q_units ? MT : 1; + nm = NM_MAX; + } + } else mk2_group_tile(work,grouped,tile_id,tok,nm); + } + int epi_hold = -1; + rt_fl<16, BATCH_PAD> accum[ACCN]; + if (hex && nm == 1) { + #pragma unroll + for (int u = 0; u < 3; ++u) warp::zero(accum[u]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + HEX_MT]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[1], pf[s * pslot + 1], pf[s * pslot + HEX_MT]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[2], pf[s * pslot + 2], pf[s * pslot + HEX_MT]); + warpgroup::mma_async_wait(); + if (i + 1 < iters) { + if (warpgroup::laneid() == 0) + arrive(drained[s], CONSUMER_WARPS); + } else epi_hold = s; + } + } else if (hex && mt == HEX_MT) { + #pragma unroll + for (int u = 0; u < ACCN; ++u) warp::zero(accum[u]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], + pf[s * pslot + HEX_MT]); + #pragma unroll + for (int u = 1; u < ACCN; ++u) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[u], pf[s * pslot + u / HEX_NM], + pf[s * pslot + HEX_MT + u % HEX_NM]); + warpgroup::mma_async_wait(); + if (i + 1 < iters) { + if (warpgroup::laneid() == 0) + arrive(drained[s], CONSUMER_WARPS); + } else epi_hold = s; + } + } else if (WIDE && mt == 3) { + #pragma unroll + for (int u = 0; u < 3; ++u) warp::zero(accum[u]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + 3]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[1], pf[s * pslot + 1], pf[s * pslot + 3]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[2], pf[s * pslot + 2], pf[s * pslot + 3]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else if (WIDE && mt == MT && nm == NM_MAX) { + #pragma unroll + for (int u = 0; u < ACCN; ++u) warp::zero(accum[u]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + MT]); + #pragma unroll + for (int u = 1; u < ACCN; ++u) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[u], pf[s * pslot + u / NM_MAX], + pf[s * pslot + MT + u % NM_MAX]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else if (MT > 1 && mt == MT) { + #pragma unroll + for (int c = 0; c < MT; ++c) warp::zero(accum[c]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + MT]); + #pragma unroll + for (int c = 1; c < MT; ++c) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[c], pf[s * pslot + c], pf[s * pslot + MT]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else if (NM_MAX == 1 || nm == NM_MAX) { + #pragma unroll + for (int m = 0; m < NM_MAX; ++m) warp::zero(accum[m]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + 1]); + #pragma unroll + for (int m = 1; m < NM_MAX; ++m) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[m], pf[s * pslot], pf[s * pslot + 1 + m]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else { + warp::zero(accum[0]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + 1]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } + out_tile_f &e_sf = epi_hold >= 0 + ? *reinterpret_cast(&pf[epi_hold * pslot]) + : o_sf; + out_tile *const e_sp = reinterpret_cast(&e_sf); + const int nacc = mt * nm; + if (epi != EPI_BIAS) { + epi_publish(e_sp[0], accum[0]); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + } + #pragma unroll + for (int u = 0; u < ACCN; ++u) { + if (u >= mt * nm) break; + const int c = (mt > 1) ? ((nm > 1) ? (u / NM_MAX) : u) : 0; + const int m = (nm > 1) ? (u % NM_MAX) : 0; + const int base = (r6_lut ? (tile_id + c)%r6_output_tiles + : (tile_id + c)) * OUT_ATOM; + const int grow0 = (tok + m) * BATCH_PAD; + if (epi == EPI_BIAS) { + const int bw = 16 * warpgroup::warpid() + + (::kittens::laneid() >> 2); + const float b0 = __bfloat162float(bias[base + bw]); + const float b1 = __bfloat162float(bias[base + bw + 8]); + #pragma unroll + for (int j = 0; j < rt_fl<16, BATCH_PAD>::width; ++j) { + accum[u].tiles[0][j].data[0].x += b0; + accum[u].tiles[0][j].data[0].y += b0; + accum[u].tiles[0][j].data[1].x += b1; + accum[u].tiles[0][j].data[1].y += b1; + accum[u].tiles[0][j].data[2].x += b0; + accum[u].tiles[0][j].data[2].y += b0; + accum[u].tiles[0][j].data[3].x += b1; + accum[u].tiles[0][j].data[3].y += b1; + } + epi_publish(e_sp[0], accum[u]); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + constexpr int OUT_OCTS = OUT_ATOM / 8; + for (int idx = warpgroup::laneid(); idx < OUT_OCTS * rows_valid; + idx += 128) { + const int row = idx >> 3, col = (idx & 7) << 3; + const size_t at = (size_t)(grow0 + row) * dest_stride + base + col; + *reinterpret_cast(dest + at) = + *reinterpret_cast( + &e_sp[0][int2{row, col}]); + } + warpgroup::sync(1); + continue; + } + if (u + 1 < ACCN && u + 1 < nacc) + epi_publish(e_sp[(u + 1) & 1], + accum[u + 1 < ACCN ? u + 1 : 0]); + const out_tile &e_sr = e_sp[u & 1]; + constexpr int OUT_OCTS = OUT_ATOM / 8; + for (int idx = warpgroup::laneid(); idx < OUT_OCTS * rows_valid; + idx += 128) { + const int row = idx >> 3; + const int col = (idx & 7) << 3; + const int grow = grow0 + row; + const size_t at = + (size_t)grow * dest_stride + base + col; + const bf16_oct src = *reinterpret_cast( + &e_sr[int2{row, col}]); + if (epi == EPI_STORE) { + *reinterpret_cast(dest + at) = src; + continue; + } + const float v0 = __bfloat162float(src.x0); + const float v1 = __bfloat162float(src.x1); + const float v2 = __bfloat162float(src.x2); + const float v3 = __bfloat162float(src.x3); + const float v4 = __bfloat162float(src.x4); + const float v5 = __bfloat162float(src.x5); + const float v6 = __bfloat162float(src.x6); + const float v7 = __bfloat162float(src.x7); + switch (epi) { + case EPI_STORE: break; + case EPI_SILU: + *reinterpret_cast(dest + at) = bf16_oct{ + __float2bfloat16(v0 / (1.f + __expf(-v0))), + __float2bfloat16(v1 / (1.f + __expf(-v1))), + __float2bfloat16(v2 / (1.f + __expf(-v2))), + __float2bfloat16(v3 / (1.f + __expf(-v3))), + __float2bfloat16(v4 / (1.f + __expf(-v4))), + __float2bfloat16(v5 / (1.f + __expf(-v5))), + __float2bfloat16(v6 / (1.f + __expf(-v6))), + __float2bfloat16(v7 / (1.f + __expf(-v7)))}; + break; + case EPI_MOE_ACC: { + const int expert = (tile_id + c) / out_tiles_per_expert; + const int local = ((tile_id + c) % out_tiles_per_expert) * OUT_ATOM + + col; + const int original_row=r6_map?r6_map[grow]:grow; + if(original_row<0) break; + const float wt = moe_w[original_row * BB_EXPERTS + expert]; + if (wt != 0.f) { + const size_t dst = + ((size_t)original_row * (BB_EXPERTS + 1) + expert) * + D_MODEL + local; + *reinterpret_cast(moe_out + dst) = make_float4( + __fmul_rn(wt, v0), __fmul_rn(wt, v1), + __fmul_rn(wt, v2), __fmul_rn(wt, v3)); + *reinterpret_cast(moe_out + dst + 4) = make_float4( + __fmul_rn(wt, v4), __fmul_rn(wt, v5), + __fmul_rn(wt, v6), __fmul_rn(wt, v7)); + } + break; + } + case EPI_ACC_F32: + *reinterpret_cast(moe_out + + ((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * + D_MODEL + base + col) = make_float4(v0, v1, v2, v3); + *reinterpret_cast(moe_out + + ((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * + D_MODEL + base + col + 4) = make_float4(v4, v5, v6, v7); + break; + case EPI_BIAS: break; + case EPI_GELU: + *reinterpret_cast(dest + at) = bf16_oct{ + __float2bfloat16(gelu_tanh_fp32(v0)), + __float2bfloat16(gelu_tanh_fp32(v1)), + __float2bfloat16(gelu_tanh_fp32(v2)), + __float2bfloat16(gelu_tanh_fp32(v3)), + __float2bfloat16(gelu_tanh_fp32(v4)), + __float2bfloat16(gelu_tanh_fp32(v5)), + __float2bfloat16(gelu_tanh_fp32(v6)), + __float2bfloat16(gelu_tanh_fp32(v7))}; + break; + case EPI_RESIDUAL: { + const bf16_oct prior = + *reinterpret_cast(residual + at); + *reinterpret_cast(residual + at) = bf16_oct{ + __float2bfloat16(__bfloat162float(prior.x0) + v0), + __float2bfloat16(__bfloat162float(prior.x1) + v1), + __float2bfloat16(__bfloat162float(prior.x2) + v2), + __float2bfloat16(__bfloat162float(prior.x3) + v3), + __float2bfloat16(__bfloat162float(prior.x4) + v4), + __float2bfloat16(__bfloat162float(prior.x5) + v5), + __float2bfloat16(__bfloat162float(prior.x6) + v6), + __float2bfloat16(__bfloat162float(prior.x7) + v7)}; + break; + } + } + } + if (u + 1 < nacc) { + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + } + } + if (epi != EPI_BIAS) warpgroup::sync(1); + if (epi_hold >= 0 && warpgroup::laneid() == 0) + arrive(drained[epi_hold], CONSUMER_WARPS); + } + (void)my_tiles; + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr, + const unsigned char *native_predictor_active_mask = nullptr + MK2_PROGRESS_PARAMETER(native_progress_mask)) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&p_s)[STAGES][SLOTS] = alloc.allocate(); + out_tile_f (&o_sf) = alloc.allocate(); + out_tile &o_s = *reinterpret_cast(&o_sf); + bf16 *const route_stage = reinterpret_cast(&o_sf); + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_SCRATCH_STRIDE]; + static_assert(STAGES * SLOTS * sizeof(weight_tile) >= sizeof(sampler_storage)); + sampler_storage &sort_storage = *reinterpret_cast(&__shm[0]); + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS, 0); + arrive(drained[s], CONSUMER_WARPS); + } + } + int gemm_iter = 0; + int gemm_tile = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + const float *fold_src = nullptr, *fold_routes = nullptr; + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + fold_src = nullptr; + fold_routes = nullptr; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + grid_barrier(g.barrier); +#endif + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); +#endif + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, BATCH, + route_stage); + } + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + } + grid_barrier(g.barrier); + r6_sparse_map(bb->moe_w,bb->r6_map,bb->r6_counts); + grid_barrier(g.barrier); + r6_sparse_gather(bb->normed,bb->r6_packed,bb->r6_map,bb->r6_counts, + bb->r6_lut); + grid_barrier(g.barrier); + gemm(L.exp_up, bb->r6_views[0], BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0,0,nullptr,nullptr,DECODE_TILES,DECODE_ROWS, + nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->r6_views[1], BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + DECODE_TILES, DECODE_ROWS, nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + if (bb->layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_selected(bb->moe_out, bb->moe_w, idx); + grid_barrier(g.barrier); + } else { + fold_src = bb->moe_out; + fold_routes = bb->moe_w; + } + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int n_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = n_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + bb->pf_h[idx] = bb->prompt[idx]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = bb->pf_n_seq > 0 ? bb->pf_n_seq : 1; + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_gated, + T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, T, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * (BB_EXPERTS + 1) * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->pf_normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->pf_moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = bb_moe_fold_deterministic(bb->pf_moe_out, idx); + grid_barrier(g.barrier); + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, bb->pf_final, ln_scratch, + n_total); + grid_barrier(g.barrier); + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, bb->phoneme_b); + grid_barrier(g.barrier); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + for (int row = blockIdx.x + warp * gridDim.x; row < BATCH; + row += gridDim.x * WARPS) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + MK2_PROGRESS_GUARD(native_progress_mask, row); + const int arg = phoneme_warp_argmax(lg); + if (lane == 0) bb->phoneme_out[row] = arg; + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + __shared__ int native_run_predictor; + if (native_predictor_active_mask && bb && g.native_codes) { + if (threadIdx.x == 0) { + int any_active = 0; + for (int row = 0; row < BATCH; ++row) + any_active |= int(native_predictor_active_mask[row] != 0 && + MK2_PROGRESS_ENABLED(native_progress_mask, row)); + native_run_predictor = any_active; + } + __syncthreads(); + } + const bool skip_native_predictor = native_predictor_active_mask && bb && + g.native_codes && native_run_predictor == 0; + if (skip_native_predictor) { + for (int index = blockIdx.x * THREADS + threadIdx.x; + index < BATCH * N_POS; index += gridDim.x * THREADS) { + MK2_PROGRESS_GUARD(native_progress_mask, index / N_POS); + g.sampled[index] = 0; + g.native_codes[index] = 0; + } + grid_barrier(g.barrier); + } else { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + constexpr int DPL = D_HEAD / 32; + float qv[DPL], kn[DPL], vn[DPL]; + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + qv[j] = __bfloat162float(q[i]); + kn[j] = __bfloat162float(k_new[i]); + vn[j] = __bfloat162float(v_new[i]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += qv[j] * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + { + float acc = 0.f; + #pragma unroll + for (int j = 0; j < DPL; ++j) acc += qv[j] * kn[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[pos] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + float ov[DPL] = {}; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *vt = g.vcache + cache_base + t * D_HEAD; + const float s = sc[t]; + #pragma unroll + for (int j = 0; j < DPL; ++j) + ov[j] += s * __bfloat162float(vt[lane + j * 32]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + ov[j] += sc[pos] * vn[j]; + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + + lane + j * 32] = __float2bfloat16(ov[j]); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, + g.sampled_pos, sort_storage, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif +#ifdef ABLATE_SAMPLE + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { + __syncthreads(); + } else { + grid_barrier(g.barrier); + } +#endif + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x == 0) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + g.sampled[row * N_POS + pos] = code; + if (g.native_codes && MK2_PROGRESS_ENABLED(native_progress_mask, row)) { +#if MK2_INFLIGHT + g.native_codes[row * N_POS + pos] = + native_predictor_active_mask && !native_predictor_active_mask[row] + ? 0 : (long long)code; +#else + g.native_codes[row * N_POS + pos] = (long long)code; +#endif + } + } + } +#ifdef ABLATE_SAMPLE + grid_barrier(g.barrier); +#else + if (pos + 1 == n_pos) grid_barrier(g.barrier); +#endif + } + } + if (g.native_seq_len && blockIdx.x == 0) + for (int row = threadIdx.x; row < BATCH; row += blockDim.x) { + MK2_PROGRESS_GUARD(native_progress_mask, row); + g.native_seq_len[row] += 1; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2_startup( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr, + const unsigned char *native_predictor_active_mask = nullptr + MK2_PROGRESS_PARAMETER(native_progress_mask)) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&p_s)[STAGES][SLOTS] = alloc.allocate(); + out_tile_f (&o_sf) = alloc.allocate(); + out_tile &o_s = *reinterpret_cast(&o_sf); + bf16 *const route_stage = reinterpret_cast(&o_sf); + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_SCRATCH_STRIDE]; + static_assert(STAGES * SLOTS * sizeof(weight_tile) >= sizeof(sampler_storage)); + sampler_storage &sort_storage = *reinterpret_cast(&__shm[0]); + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS, 0); + arrive(drained[s], CONSUMER_WARPS); + } + } + int gemm_iter = 0; + int gemm_tile = 0; + __syncthreads(); + constexpr bool prefilling = false; + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + const float *fold_src = nullptr, *fold_routes = nullptr; + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + fold_src = nullptr; + fold_routes = nullptr; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + grid_barrier(g.barrier); +#endif + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); +#endif + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, BATCH, + route_stage); + } + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + } + grid_barrier(g.barrier); + r6_sparse_map(bb->moe_w,bb->r6_map,bb->r6_counts); + grid_barrier(g.barrier); + r6_sparse_gather(bb->normed,bb->r6_packed,bb->r6_map,bb->r6_counts, + bb->r6_lut); + grid_barrier(g.barrier); + gemm(L.exp_up, bb->r6_views[0], BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0,0,nullptr,nullptr,DECODE_TILES,DECODE_ROWS, + nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->r6_views[1], BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + DECODE_TILES, DECODE_ROWS, nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + if (bb->layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_selected(bb->moe_out, bb->moe_w, idx); + grid_barrier(g.barrier); + } else { + fold_src = bb->moe_out; + fold_routes = bb->moe_w; + } + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if constexpr (false) { + const int n_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = n_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + bb->pf_h[idx] = bb->prompt[idx]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = bb->pf_n_seq > 0 ? bb->pf_n_seq : 1; + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_gated, + T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, T, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * (BB_EXPERTS + 1) * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->pf_normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->pf_moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = bb_moe_fold_deterministic(bb->pf_moe_out, idx); + grid_barrier(g.barrier); + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, bb->pf_final, ln_scratch, + n_total); + grid_barrier(g.barrier); + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, bb->phoneme_b); + grid_barrier(g.barrier); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + for (int row = blockIdx.x + warp * gridDim.x; row < BATCH; + row += gridDim.x * WARPS) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + MK2_PROGRESS_GUARD(native_progress_mask, row); + const int arg = phoneme_warp_argmax(lg); + if (lane == 0) bb->phoneme_out[row] = arg; + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + __shared__ int native_run_predictor; + if (native_predictor_active_mask && bb && g.native_codes) { + if (threadIdx.x == 0) { + int any_active = 0; + for (int row = 0; row < BATCH; ++row) + any_active |= int(native_predictor_active_mask[row] != 0 && + MK2_PROGRESS_ENABLED(native_progress_mask, row)); + native_run_predictor = any_active; + } + __syncthreads(); + } + const bool skip_native_predictor = native_predictor_active_mask && bb && + g.native_codes && native_run_predictor == 0; + if constexpr (true) { + for (int index = blockIdx.x * THREADS + threadIdx.x; + index < BATCH * N_POS; index += gridDim.x * THREADS) { + MK2_PROGRESS_GUARD(native_progress_mask, index / N_POS); + g.sampled[index] = 0; + g.native_codes[index] = 0; + } + grid_barrier(g.barrier); + } else { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + constexpr int DPL = D_HEAD / 32; + float qv[DPL], kn[DPL], vn[DPL]; + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + qv[j] = __bfloat162float(q[i]); + kn[j] = __bfloat162float(k_new[i]); + vn[j] = __bfloat162float(v_new[i]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += qv[j] * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + { + float acc = 0.f; + #pragma unroll + for (int j = 0; j < DPL; ++j) acc += qv[j] * kn[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[pos] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + float ov[DPL] = {}; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *vt = g.vcache + cache_base + t * D_HEAD; + const float s = sc[t]; + #pragma unroll + for (int j = 0; j < DPL; ++j) + ov[j] += s * __bfloat162float(vt[lane + j * 32]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + ov[j] += sc[pos] * vn[j]; + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + + lane + j * 32] = __float2bfloat16(ov[j]); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, + g.sampled_pos, sort_storage, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif +#ifdef ABLATE_SAMPLE + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { + __syncthreads(); + } else { + grid_barrier(g.barrier); + } +#endif + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x == 0) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + g.sampled[row * N_POS + pos] = code; + if (g.native_codes && MK2_PROGRESS_ENABLED(native_progress_mask, row)) { +#if MK2_INFLIGHT + g.native_codes[row * N_POS + pos] = + native_predictor_active_mask && !native_predictor_active_mask[row] + ? 0 : (long long)code; +#else + g.native_codes[row * N_POS + pos] = (long long)code; +#endif + } + } + } +#ifdef ABLATE_SAMPLE + grid_barrier(g.barrier); +#else + if (pos + 1 == n_pos) grid_barrier(g.barrier); +#endif + } + } + if (g.native_seq_len && blockIdx.x == 0) + for (int row = threadIdx.x; row < BATCH; row += blockDim.x) { + MK2_PROGRESS_GUARD(native_progress_mask, row); + g.native_seq_len[row] += 1; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 64 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * SLOTS * sizeof(weight_tile) + + sizeof(out_tile_f) + 2048; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; + const unsigned char *native_predictor_active_mask = nullptr; + MK2_PROGRESS_MEMBER + int maximum_grid = 0; + bool grid_reset_pending = false; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * SLOTS * sizeof(weight_tile) + + sizeof(out_tile_f) + 2048), 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaFuncSetAttribute(predictor_mk2_startup, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = std::min(out->sms * out->blocks_per_sm, 4096); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 64 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->maximum_grid = info.sms * info.blocks_per_sm; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + mk2_check(owned.alloc(&hb.r6_packed,(size_t)BB_EXPERTS*BATCH*D_MODEL*sizeof(bf16))); + mk2_check(owned.alloc(&hb.r6_map,(size_t)BB_EXPERTS*BATCH*sizeof(int))); + mk2_check(owned.alloc(&hb.r6_counts,BB_EXPERTS*sizeof(int))); + mk2_check(owned.alloc(&hb.r6_lut,(1+BB_EXPERTS*DECODE_TILES)*sizeof(int))); + act_gl sparse_views[]={ + act_gl{hb.r6_packed,nullptr,nullptr,(size_t)BB_EXPERTS*BATCH,(size_t)D_MODEL}, + act_gl{moe_h,nullptr,nullptr,(size_t)BB_EXPERTS*BATCH,(size_t)BB_MOE_INTER}}; + mk2_check(owned.alloc(&hb.r6_views,sizeof(sparse_views))); + mk2_check(cudaMemcpy(hb.r6_views,sparse_views,sizeof(sparse_views),cudaMemcpyHostToDevice)); + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.maximum_grid = pred->maximum_grid; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native, bool startup = false) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(s, &capture_status)); + if (st->pred.grid_reset_pending || capture_status != cudaStreamCaptureStatusNone) { + mk2_check(cudaMemsetAsync(st->pred.g.barrier, 0, 64 * sizeof(unsigned int), s)); + st->pred.grid_reset_pending = false; + } + if (startup) { + predictor_mk2_startup<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS, nullptr, + st->pred.native_predictor_active_mask + MK2_PROGRESS_ARGUMENT(st->pred.native_progress_mask)); + } else { + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS, nullptr, + native ? st->pred.native_predictor_active_mask : nullptr + MK2_PROGRESS_ARGUMENT(native ? st->pred.native_progress_mask : nullptr)); + } + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" int mk2_run_native_startup_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream) { + return mk2_enqueue_stage0(p, b, stream, true, true); +} +extern "C" int mk2_set_native_predictor_active_mask(MK2Predictor *p, + const unsigned char *mask) { + try { + if (!p || reinterpret_cast(p)->kind != 1) + throw std::runtime_error("startup mask requires a full Stage0 handle"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && st->pred.native_predictor_active_mask != mask) + throw std::runtime_error("startup mask pointer cannot change after first Stage0 enqueue"); + if (mask) + mk2_stage0_device_pointer(mask, st->pred.device, "native_predictor_active_mask"); + st->pred.native_predictor_active_mask = mask; + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown startup mask binding failure"; return 1; } +} +extern "C" int mk2_native_startup_skip_abi_version() { return 1; } +#include "../../../../../inflight/runtime.cuh" +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "native_prefill_impl.inc" +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *p, int grid) { + if (!p) { g_last_error = "grid setter requires full Stage0 handle"; return 1; } + auto *st = reinterpret_cast(p); + if (st->kind != 1 || grid < 1 || grid > st->maximum_grid) { + g_last_error = "requested Stage0 grid outside independently resident capacity"; + return 1; + } + if (st->grid != grid) st->grid_reset_pending = true; + st->grid = grid; + g_last_error.clear(); + return 0; +} +extern "C" int mk2_r8_bind_initial_state(MK2Predictor *p, const void *const *initial, const void *ready, const void *seq_len, int positions) { + try { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || st->bindings_initialized || !initial || !ready || !seq_len || positions <= 0) + throw std::runtime_error("invalid immutable initial-state binding"); + for (int i = 0; i < 14; ++i) { + mk2_stage0_device_pointer(initial[i], st->pred.device, "initial state"); + st->mb_host[i].r8_initial_ssm = static_cast(initial[i]); + st->mb_host[i].r8_initial_ready = static_cast(ready); + st->mb_host[i].r8_seq_len = static_cast(seq_len); + st->mb_host[i].r8_initial_positions = positions; + } + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" int mk2_r9_state_bind(void* p,const void*const* scales,const void*const* maps) {auto*st=(MK2Stage0State*)p;if(st->bindings_initialized)return 1;for(int i=0;i<14;++i){st->mb_host[i].r9_scales=(const float*)scales[i];st->mb_host[i].r9_map=(const int*)maps[i];}return 0;} +extern "C" int mk2_r9_state_bytes(){return sizeof(r9_state_t);} diff --git a/deploy/nemotron_native/kernels/b256/regions/stage0/stage0.cu b/deploy/nemotron_native/kernels/b256/regions/stage0/stage0.cu new file mode 100644 index 0000000..cfb714b --- /dev/null +++ b/deploy/nemotron_native/kernels/b256/regions/stage0/stage0.cu @@ -0,0 +1,35 @@ +#define MK2_BATCH 256 +#define STAGES 3 +#define KCHUNK 128 +#define MK2_NO_MAIN 1 +#define MK2_DIRECT_INPUT 1 +#define MK2_B1_LOCAL_BARRIERS 1 +#define MK2_SKIP_HIDDEN_SNAPSHOT 1 +#define MK2_ACTIVE_EXPERTS 1 +#define MK2_SSM_SPLITS 1 +#define NMULT 2 +#define MK2_OUT_MULT 2 +#define SSM_PREFETCH_MAX 12 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#if MK2_INFLIGHT +#include +#endif +#include +#include +#include +#include +#include +#include +#include +#include +namespace b256_stage0 { +#include "native/stage0/predictor_mk2.cu" +} diff --git a/deploy/nemotron_native/kernels/b256/stage0.cu b/deploy/nemotron_native/kernels/b256/stage0.cu new file mode 100644 index 0000000..6c04cd7 --- /dev/null +++ b/deploy/nemotron_native/kernels/b256/stage0.cu @@ -0,0 +1 @@ +#include "regions/stage0/stage0.cu" diff --git a/deploy/nemotron_native/kernels/b32/codec.cu b/deploy/nemotron_native/kernels/b32/codec.cu new file mode 100644 index 0000000..7cf7265 --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/codec.cu @@ -0,0 +1,4 @@ +#define NATIVE_CODEC_PREFETCH 1 +#define NATIVE_CODEC_PACK_SUBCTAS 1 +#define D7W 96 +#include "native/codec/native_stage1.cu" diff --git a/deploy/nemotron_native/kernels/b32/native/codec/kernel_kf_codec.h b/deploy/nemotron_native/kernels/b32/native/codec/kernel_kf_codec.h new file mode 100644 index 0000000..9a8b4d5 --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/codec/kernel_kf_codec.h @@ -0,0 +1,13 @@ +#pragma once +#include +struct CodecArgs { + const void *inp[41]; + const void *st[26]; + void *ns[26]; + void *waveform; + void *scratch; + int B; + int F; +}; +void codec_launch(const CodecArgs &a, void *stream); +size_t codec_scratch_bytes(int B, int F); diff --git a/deploy/nemotron_native/kernels/b32/native/codec/legacy_codec.cuh b/deploy/nemotron_native/kernels/b32/native/codec/legacy_codec.cuh new file mode 100644 index 0000000..eb6d793 --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/codec/legacy_codec.cuh @@ -0,0 +1,1390 @@ +#include "kernel_kf_codec.h" +#include +#include +#include +#include +using namespace nvcuda; +template __device__ void prepare_attribute(F,int) {} +#ifndef TC2_QPT +#define TC2_QPT 4 +#endif +#define DIVUP(a,b) (((a) + (b) - 1) / (b)) +__device__ __forceinline__ float half_snake(float v, int c, int half, + const float *__restrict__ alpha) { + if (c < half) { + const float a = alpha[c]; + const float s = __sinf(a * v); + return __fmaf_rn(s, s / (a + 1e-9f), v); + } + return v > 0.f ? v : 0.01f * v; +} +#ifndef WG_ST3 +#define WG_ST3 1 +#define WG_MB3 3 +#define WG_ST7 1 +#define WG_MB7 3 +#define WG_BN3 128 +#define WG_BN7 128 +#endif +#define PDL_ENTER() do { (void)0; if (threadIdx.x == 0) (void)0; } while (0) +#define PDL_TRIGGER() do { if (threadIdx.x == 0) (void)0; } while (0) +template +__device__ __forceinline__ void pdl_launch(KFn kern, dim3 grid, dim3 blk, size_t shmem, + cudaStream_t s, KArgs... args) { + native_invoke(kern, grid, blk, shmem, args...); +} +template +__device__ __forceinline__ void pdl_launch_pf(KFn kern, dim3 grid, dim3 blk, size_t shmem, + cudaStream_t s, const void *pf, size_t pfb, KArgs... args) { + native_invoke_pf(kern, grid, blk, shmem, pf, pfb, args...); +} +#include "wg_kernel.inc" +#include "wg2_kernel.inc" +#include "wg3_kernel.inc" +#include "wg4_step_widths.inc" +#include "wg5_step_widths.inc" +#include "repack3.inc" +#include "tail_kernel.inc" +template +__device__ __noinline__ void dense_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + PDL_TRIGGER(); + constexpr int BK = 16 * KT; + constexpr int NT = NWARPS * 32; + constexpr int WTM = BM / WM; + constexpr int WTN = BN / WN; + constexpr int FM = WTM / 16, FN = WTN / 16; + constexpr int LDA = BK + 8, LDB = BN + 8; + constexpr int GEMM_SM = BM * LDA + BK * LDB; + constexpr int EPI_SM = NWARPS * WTM * WTN * 2; + constexpr int SM_H = GEMM_SM > EPI_SM ? GEMM_SM : EPI_SM; + __half *smem=reinterpret_cast<__half*>(native_arena()); + char *colok=reinterpret_cast(native_aux()); + __half *As = smem; + __half *Bs = smem + BM * LDA; + float *Cw = reinterpret_cast(smem); + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int wm = warp / WN, wn = warp % WN; + const int m0 = blockIdx.x * BM; + const int n0 = blockIdx.y * BN; + const int H = KT - 1; + const int Kg = Cin * KT; + if (tid < BN) { + const int n = n0 + tid; + colok[tid] = (n < N && (n % L) < T) ? 1 : 0; + } + wmma::fragment acc[FM][FN]; +#pragma unroll + for (int i = 0; i < FM; ++i) +#pragma unroll + for (int j = 0; j < FN; ++j) wmma::fill_fragment(acc[i][j], 0.f); + constexpr int ATOT = BM * (BK / 8); + constexpr int AREG = (ATOT + NT - 1) / NT; + constexpr int BTOT = 16 * (BN / 8); + constexpr int BREG = (BTOT + NT - 1) / NT; + int4 ar[AREG]; + int4 br0[BREG], br1[BREG]; +#define CODEC_LOADG(kg0) { const int _ci0 = (kg0) / KT; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * NT; ar[r] = make_int4(0, 0, 0, 0); if (ATOT % NT == 0 || idx < ATOT) { const int m = idx / (BK / 8), v = idx % (BK / 8); const int co = m0 + m, kg = (kg0) + v * 8; if (co < Cout && kg < Kg) ar[r] = *reinterpret_cast(w + (size_t)co * Kg + kg); } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * NT; br0[r] = make_int4(0, 0, 0, 0); br1[r] = make_int4(0, 0, 0, 0); if (BTOT % NT == 0 || idx < BTOT) { const int cil = idx / (BN / 8), nc = idx % (BN / 8); const int ci = _ci0 + cil; if (ci < Cin) { const int4 *src = reinterpret_cast( in + (size_t)ci * SCin + n0 + nc * 8); br0[r] = src[0]; br1[r] = src[1]; } } } } + CODEC_LOADG(0) + for (int kg0 = 0; kg0 < Kg; kg0 += BK) { + __syncthreads(); +#pragma unroll + for (int r = 0; r < AREG; ++r) { + const int idx = tid + r * NT; + if (ATOT % NT == 0 || idx < ATOT) + *reinterpret_cast(As + (idx / (BK / 8)) * LDA + + (idx % (BK / 8)) * 8) = ar[r]; + } +#pragma unroll + for (int r = 0; r < BREG; ++r) { + const int idx = tid + r * NT; + if (BTOT % NT != 0 && idx >= BTOT) continue; + const int cil = idx / (BN / 8), nc = idx % (BN / 8); + unsigned int u[8]; + u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; + u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; +#pragma unroll + for (int k = 0; k < KT; ++k) { + const int kh = k >> 1; + int4 o; + if (k & 1) { + o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); + o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); + o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); + o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); + } else { + o.x = u[kh + 0]; o.y = u[kh + 1]; + o.z = u[kh + 2]; o.w = u[kh + 3]; + } + *reinterpret_cast(Bs + (cil * KT + k) * LDB + nc * 8) = o; + } + } + __syncthreads(); + if (kg0 + BK < Kg) CODEC_LOADG(kg0 + BK) +#pragma unroll + for (int kf = 0; kf < BK; kf += 16) { + wmma::fragment af[FM]; + wmma::fragment bf[FN]; +#pragma unroll + for (int i = 0; i < FM; ++i) + wmma::load_matrix_sync(af[i], As + (wm * WTM + i * 16) * LDA + kf, LDA); +#pragma unroll + for (int j = 0; j < FN; ++j) + wmma::load_matrix_sync(bf[j], Bs + kf * LDB + wn * WTN + j * 16, LDB); +#pragma unroll + for (int i = 0; i < FM; ++i) +#pragma unroll + for (int j = 0; j < FN; ++j) + wmma::mma_sync(acc[i][j], af[i], bf[j], acc[i][j]); + } + } +#undef CODEC_LOADG + __syncthreads(); + float *Cme = Cw + warp * WTM * WTN; +#pragma unroll + for (int i = 0; i < FM; ++i) +#pragma unroll + for (int j = 0; j < FN; ++j) + wmma::store_matrix_sync(Cme + (i * 16) * WTN + j * 16, acc[i][j], WTN, + wmma::mem_row_major); + __syncwarp(); + const int half = Cout >> 1; + const int lane = tid & 31; + for (int e = lane; e < WTM * WTN; e += 32) { + const int mm = e / WTN, nn = e % WTN; + const int co = m0 + wm * WTM + mm; + const int nl = wn * WTN + nn; + if (co >= Cout || !colok[nl]) continue; + const int n = n0 + nl; + float v = Cme[e] + bias[co]; + if (resid != nullptr) + v += __half2float(resid[(size_t)co * SCout + n + H]); + if (do_snake) v = half_snake(v, co, half, alpha); + out[(size_t)co * SCout + n + H] = __float2half(v); + } + if (blockIdx.y == 0) { + const int nbx = gridDim.x; + if (ns != nullptr) { + const int B = N / L; + const int total = Cin * B * H; + for (int idx = blockIdx.x * NT + tid; idx < total; idx += nbx * NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int B = N / L; + const int total = Cout * B * pad_len; + for (int idx = blockIdx.x * NT + tid; idx < total; idx += nbx * NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} +template +__device__ __noinline__ void transpose_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + __half *__restrict__ out, const __half *__restrict__ pad_src, + __half *__restrict__ ns, int Cin, int Cout, int Tin, + int Lin, int SCin, int jbase, int Lout, int SCout, int Hout, int pad_off, + int pad_len, int B) { + if constexpr (PDL) PDL_ENTER(); + constexpr int Ktaps = 2 * STRIDE; + constexpr int QPT = QPT_; + constexpr int VPT = QPT * STRIDE; + static_assert(QPT % 2 == 0, "half2 input loads need an even frame count"); + static_assert(VPT % 2 == 0, "half2 staging needs an even slot"); + float *ws=reinterpret_cast(native_aux()); + __half *sm=reinterpret_cast<__half*>(native_arena()); + const int co = blockIdx.y; + const int b = blockIdx.z; + const int tid = threadIdx.x, BLK = blockDim.x; + for (int i = tid; i < G * Ktaps; i += BLK) + ws[i] = __half2float(w[(size_t)(co * G + i / Ktaps) * Ktaps + (i % Ktaps)]); + __syncthreads(); + const float bs = bias[0 + co]; + const int half = Cout >> 1; + const bool snake = co < half; + const float av = snake ? alpha[co] : 0.f; + const float ri = snake ? __frcp_rn(av + 1e-9f) : 0.f; + const __half *ic = in + (size_t)(co * G) * SCin + b * Lin + jbase; + __half *oc = out + (size_t)co * SCout + b * Lout + Hout; + constexpr bool VEC16 = (VPT % 8) == 0; + const int ophase = + VEC16 ? 0 : (int)((((uintptr_t)oc) >> 1) & 7u); + __half *const sbase = sm + ophase; + const int hd = (8 - ophase) & 7; + const int lane = tid & 31; + for (int qb = 1 + blockIdx.x * BLK * QPT; qb <= Tin; + qb += gridDim.x * BLK * QPT) { + const int q = qb + QPT * tid; + const bool a0 = (q <= Tin); + float jq[QPT][G], jp0[G]; +#pragma unroll + for (int g = 0; g < G; ++g) { +#pragma unroll + for (int e2 = 0; e2 < QPT / 2; ++e2) { + const __half2 v = a0 ? *reinterpret_cast( + ic + (size_t)g * SCin + q + 2 * e2) + : __floats2half2_rn(0.f, 0.f); + jq[2 * e2][g] = __half2float(v.x); + jq[2 * e2 + 1][g] = __half2float(v.y); + } + } +#pragma unroll + for (int g = 0; g < G; ++g) { + const float up = __shfl_up_sync(0xffffffffu, jq[QPT - 1][g], 1); + jp0[g] = (lane == 0) + ? (a0 ? __half2float(ic[(size_t)g * SCin + q - 1]) : 0.f) + : up; + } + __half *const sp = sbase + tid * VPT; + __align__(16) __half vreg[VPT]; +#pragma unroll + for (int r = 0; r < STRIDE; ++r) { + float wc[G], wp[G]; +#pragma unroll + for (int g = 0; g < G; ++g) { + wc[g] = ws[g * Ktaps + r]; + wp[g] = ws[g * Ktaps + r + STRIDE]; + } +#pragma unroll + for (int e = 0; e < QPT; ++e) { + float v = bs; +#pragma unroll + for (int g = 0; g < G; ++g) { + const float cur = jq[e][g]; + const float prv = (e == 0) ? jp0[g] : jq[e - 1][g]; + v = __fmaf_rn(wc[g], cur, __fmaf_rn(wp[g], prv, v)); + } + if (snake) { + const float sn = __sinf(av * v); + v = __fmaf_rn(sn * sn, ri, v); + } else { + v = v > 0.f ? v : 0.01f * v; + } + vreg[e * STRIDE + r] = __float2half(v); + } + } + if constexpr (VEC16) { +#pragma unroll + for (int m = 0; m < VPT / 8; ++m) + *reinterpret_cast(sp + m * 8) = + *reinterpret_cast(vreg + m * 8); + } else { +#pragma unroll + for (int i = 0; i < VPT; ++i) sp[i] = vreg[i]; + } + __syncthreads(); + const int nval = min(BLK * QPT, Tin - qb + 1); + const int tot = nval * STRIDE; + __half *dst = oc + (size_t)(qb - 1) * STRIDE; + if constexpr (VEC16) { + const int tot2 = tot >> 1; + for (int i = tid; i < tot2; i += BLK) + reinterpret_cast<__half2 *>(dst)[i] = + reinterpret_cast(sbase)[i]; + if ((tot & 1) && tid == 0) dst[tot - 1] = sbase[tot - 1]; + } else { + const int head = hd < tot ? hd : tot; + for (int i = tid; i < head; i += BLK) dst[i] = sbase[i]; + const int nb4 = (tot - head) >> 3; + const int4 *s4 = reinterpret_cast(sbase + head); + int4 *d4 = reinterpret_cast(dst + head); + for (int i = tid; i < nb4; i += BLK) d4[i] = s4[i]; + for (int i = head + (nb4 << 3) + tid; i < tot; i += BLK) + dst[i] = sbase[i]; + } + __syncthreads(); + } + if (blockIdx.x == 0) { + if (ns != nullptr) { + for (int ci = co + tid * Cout; ci < Cin; ci += BLK * Cout) + ns[(size_t)b * Cin + ci] = + in[(size_t)ci * SCin + b * Lin + jbase + native_true_tokens(b,Tin)]; + } + if (pad_src != nullptr) { + for (int j = tid; j < pad_len; j += BLK) + out[(size_t)co * SCout + b * Lout + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } +} +template +__device__ __noinline__ void transpose_conv_early_bpack_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + __half *__restrict__ out, const __half *__restrict__ pad_src, + __half *__restrict__ ns, int Cin, int Cout, int Tin, + int Lin, int SCin, int jbase, int Lout, int SCout, int Hout, int pad_off, + int pad_len, int B) { + if constexpr (PDL) PDL_ENTER(); + static_assert(BPACK <= 32); + static_assert(LPB == 8 || LPB == 16 || LPB == 32); + constexpr int Ktaps = 2 * STRIDE; + constexpr int WARP = 32; + constexpr int NT = WARP * BPACK; + constexpr int BPB = WARP / LPB; + float *ws=reinterpret_cast(native_aux()); + __half *sm=reinterpret_cast<__half*>(native_arena()); + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int lane = tid & 31; + const int sub = lane % LPB; + const int bsub = lane / LPB; + const int co = blockIdx.y; + const int b = blockIdx.z * (BPACK * BPB) + warp * BPB + bsub; + const bool active = b < B; + for (int i = tid; i < G * Ktaps; i += NT) + ws[i] = __half2float(w[(size_t)(co * G + i / Ktaps) * Ktaps + + (i % Ktaps)]); + __syncthreads(); + const int qb = 1 + blockIdx.x * LPB; + const int q = qb + sub; + __half *sp = sm + (warp * WARP + lane) * STRIDE; + const bool live = active && q <= Tin; + { + const float bs = bias[co]; + const int half = Cout >> 1; + const bool snake = co < half; + const float av = snake ? alpha[co] : 0.f; + const float ri = snake ? __frcp_rn(av + 1e-9f) : 0.f; + const __half *ic = in + (size_t)(co * G) * SCin + b * Lin + jbase; + float jq[G], jp[G]; +#pragma unroll + for (int g = 0; g < G; ++g) + jq[g] = live ? __half2float(ic[(size_t)g * SCin + q]) : 0.f; +#pragma unroll + for (int g = 0; g < G; ++g) { + const float up = __shfl_up_sync(0xffffffffu, jq[g], 1, LPB); + jp[g] = (sub == 0) + ? (live ? __half2float(ic[(size_t)g * SCin + q - 1]) : 0.f) + : up; + } + if (live) +#pragma unroll + for (int r = 0; r < STRIDE; ++r) { + float v = bs; +#pragma unroll + for (int g = 0; g < G; ++g) + v = __fmaf_rn(ws[g * Ktaps + r], jq[g], + __fmaf_rn(ws[g * Ktaps + r + STRIDE], jp[g], v)); + if (snake) { + const float sn = __sinf(av * v); + v = __fmaf_rn(sn * sn, ri, v); + } else { + v = v > 0.f ? v : 0.01f * v; + } + sp[r] = __float2half(v); + } + } + __syncwarp(); + if (active) { + const int nval = min(LPB, Tin - (qb - 1)); + const int tot2 = (nval * STRIDE) >> 1; + const __half *src = sm + (warp * WARP + bsub * LPB) * STRIDE; + __half *dst = out + (size_t)co * SCout + b * Lout + Hout + + (size_t)(qb - 1) * STRIDE; + for (int i = sub; i < tot2; i += LPB) + reinterpret_cast<__half2 *>(dst)[i] = + reinterpret_cast(src)[i]; + if ((nval * STRIDE) & 1) + if (sub == 0) dst[nval * STRIDE - 1] = src[nval * STRIDE - 1]; + } + if (blockIdx.x == 0 && active) { + if (ns != nullptr) { + for (int ci = co + sub * Cout; ci < Cin; ci += LPB * Cout) + ns[(size_t)b * Cin + ci] = + in[(size_t)ci * SCin + b * Lin + jbase + native_true_tokens(b,Tin)]; + } + if (pad_src != nullptr) { + for (int j = sub; j < pad_len; j += LPB) + out[(size_t)co * SCout + b * Lout + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } +} +template +__device__ __noinline__ void transpose_conv_bpack2_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + __half *__restrict__ out, const __half *__restrict__ pad_src, + __half *__restrict__ ns, int Cin, int Cout, int Tin, + int Lin, int SCin, int jbase, int Lout, int SCout, int Hout, int pad_off, + int pad_len, int B) { + if constexpr (PDL) PDL_ENTER(); + static_assert(BPACK <= 32); + static_assert(LPB == 8 || LPB == 16 || LPB == 32); + constexpr int Ktaps = 2 * STRIDE; + constexpr int WARP = 32; + constexpr int NT = WARP * BPACK; + constexpr int BPB = WARP / LPB; + constexpr int QPT = 2; + constexpr int VPT = QPT * STRIDE; + float *ws=reinterpret_cast(native_aux()); + __half *sm=reinterpret_cast<__half*>(native_arena()); + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int lane = tid & 31; + const int sub = lane % LPB; + const int bsub = lane / LPB; + const int co = blockIdx.y; + const int b = blockIdx.z * (BPACK * BPB) + warp * BPB + bsub; + const bool active = b < B; + for (int i = tid; i < G * Ktaps; i += NT) + ws[i] = __half2float(w[(size_t)(co * G + i / Ktaps) * Ktaps + + (i % Ktaps)]); + __syncthreads(); + const int qb = 1 + blockIdx.x * (LPB * QPT); + const int q = qb + sub * QPT; + __half *sp = sm + (warp * WARP + lane) * VPT; + const bool a0 = active && (q <= Tin); + { + const float bs = bias[co]; + const int half = Cout >> 1; + const bool snake = co < half; + const float av = snake ? alpha[co] : 0.f; + const float ri = snake ? __frcp_rn(av + 1e-9f) : 0.f; + const __half *ic = in + (size_t)(co * G) * SCin + b * Lin + jbase; + float jq0[G], jq1[G], jp0[G]; +#pragma unroll + for (int g = 0; g < G; ++g) { + const __half2 v = a0 ? *reinterpret_cast( + ic + (size_t)g * SCin + q) + : __floats2half2_rn(0.f, 0.f); + jq0[g] = __half2float(v.x); + jq1[g] = __half2float(v.y); + } +#pragma unroll + for (int g = 0; g < G; ++g) { + const float up = __shfl_up_sync(0xffffffffu, jq1[g], 1, LPB); + jp0[g] = (sub == 0) + ? (a0 ? __half2float(ic[(size_t)g * SCin + q - 1]) + : 0.f) + : up; + } + if (a0) { + float vv[VPT]; +#pragma unroll + for (int r = 0; r < STRIDE; ++r) { + float wc[G], wp[G]; +#pragma unroll + for (int g = 0; g < G; ++g) { + wc[g] = ws[g * Ktaps + r]; + wp[g] = ws[g * Ktaps + r + STRIDE]; + } +#pragma unroll + for (int e = 0; e < QPT; ++e) { + float v = bs; +#pragma unroll + for (int g = 0; g < G; ++g) { + const float cur = (e == 0) ? jq0[g] : jq1[g]; + const float prv = (e == 0) ? jp0[g] : jq0[g]; + v = __fmaf_rn(wc[g], cur, __fmaf_rn(wp[g], prv, v)); + } + if (snake) { + const float sn = __sinf(av * v); + v = __fmaf_rn(sn * sn, ri, v); + } else { + v = v > 0.f ? v : 0.01f * v; + } + vv[e * STRIDE + r] = v; + } + } + __half2 *sp2 = reinterpret_cast<__half2 *>(sp); +#pragma unroll + for (int i = 0; i < VPT / 2; ++i) + sp2[i] = __floats2half2_rn(vv[2 * i], vv[2 * i + 1]); + } + } + __syncwarp(); + if (active) { + const int nval = min(LPB * QPT, Tin - (qb - 1)); + const int tot = nval * STRIDE; + const __half *src = sm + (warp * WARP + bsub * LPB) * VPT; + __half *dst = out + (size_t)co * SCout + b * Lout + Hout + + (size_t)(qb - 1) * STRIDE; + for (int i = sub; i < (tot >> 1); i += LPB) + reinterpret_cast<__half2 *>(dst)[i] = + reinterpret_cast(src)[i]; + if (tot & 1) + if (sub == 0) dst[tot - 1] = src[tot - 1]; + } + if (blockIdx.x == 0 && active) { + if (ns != nullptr) { + for (int ci = co + sub * Cout; ci < Cin; ci += LPB * Cout) + ns[(size_t)b * Cin + ci] = + in[(size_t)ci * SCin + b * Lin + jbase + native_true_tokens(b,Tin)]; + } + if (pad_src != nullptr) { + for (int j = sub; j < pad_len; j += LPB) + out[(size_t)co * SCout + b * Lout + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } +} +#define TCONV_BPACK(G,STRIDE,BPACK,LPB,TIN,COUT,BB,...) pdl_launch(transpose_conv_early_bpack_kernel, dim3(DIVUP(TIN, LPB), COUT, DIVUP(BB, BPACK * (32 / LPB))), dim3(32 * BPACK), (size_t)(32 * (BPACK)) * (STRIDE) * 2 + 64, s, __VA_ARGS__) +#define TCONV_BP2(G,STRIDE,LPB,TIN,COUT,BB,...) pdl_launch(transpose_conv_bpack2_kernel, dim3(DIVUP(TIN, (LPB) * 2), COUT, DIVUP(BB, 32)), dim3(32 * (LPB)), (size_t)(32 * (LPB)) * (4 * (STRIDE)) + 64, s, __VA_ARGS__) +#define TCONV_BPACK_DISPATCH1(G,STRIDE,BPACK,TIN,COUT,BB,...) do { if ((TIN) >= 32) TCONV_BP2(G, STRIDE, 16, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) >= 16) TCONV_BP2(G, STRIDE, 8, TIN, COUT, BB, __VA_ARGS__); else TCONV_BPACK(G, STRIDE, BPACK, 16, TIN, COUT, BB, __VA_ARGS__); } while (0) +#define TCONV_BPACK_DISPATCH(G,STRIDE,BPACK,TIN,COUT,BB,...) do { if ((TIN) >= 32) TCONV_BP2(G, STRIDE, 16, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) >= 16) TCONV_BP2(G, STRIDE, 8, TIN, COUT, BB, __VA_ARGS__); else TCONV_BPACK(G, STRIDE, BPACK, 16, TIN, COUT, BB, __VA_ARGS__); } while (0) +#define TCONV_BPACK_DISPATCH8(G,STRIDE,BPACK,TIN,COUT,BB,...) do { if ((TIN) >= 32) TCONV_BP2(G, STRIDE, 16, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) >= 16) TCONV_BP2(G, STRIDE, 8, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) <= 8) TCONV_BPACK(G, STRIDE, BPACK, 8, TIN, COUT, BB, __VA_ARGS__); else TCONV_BPACK(G, STRIDE, BPACK, 16, TIN, COUT, BB, __VA_ARGS__); } while (0) +__device__ __noinline__ void pack_latent_kernel(const __half *__restrict__ latent, + const __half *__restrict__ state, + __half *__restrict__ out, int Cin, int T, + int L, int SC, int H, int B) { + const int total = Cin * B * L; + for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; + idx += gridDim.x * blockDim.x) { + const int j = idx % L, r = idx / L; + const int b = r % B, ci = r / B; + out[(size_t)ci * SC + b * L + j] = + (j < H) ? state[((size_t)b * Cin + ci) * H + j] + : latent[((size_t)b * Cin + ci) * T + (j - H)]; + } +} +#define PC_VEC 8 +__device__ __noinline__ void post_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, __half *__restrict__ y, + __half *__restrict__ ns, int Cin, int T, int L, int SC, int KT, int B, + int N) { + float *pw=reinterpret_cast(native_arena()); + for (int i = threadIdx.x; i < Cin * KT; i += 256) pw[i] = __half2float(w[i]); + __syncthreads(); + const int H = KT - 1; + const float bs = bias[0]; + for (int n = (blockIdx.x * 256 + threadIdx.x) * PC_VEC; n < N; + n += gridDim.x * 256 * PC_VEC) { + float v[PC_VEC]; +#pragma unroll + for (int i = 0; i < PC_VEC; ++i) v[i] = bs; + for (int ci = 0; ci < Cin; ++ci) { + const __half *xr = in + (size_t)ci * SC + n; + const float *wr = pw + ci * KT; + const int4 lo = *reinterpret_cast(xr); + const int4 hi = *reinterpret_cast(xr + 8); + unsigned u[8]; + u[0] = lo.x; u[1] = lo.y; u[2] = lo.z; u[3] = lo.w; + u[4] = hi.x; u[5] = hi.y; u[6] = hi.z; u[7] = hi.w; + float x[PC_VEC + 6]; +#pragma unroll + for (int j = 0; j < (PC_VEC + 6) / 2; ++j) { + const float2 f = __half22float2( + *reinterpret_cast(&u[j])); + x[2 * j] = f.x; + x[2 * j + 1] = f.y; + } +#pragma unroll + for (int i = 0; i < PC_VEC; ++i) + for (int k = 0; k < 7; ++k) + if (k < KT) v[i] = __fmaf_rn(wr[k], x[i + k], v[i]); + } + const int b0 = n / L, t0 = n - (n / L) * L; +#pragma unroll + for (int i = 0; i < PC_VEC; ++i) { + int t = t0 + i, b = b0; + if (t >= L) { t -= L; b = b0 + 1; } + if (t < T && b < B) + y[(size_t)b * T + t] = + __float2half(fminf(fmaxf(v[i], -1.f), 1.f)); + } + } + if (ns != nullptr) { + const int total = B * Cin * H; + for (int idx = blockIdx.x * 256 + threadIdx.x; idx < total; + idx += gridDim.x * 256) { + const int j = idx % H, r = idx / H; + const int ci = r % Cin, b = r / Cin; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SC + b * L + native_true_tokens(b,T) + j]; + } + } +} +#define RP_JOBS 6 +struct RepackSet { + const __half *src[RP_JOBS]; + __half *dst[RP_JOBS]; + int Cout[RP_JOBS], CoutP[RP_JOBS], Cin[RP_JOBS]; + int base[RP_JOBS], nch[RP_JOBS], rs[RP_JOBS]; + int k3base; + const __half *k3src[2]; + __half *k3dst[2]; +}; +__host__ __device__ inline int native_RP_COUT(int j) {const int v[6]={768, 384, 256, 128, 64, 32};return v[j];} +__host__ __device__ inline int native_RP_CIN(int j) {const int v[6]={384, 768, 128, 256, 32, 64};return v[j];} +__host__ __device__ inline int native_RP_COUTP(int j) {const int v[6]={768, 384, 256, 128, 128, 128};return v[j];} +template +__device__ __noinline__ void repack_kernel(RepackSet js) { + const int bx = blockIdx.x; + int job, job_bx; + if (bx < 1152) { + job = 0; job_bx = 0; + } else if (bx < 2304) { + job = 1; job_bx = 1152; + } else if (bx < 2432) { + job = 2; job_bx = 2304; + } else if (bx < 2560) { + job = 3; job_bx = 2432; + } else if (bx < 2576) { + job = 4; job_bx = 2560; + } else { + job = 5; job_bx = 2576; + } + const int Cout = js.Cout[job], CoutP = js.CoutP[job], Cin = js.Cin[job]; + const int idx = (bx - job_bx) * blockDim.x + threadIdx.x; + const int co = idx / Cin, ci = idx - co * Cin; + const int c = ci / CIB, cl = ci - c * CIB; + __half *d = js.dst[job] + (size_t)c * (KT * CoutP * CIB) + + (size_t)(co >> 3) * (CIB * 8) + (cl >> 3) * 64 + (co & 7) * 8 + + (cl & 7); + const __half *sp = js.src[job] + (size_t)(co * Cin + ci) * KT; + const bool live = co < Cout; +#pragma unroll + for (int k = 0; k < KT; ++k) + d[(size_t)k * CoutP * CIB] = live ? sp[k] : __float2half(0.f); +} +template +__device__ __noinline__ void repack_rs_kernel( + const RepackSet &js, const __half *__restrict__ latent, const __half *__restrict__ latent_state, + __half *__restrict__ latent_out, int latent_C, int latent_T, int latent_L, + int latent_SC, int latent_H, int latent_B) { + PDL_TRIGGER(); + constexpr int CO_TILE = 16; + constexpr int CI_TILE = 64; + constexpr int ROW_H = CI_TILE * KT + 8; + constexpr int ROW4 = CI_TILE * KT / 8; + __half *tile=reinterpret_cast<__half*>(native_arena()); + const int bx = blockIdx.x; + { + const int idx = bx + threadIdx.x * gridDim.x; + const int total = latent_C * latent_B * latent_L; + if (idx < total) { + const int j = idx % latent_L, r = idx / latent_L; + const int b = r % latent_B, ci = r / latent_B; + latent_out[(size_t)ci * latent_SC + b * latent_L + j] = + (j < latent_H) + ? latent_state[((size_t)b * latent_C + ci) * latent_H + j] + : latent[((size_t)b * latent_C + ci) * latent_T + + (j - latent_H)]; + } + } + if (bx >= js.k3base) { + constexpr int K3 = 3; + constexpr int R3_H = CI_TILE * K3 + 8; + constexpr int R3_4 = CI_TILE * K3 / 8; + const int r = bx - js.k3base; + const int conv = r / 1152; + const int loc = r - conv * 1152; + const int Cout3 = conv ? 768 : 1536; + const int Cin3 = conv ? 1536 : 768; + const int nci3 = Cin3 / CI_TILE; + const int cot3 = loc / nci3; + if (cot3 * CO_TILE >= Cout3) return; + const int co03 = cot3 * CO_TILE; + const int ci03 = (loc - cot3 * nci3) * CI_TILE; + const int NCH3 = Cin3 >> 4; + const __half *const s3 = js.k3src[conv]; + __half *const d3 = js.k3dst[conv]; + const int t3 = threadIdx.x; +#pragma unroll + for (int n = 0; n < CO_TILE * R3_4 / 128; ++n) { + const int i = t3 + n * 128; + const int co = i / R3_4, q = i - co * R3_4; + reinterpret_cast(tile + co * R3_H)[q] = + *reinterpret_cast( + s3 + ((size_t)(co03 + co) * Cin3 + ci03) * K3 + q * 8); + } + __syncthreads(); + const int w3 = t3 >> 5, l3 = t3 & 31; + const int rA3 = l3 >> 2, q43 = l3 & 3; + const int cb3 = co03 >> 4; + __half *const db3 = d3 + ((size_t)(cb3 >> 2) * NCH3 * K3) * 1024 + + ((cb3 & 3) * 32 + l3) * 8; +#pragma unroll + for (int u = w3; u < (CI_TILE / 16) * K3; u += 4) { + const int cl = u / K3, tap = u - cl * K3; + const int cgi = (ci03 >> 4) + cl; + const __half *sa = tile + rA3 * R3_H + (cl * 16 + 2 * q43) * K3 + tap; + const __half *sb = sa + 8 * R3_H; + uint32_t o[4]; + o[0] = __half_as_ushort(sa[0]) | + ((uint32_t)__half_as_ushort(sa[K3]) << 16); + o[1] = __half_as_ushort(sb[0]) | + ((uint32_t)__half_as_ushort(sb[K3]) << 16); + o[2] = __half_as_ushort(sa[8 * K3]) | + ((uint32_t)__half_as_ushort(sa[9 * K3]) << 16); + o[3] = __half_as_ushort(sb[8 * K3]) | + ((uint32_t)__half_as_ushort(sb[9 * K3]) << 16); + *reinterpret_cast(db3 + (size_t)(cgi * K3 + tap) * 1024) = + make_int4(o[0], o[1], o[2], o[3]); + } + return; + } + int job = 0; +#pragma unroll + for (int q = 1; q < RP_JOBS; ++q) if (bx >= js.base[q]) job = q; + const int local = bx - js.base[job]; + const int Cout = js.Cout[job], Cin = js.Cin[job]; + const int nci = (Cin + CI_TILE - 1) / CI_TILE; + const int cot = local / nci; + const int co0 = cot * CO_TILE; + const int ci0 = (local - cot * nci) * CI_TILE; + const int NCH = js.nch[job]; + const int t = threadIdx.x; + const int valid4 = (min(Cin - ci0, CI_TILE) * KT) >> 3; + for (int i = t; i < CO_TILE * ROW4; i += 128) { + const int co = i / ROW4, q = i - co * ROW4; + int4 v = make_int4(0, 0, 0, 0); + if (co0 + co < Cout && q < valid4) + v = *reinterpret_cast( + js.src[job] + ((size_t)(co0 + co) * Cin + ci0) * KT + q * 8); + reinterpret_cast(tile + co * ROW_H)[q] = v; + } + __syncthreads(); + const int warp = t >> 5, lane = t & 31; + const int rA = lane >> 2, q4 = lane & 3; + const int cb = co0 >> 4; + const int mtile = cb >> 2, wsub = cb & 3; + __half *const dbase = js.dst[job] + + ((size_t)mtile * NCH * KT) * 1024 + (wsub * 32 + lane) * 8; + if (!js.rs[job]) { + const int CoutP = js.CoutP[job]; + constexpr int UN2 = 16 * (CI_TILE / 8) * KT; +#pragma unroll + for (int n = 0; n < UN2 / 128; ++n) { + const int u = t + n * 128; + const int col = u & 15, cig = (u >> 4) & (CI_TILE / 8 - 1); + const int tap = u / (16 * (CI_TILE / 8)); + const int gci = ci0 + cig * 8; + if (gci >= Cin) continue; + const int gco = co0 + col; + uint32_t o[4]; + const __half *sp = tile + col * ROW_H; +#pragma unroll + for (int e = 0; e < 4; ++e) + o[e] = __half_as_ushort(sp[(cig * 8 + 2 * e) * KT + tap]) | + ((uint32_t)__half_as_ushort( + sp[(cig * 8 + 2 * e + 1) * KT + tap]) << 16); + __half *d = js.dst[job] + + (size_t)(gci >> 4) * ((size_t)KT * CoutP * 16) + + (size_t)tap * (CoutP * 16) + (size_t)(gco >> 3) * 128 + + ((gci >> 3) & 1) * 64 + (gco & 7) * 8; + *reinterpret_cast(d) = make_int4(o[0], o[1], o[2], o[3]); + } + return; + } + constexpr int UNITS = (CI_TILE / 16) * KT; +#pragma unroll + for (int u = warp; u < UNITS; u += 4) { + const int cl = u / KT, tap = u - cl * KT; + const int cgi = (ci0 >> 4) + cl; + if (cgi >= NCH) continue; + const __half *sa = tile + rA * ROW_H + (cl * 16 + 2 * q4) * KT + tap; + const __half *sb = sa + 8 * ROW_H; + uint32_t o[4]; + o[0] = __half_as_ushort(sa[0]) | + ((uint32_t)__half_as_ushort(sa[KT]) << 16); + o[1] = __half_as_ushort(sb[0]) | + ((uint32_t)__half_as_ushort(sb[KT]) << 16); + o[2] = __half_as_ushort(sa[8 * KT]) | + ((uint32_t)__half_as_ushort(sa[9 * KT]) << 16); + o[3] = __half_as_ushort(sb[8 * KT]) | + ((uint32_t)__half_as_ushort(sb[9 * KT]) << 16); + *reinterpret_cast(dbase + (size_t)(cgi * KT + tap) * 1024) = + make_int4(o[0], o[1], o[2], o[3]); + } +} +__device__ __noinline__ void repack3_rs_kernel( + const __half *__restrict__ win, const __half *__restrict__ wsk, + __half *__restrict__ din, __half *__restrict__ dsk, + const __half *__restrict__ asrc, __half *__restrict__ adst, int SCA, + int ncol, int ntile) { + PDL_TRIGGER(); + if (blockIdx.y == 12) { + if (adst == nullptr) return; + for (int u = blockIdx.x; u < 96 * ntile; u += gridDim.x) { + const int g = u / ntile; + const int col = (u - g * ntile) * 128 + (int)threadIdx.x; + if (col >= ncol) continue; + __half v[8]; +#pragma unroll + for (int h = 0; h < 8; ++h) + v[h] = asrc[(size_t)(g * 8 + h) * SCA + col]; + *reinterpret_cast(adst + (size_t)g * ((size_t)SCA * 8) + + (size_t)col * 8) = + *reinterpret_cast(v); + } + return; + } + constexpr int KT = 3; + constexpr int CO_TILE = 16; + constexpr int CI_TILE = 64; + constexpr int ROW_H = CI_TILE * KT + 8; + constexpr int ROW4 = CI_TILE * KT / 8; + __half *tile=reinterpret_cast<__half*>(native_arena()); + const int jy = blockIdx.y; + const int conv = jy & 1, b = 5 - (jy >> 1); + const int Cout = conv ? 768 : 1536; + const int Cin = conv ? 1536 : 768; + const __half *src = (conv ? wsk : win) + (size_t)(b + 1) * Cout * Cin * KT; + __half *dst = (conv ? dsk : din) + (size_t)b * ((size_t)Cout * Cin * KT); + const int nci = Cin / CI_TILE; + const int cot = blockIdx.x / nci; + const int co0 = cot * CO_TILE; + const int ci0 = (blockIdx.x - cot * nci) * CI_TILE; + const int NCH = Cin >> 4; + const int t = threadIdx.x; +#pragma unroll + for (int n = 0; n < CO_TILE * ROW4 / 128; ++n) { + const int i = t + n * 128; + const int co = i / ROW4, q = i - co * ROW4; + reinterpret_cast(tile + co * ROW_H)[q] = + *reinterpret_cast( + src + ((size_t)(co0 + co) * Cin + ci0) * KT + q * 8); + } + __syncthreads(); + const int warp = t >> 5, lane = t & 31; + const int rA = lane >> 2, q4 = lane & 3; + const int cb = co0 >> 4; + __half *const dbase = dst + ((size_t)(cb >> 2) * NCH * KT) * 1024 + + ((cb & 3) * 32 + lane) * 8; +#pragma unroll + for (int u = warp; u < (CI_TILE / 16) * KT; u += 4) { + const int cl = u / KT, tap = u - cl * KT; + const int cgi = (ci0 >> 4) + cl; + const __half *sa = tile + rA * ROW_H + (cl * 16 + 2 * q4) * KT + tap; + const __half *sb = sa + 8 * ROW_H; + uint32_t o[4]; + o[0] = __half_as_ushort(sa[0]) | + ((uint32_t)__half_as_ushort(sa[KT]) << 16); + o[1] = __half_as_ushort(sb[0]) | + ((uint32_t)__half_as_ushort(sb[KT]) << 16); + o[2] = __half_as_ushort(sa[8 * KT]) | + ((uint32_t)__half_as_ushort(sa[9 * KT]) << 16); + o[3] = __half_as_ushort(sb[8 * KT]) | + ((uint32_t)__half_as_ushort(sb[9 * KT]) << 16); + *reinterpret_cast(dbase + (size_t)(cgi * KT + tap) * 1024) = + make_int4(o[0], o[1], o[2], o[3]); + } +} +__device__ __noinline__ void repack_activation_only(const __half *asrc,__half *adst,int SCA,int ncol,int ntile){ + if (adst == nullptr) return; + for (int u = blockIdx.x; u < 96 * ntile; u += gridDim.x) { + const int g = u / ntile; + const int col = (u - g * ntile) * 128 + (int)threadIdx.x; + if (col >= ncol) continue; + __half v[8]; +#pragma unroll + for (int h = 0; h < 8; ++h) + v[h] = asrc[(size_t)(g * 8 + h) * SCA + col]; + *reinterpret_cast(adst + (size_t)g * ((size_t)SCA * 8) + + (size_t)col * 8) = + *reinterpret_cast(v); + } + return; +} +namespace { +struct Stack { + int T, L, N, SC, H; +}; +inline __host__ __device__ int tc_blk(int Tin, int qpt = 2) { + int best = 256, bw = 1 << 30; + for (int b = 256; b >= 32; b >>= 1) { + const int per = qpt * b; + const int waste = (Tin + per - 1) / per * per - Tin; + if (waste < bw) { bw = waste; best = b; } + } + return best; +} +inline __host__ __device__ int roundup(int a, int b) { return (a + b - 1) / b * b; } +__host__ __device__ Stack mk(int B, int T, int K) { + Stack s; + s.T = T; + s.H = K - 1; + s.L = s.H + T; + s.N = B * s.L; + s.SC = roundup(s.N, 128) + 128; + return s; +} +struct Plan { + Stack s0, s1, s2, s3, s4; + __half *LB, *X0, *M0, *Y0; + __half *A, *Ac, *Bb, *M1; + __half *P, *Q, *M2; + __half *R, *S, *M3; + __half *U; + __half *W7[RP_JOBS]; + __half *W3i, *W3s; + size_t total; +}; +static constexpr int kTuneB = 32; +inline __host__ __device__ int nref(const Stack &st) { return kTuneB * st.L; } +inline __host__ __device__ bool use_rp3i(int, int F) { return (size_t)kTuneB * (4 * F + 2) > 400; } +inline __host__ __device__ bool use_rp3s(int, int F) { return (size_t)kTuneB * (4 * F + 2) > 400; } +__host__ __device__ Plan build_plan(int B, int F, __half *base) { + Plan p; + p.s0 = mk(B, 2 * F, 3); + p.s1 = mk(B, 4 * F, 3); + p.s2 = mk(B, 36 * F, 7); + p.s3 = mk(B, 252 * F, 7); + p.s4 = mk(B, 1764 * F, 7); + size_t off = 0; + auto alloc = [&](int ch, const Stack &s) -> __half * { + __half *r = base ? base + off : nullptr; + off += (size_t)ch * s.SC; + return r; + }; + p.LB = alloc(40, p.s0); + p.X0 = alloc(768, p.s0); + p.M0 = alloc(1536, p.s0); + p.Y0 = alloc(768, p.s0); + p.A = alloc(768, p.s1); + p.Ac = use_rp3i(B, F) ? alloc(768, p.s1) : nullptr; + p.Bb = alloc(768, p.s1); + p.M1 = alloc(1536, p.s1); + p.P = alloc(384, p.s2); + p.Q = alloc(384, p.s2); + p.M2 = alloc(768, p.s2); + p.R = alloc(128, p.s3); + p.S = alloc(128, p.s3); + p.M3 = alloc(256, p.s3); + p.U = alloc(32, p.s4); + for (int j = 0; j < RP_JOBS; ++j) { + p.W7[j] = base ? base + off : nullptr; + off += (size_t)native_RP_COUTP(j) * native_RP_CIN(j) * 7; + } + p.W3i = nullptr; + if (use_rp3i(B, F)) { + p.W3i = base ? base + off + (size_t)3 * 1536 * 768 : nullptr; + off += (size_t)7 * 3 * 1536 * 768; + } + p.W3s = nullptr; + if (use_rp3s(B, F)) { + p.W3s = base ? base + off + (size_t)3 * 768 * 1536 : nullptr; + off += (size_t)7 * 3 * 768 * 1536; + } + p.total = off; + return p; +} +__host__ __device__ const __half *H16(const void *p) { return static_cast(p); } +__host__ __device__ const float *F32(const void *p) { return static_cast(p); } +} +size_t codec_scratch_bytes(int B, int F) { + return build_plan(B, F, nullptr).total * sizeof(__half); +} +template +__device__ __forceinline__ void wg_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, bool prep = false) { + constexpr int WG_BM = 64 * NWG, WG_NT = 128 * NWG, WG_CLD = WG_BN + 8; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int GEMM_B = STAGES * (A_E + B_E) * 2; + constexpr int EPI_B = WG_BM * WG_CLD * 4; + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute(wg_conv_kernel, SH); + return; + } + native_invoke(wg_conv_kernel,dim3(dim3(DIVUP(Cout, WG_BM), DIVUP(st.N, WG_BN))),dim3(WG_NT),SH, + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg2_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, + bool prep = false) { + constexpr int BK = CIB * KT; + constexpr int WG_BM = 64 * MSPLIT; + constexpr int WG_NT = 128 * MSPLIT * KSPLIT * NSPLIT; + constexpr int WG_BNT = WG_BN * NSPLIT; + constexpr int NCRAW = WG_BNT / 4; + constexpr int EPN = + (WG_NT % NCRAW == 0) ? WG_BNT : (((WG_BNT + 127) / 128) * 128); + constexpr int WG_CLD = EPN + 8; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int GEMM_B = KSPLIT * (SA * A_E + SB * NSPLIT * B_E) * 2; + constexpr int EPI_B = KSPLIT * WG_BM * WG_CLD * 4; + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute( + wg2_conv_kernel, + SH); + return; + } + native_invoke_pf(wg2_conv_kernel,dim3(dim3(DIVUP(Cout, WG_BM), DIVUP(st.N, WG_BNT))),dim3(WG_NT),SH, + w, (size_t)KT * Cin * Cout * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg3_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, + bool prep = false) { + constexpr int WG_BM = 64 * NWG, WG_NT = 128 * NWG * NSG; + constexpr int WG_BNT = WG_BN * NSG, WG_CLD = WG_BNT + 8; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int GEMM_B = STAGES * (A_E + NSG * B_E) * 2; + constexpr int EPI_B = WG_BM * WG_CLD * 4; + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute(wg3_conv_kernel, SH); + return; + } + native_invoke(wg3_conv_kernel,dim3(dim3(DIVUP(Cout, WG_BM), DIVUP(st.N, WG_BNT))),dim3(WG_NT),SH, + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg4_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, int CoutP, + const Stack &st, int SCin, int do_snake, int pad_off, + int pad_len, bool prep = false) { + constexpr int NT = 128 * MSPLIT * NSPLIT; + constexpr int BM = 64 * MSPLIT, BNT = WG_BN * NSPLIT; + constexpr int NBR = BNT + KT - 1; + constexpr int NBRP = BC8 ? ((NBR + 7) & ~7) : ((NBR + 31) & ~31); + constexpr int NCRAW = BNT / 4; + constexpr int EPN = + (NT % NCRAW == 0) ? BNT : (((BNT + 127) / 128) * 128); + constexpr int B_E = CIB * NBRP; + constexpr int GEMM_B = STAGES * B_E * 2; + constexpr int EPI_B = BM * (EPN + 8) * (PREACT ? 2 : 4); + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute( + wg4_conv_kernel, + SH); + return; + } + pdl_launch_pf(wg4_conv_kernel, + dim3(DIVUP(Cout, BM), DIVUP(st.N, BNT)), dim3(NT), SH, s, + w, (size_t)KT * CoutP * Cin * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, + st.T, st.L, SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg5_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, + bool prep = false) { + constexpr int BM = 64 * MSPLIT; + constexpr int NT = 128 * MSPLIT * KSPLIT; + constexpr int NCRAW = WG_BN / 4; + constexpr int EPN = (NT % NCRAW == 0) ? WG_BN : (((WG_BN + 127) / 128) * 128); + constexpr int NBRP = BC8 ? (((WG_BN + 2) + 7) & ~7) + : (((WG_BN + 2) + 31) & ~31); + constexpr int B_E = CIB * NBRP; + constexpr int GEMM_B = KSPLIT * SB * B_E * 2; + constexpr int EPI_COLS = OC8 ? WG_BN : EPN; + constexpr int EPI_CLD_R = EPI_COLS + 8; + constexpr int EPI_CLD = (EPI_CLD_R % 16 == 0) ? EPI_CLD_R + 4 : EPI_CLD_R; + constexpr int EPI_B = KSPLIT * BM * EPI_CLD * (OC8 ? 2 : 4); + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + static_assert(SH <= 227 * 1024, "wg5 shared over budget"); + if (prep) { + prepare_attribute( + wg5_conv_kernel, + SH); + return; + } + if constexpr (PDL) { + cudaLaunchConfig_t cfg = {}; + cudaLaunchAttribute attr[1] = {}; + native_config_grid(cfg) = dim3(DIVUP(Cout, BM), DIVUP(st.N, WG_BN)); + native_config_threads(cfg) = dim3(NT); + cfg.dynamicSmemBytes = SH; + cfg.stream = s; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = 1; + cfg.attrs = attr; + cfg.numAttrs = 1; + native_cudaLaunchKernelEx_pf( + &cfg, + wg5_conv_kernel, + w, (size_t)3 * Cin * Cout * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); + } else { + native_invoke_pf(wg5_conv_kernel,dim3(dim3(DIVUP(Cout, BM), DIVUP(st.N, WG_BN))),dim3(NT),SH, + w, (size_t)3 * Cin * Cout * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, + st.L, SCin, st.SC, st.N, do_snake, pad_off, pad_len); + } +} +#ifndef D7W +#define D7W 96 +#endif +#define D7_M2_C8I(BN,LATE) wg4_launch<7, 16, 2, 1, D7W, 4, 1, false, 2, true, false, LATE>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_M2P_C8O(BN,LATE) wg4_launch<7, 16, 2, 1, D7W, 4, 1, true, 2, false, true, LATE>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_M2N2(BN) wg4_launch<7, 16, 2, 2, BN, 4, 1, false, 2>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_M1(BN) wg4_launch<7, 16, 1, 1, BN, 4, 2, false, 2>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_C8I(W) \ + wg4_launch<7, 16, 2, 1, W, 4, 1, false, 2, true, false, false>( \ + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, \ + SCin, 1, pad_off, pad_len, prep) +#define D7_C8O(W) \ + wg4_launch<7, 16, 2, 1, W, 4, 1, true, 2, false, true, false>( \ + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, \ + SCin, 1, pad_off, pad_len, prep) +template +__device__ __forceinline__ int d7_band_cost(int N, int mtiles, int resident) { + const int v = mtiles * DIVUP(N, BN); + const int sub = (DIVUP(v, resident) >= 2) ? 2 : 1; + return DIVUP(DIVUP(v, sub), resident) * sub * (128 + BN + 6); +} +__device__ __forceinline__ int d7_band(int N, int mtiles, int resident) { + int best = 64, bc = d7_band_cost<64>(N, mtiles, resident); + const int c96 = d7_band_cost<96>(N, mtiles, resident); + if (c96 < bc) { bc = c96; best = 96; } + if (d7_band_cost<112>(N, mtiles, resident) < bc) best = 112; + return best; +} +__device__ __forceinline__ void dense7(cudaStream_t s, int job, const __half *in, const __half *w, + const float *bias, const float *alpha, const __half *resid, + __half *out, const __half *pad_src, __half *ns, + const Stack &st, int SCin, int pad_off, int pad_len, + bool prep = false) { + const int Cin = native_RP_CIN(job), Cout = native_RP_COUT(job), CoutP = native_RP_COUTP(job); + const int b_ = d7_band(st.N, DIVUP(Cout, 128), native_resident()); + if ((job & 1) == 0) { + if (b_ == 64) D7_C8O(64); + else if (b_ == 96) D7_C8O(96); + else D7_C8O(112); + } else { + if (b_ == 64) D7_C8I(64); + else if (b_ == 96) D7_C8I(96); + else D7_C8I(112); + } +} +__device__ __forceinline__ int wg5_band_skip(int N, int mtiles, int resident) { + const int wmin = DIVUP(mtiles * DIVUP(N, 104), resident); + if (DIVUP(mtiles * DIVUP(N, 24), resident) == wmin) return 24; + if (DIVUP(mtiles * DIVUP(N, 40), resident) == wmin) return 40; + if (DIVUP(mtiles * DIVUP(N, 48), resident) == wmin) return 48; + if (DIVUP(mtiles * DIVUP(N, 64), resident) == wmin) return 64; + return 104; +} +__device__ __forceinline__ int wg5_band_main_long(int N, int mtiles, + int resident) { + const int wmin = DIVUP(mtiles * DIVUP(N, 104), resident); + if (DIVUP(mtiles * DIVUP(N, 64), resident) == wmin) return 64; + return 104; +} +__device__ __forceinline__ int wg5_band_main_short(int N, int mtiles, + int resident) { + const int wmin = DIVUP(mtiles * DIVUP(N, 96), resident); + if (DIVUP(mtiles * DIVUP(N, 56), resident) == wmin) return 56; + if (DIVUP(mtiles * DIVUP(N, 80), resident) == wmin) return 80; + return 96; +} +#define DENSE_LAUNCH(KT,BM,BN,NW,WM,WN) native_invoke(dense_conv_kernel,dim3(dim3(DIVUP(Cout, BM), DIVUP(st.N, BN))),dim3(NW * 32),0, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, SCin, st.SC, st.N, do_snake, pad_off, pad_len) +__device__ __forceinline__ void dense(cudaStream_t s, int KT, const __half *in, const __half *w, + const float *bias, const float *alpha, const __half *resid, + __half *out, const __half *pad_src, __half *ns, int Cin, + int Cout, const Stack &st, int SCin, int do_snake, + int pad_off, int pad_len, const __half *w3 = nullptr, + int layout_mode = 0, bool prep = false) { + if (KT == 3 && w3 != nullptr) { + if (Cin == 1536) { +#define WG5S(BN,O,R,PDL) wg5_launch<16, 1, 4, BN, 4, 3, 1, 1, true, O, R, 1536, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5SW(BN,O,R,PDL) wg5_launch<32, 1, 2, BN, 4, 3, 1, 1, true, O, R, 1536, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5S3(O,R) do { const int n_ = nref(st); if (n_ > 500 && n_ <= 640) { WG5SW(56, O, R, false); } else { const int b_ = wg5_band_skip(st.N, DIVUP(Cout, 64), native_resident()); if (b_ == 24) WG5S(24, O, R, true); else if (b_ == 40) WG5S(40, O, R, true); else if (b_ == 48) WG5S(48, O, R, true); else if (b_ == 64) WG5S(64, O, R, true); else WG5S(104, O, R, true); } } while (0) + if (layout_mode == 1) WG5S3(true, false); + else if (layout_mode == 2) WG5S3(true, true ); + else if (layout_mode == 3) WG5S3(false, true ); + else WG5S3(false, false); +#undef WG5S3 +#undef WG5SW +#undef WG5S + } else { +#define WG5I(BN,BC,MS,PDL) wg5_launch<16, MS, 2, BN, 4, 3, 1, 1, BC, true, false, 768, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5I4(BN,BC,PDL) wg5_launch<16, 1, 4, BN, 4, 3, 1, 1, BC, true, false, 768, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5I2W(BN,BC,PDL) wg5_launch<32, 1, 2, BN, 4, 3, 1, 1, BC, true, false, 768, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5I3(BC) do { const int n_ = nref(st); if (n_ > 640) { const int b_ = wg5_band_main_long(st.N, DIVUP(Cout, 128), native_resident()); if (b_ == 64) WG5I(64, BC, 2, true); else WG5I(104, BC, 2, true); } else if (n_ > 500) { WG5I2W(64, BC, false); } else { const int b_ = wg5_band_main_short(st.N, DIVUP(Cout, 64), native_resident()); if (b_ == 56) WG5I4(56, BC, true); else if (b_ == 80) WG5I4(80, BC, true); else WG5I4(96, BC, true); } } while (0) + if (layout_mode == 1) WG5I3(true); + else WG5I3(false); +#undef WG5I3 +#undef WG5I2W +#undef WG5I4 +#undef WG5I + } + } else if (KT == 3) { + const int n_ = nref(st); + if (Cin == 1536 && Cout == 768) { + if (n_ <= 704) + wg2_launch<3, 16, 1, 4, 1, 64, 2, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else + wg2_launch<3, 16, 1, 4, 1, 128, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + } else if (Cin == 768 && Cout == 1536) { + if (n_ > 1000) + wg2_launch<3, 16, 1, 2, 1, 128, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ > 640 && n_ <= 768) + wg2_launch<3, 16, 1, 2, 1, 64, 2, 2, 2>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ <= 320) + wg2_launch<3, 16, 1, 4, 1, 64, 2, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ > 384 && n_ <= 512) + wg2_launch<3, 16, 1, 4, 1, 112, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ > 512 && n_ <= 640) + wg2_launch<3, 16, 1, 4, 1, 120, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else + wg2_launch<3, 16, 1, 4, 1, 128, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + } else { + wg2_launch<3, 16, 1, 1, 1, 64, 3, 2, 2>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, + SCin, do_snake, pad_off, pad_len, prep); + } + } else if (Cout >= 128 && + ((Cin >= 384 && st.T == 288) || + (Cin < 384 && st.T == 2016))) { + wg3_launch<7, 112, 2, 2, 1, 128, 2>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, + SCin, do_snake, pad_off, pad_len); + } else { + wg3_launch<7, 112, 1, 4, 1, 128, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, + SCin, do_snake, pad_off, pad_len); + } +} +__device__ void codec_program(const CodecArgs &a, const Plan &p, const RepackSet &js, int repack_grid) { + cudaStream_t s = nullptr; + const int B = a.B; + const void *const *W = a.inp; + const __half *const *ST = reinterpret_cast(a.st); + __half *const *NS = reinterpret_cast<__half *const *>(a.ns); + { + native_invoke(pack_latent_kernel,dim3(DIVUP(40*B*p.s0.L,256)),dim3(256),0, + H16(W[0]), ST[0], p.LB, 40, p.s0.T, p.s0.L, p.s0.SC, 2, B); + } + dense(s, 3, p.LB, H16(W[1]), F32(W[2]), nullptr, nullptr, p.X0, ST[1], NS[0], + 40, 768, p.s0, p.s0.SC, 0, 0, 2); + const __half *k3iw = H16(W[3]); + const float *k3ib = F32(W[4]); + const float *k3ia = F32(W[5]); + const __half *k3sw = H16(W[6]); + const float *k3sb = F32(W[7]); + const float *k3sa = F32(W[8]); + const size_t IW = (size_t)1536 * 768 * 3, SW = (size_t)768 * 1536 * 3; + dense(s, 3, p.X0, k3iw, k3ib, k3ia, nullptr, p.M0, ST[2], NS[1], 768, 1536, + p.s0, p.s0.SC, 1, 0, 2, p.W3i ? p.W3i - IW : nullptr, 0); + dense(s, 3, p.M0, k3sw, k3sb, k3sa, p.X0, p.Y0, ST[3], NS[2], 1536, 768, + p.s0, p.s0.SC, 1, 1, 1, p.W3s ? p.W3s - SW : nullptr, 0); + TCONV_BPACK_DISPATCH8(1, 2, 8, p.s0.T, 768, B, + p.Y0, H16(W[9]), F32(W[10]), F32(W[11]), p.A, ST[4], NS[3], 768, 768, + p.s0.T, p.s0.L, p.s0.SC, 1, p.s1.L, p.s1.SC, 2, 0, 2, B); + if (p.W3i != nullptr) + native_invoke(repack_activation_only,dim3(min(1152,96*DIVUP(p.s1.SC,128))),dim3(128),0, + p.A, p.Ac, p.s1.SC, p.s1.SC, DIVUP(p.s1.SC,128)); + __half *cur = p.A, *nxt = p.Bb; + for (int i = 0; i < 6; ++i) { + const int bi = i + 1; + const __half *ain = (i == 0 && p.Ac) ? p.Ac : cur; + dense(s, 3, ain, k3iw + bi * IW, k3ib + bi * 1536, k3ia + bi * 768, + nullptr, p.M1, ST[5 + 2 * i], NS[4 + 2 * i], 768, 1536, p.s1, + p.s1.SC, 1, 0, 2, p.W3i ? p.W3i + (bi - 1) * IW : nullptr, + (i > 0 || p.Ac) ? 1 : 0); + const __half *ps = (i < 5) ? ST[6 + 2 * i] : ST[16]; + const int po = (i < 5) ? 0 : 1, pl = 1 + (i < 5); + dense(s, 3, p.M1, k3sw + bi * SW, k3sb + bi * 768, k3sa + bi * 384, cur, + nxt, ps, NS[5 + 2 * i], 1536, 768, p.s1, p.s1.SC, 1, po, pl, + p.W3s ? p.W3s + (bi - 1) * SW : nullptr, + i == 0 ? 1 : (i < 5 ? 2 : 3)); + __half *t = cur; cur = nxt; nxt = t; + } + TCONV_BPACK_DISPATCH1(2, 9, 16, p.s1.T, 384, B, + cur, H16(W[12]), F32(W[13]), F32(W[14]), p.P, ST[17], NS[16], 768, 384, + p.s1.T, p.s1.L, p.s1.SC, 1, p.s2.L, p.s2.SC, 6, 0, 6, B); + dense7(s, 0, p.P, p.W7[0], F32(W[22]), F32(W[23]), nullptr, p.M2, ST[18], + NS[17], p.s2, p.s2.SC, 0, 6); + dense7(s, 1, p.M2, p.W7[1], F32(W[25]), F32(W[26]), p.P, p.Q, ST[19], + NS[18], p.s2, p.s2.SC, 5, 1); + TCONV_BPACK_DISPATCH(3, 7, 16, p.s2.T, 128, B, + p.Q, H16(W[15]), F32(W[16]), F32(W[17]), p.R, ST[20], NS[19], 384, 128, + p.s2.T, p.s2.L, p.s2.SC, 5, p.s3.L, p.s3.SC, 6, 0, 6, B); + dense7(s, 2, p.R, p.W7[2], F32(W[28]), F32(W[29]), nullptr, p.M3, ST[21], + NS[20], p.s3, p.s3.SC, 0, 6); + dense7(s, 3, p.M3, p.W7[3], F32(W[31]), F32(W[32]), p.R, p.S, ST[22], + NS[21], p.s3, p.s3.SC, 5, 1); + if (p.s3.T >= 1500) + pdl_launch(transpose_conv_kernel<4, 7, 8, true>, dim3(1, 32, B), + dim3(tc_blk(p.s3.T, 8)), + (size_t)tc_blk(p.s3.T, 8) * (8 * 7) * 2 + 64, s, + p.S, H16(W[18]), F32(W[19]), F32(W[20]), p.U, ST[23], + NS[22], 128, 32, p.s3.T, p.s3.L, p.s3.SC, 5, p.s4.L, + p.s4.SC, 6, 0, 6, B); + else + pdl_launch(transpose_conv_kernel<4, 7, 4, true>, dim3(1, 32, B), + dim3(tc_blk(p.s3.T, 4)), + (size_t)tc_blk(p.s3.T, 4) * (4 * 7) * 2 + 64, s, + p.S, H16(W[18]), F32(W[19]), F32(W[20]), p.U, ST[23], + NS[22], 128, 32, p.s3.T, p.s3.L, p.s3.SC, 5, p.s4.L, + p.s4.SC, 6, 0, 6, B); + { + const int ntile = DIVUP(p.s4.N, TL_BN); + const int resident = native_resident(); + const int cap = resident > 0 ? resident : 132; + const int grid = ntile < cap ? ntile : cap; + pdl_launch_pf(tail_fused_kernel, dim3(grid), dim3(TL_NT), TL_SHARED, s, + p.U, (size_t)32 * p.s4.SC * sizeof(__half), + p.U, p.W7[4], p.W7[5], H16(W[39]), F32(W[34]), F32(W[35]), + F32(W[37]), F32(W[38]), F32(W[40]), ST[24], ST[25], + static_cast<__half *>(a.waveform), NS[23], NS[24], NS[25], B, + p.s4.T, p.s4.L, p.s4.SC, p.s4.N, ntile); + } +} diff --git a/deploy/nemotron_native/kernels/b32/native/codec/native_stage1.cu b/deploy/nemotron_native/kernels/b32/native/codec/native_stage1.cu new file mode 100644 index 0000000..e7ca5ce --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/codec/native_stage1.cu @@ -0,0 +1,329 @@ +#include "native_stage1.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#ifndef NATIVE_CODEC_PREFETCH +#define NATIVE_CODEC_PREFETCH 1 +#endif +#ifndef NATIVE_CODEC_PACK_SUBCTAS +#define NATIVE_CODEC_PACK_SUBCTAS 1 +#endif +namespace cg=cooperative_groups; +constexpr int PHYSICAL_THREADS=512,MAIN_ARENA=227840,AUX_BYTES=4096,META_OFFSET=231936,TOTAL_SHARED=232192; +constexpr int SLICE_BYTES=MAIN_ARENA+AUX_BYTES; +constexpr int MAX_SUBCTA=8; +struct NativeMeta { + dim3 grid,threads; + uint3 block[MAX_SUBCTA]; + NativeStage1Status *status;const int32_t *lengths;int padded_frames; + int submask,tshift,tmask,slice; + int resident; +}; +static_assert(sizeof(NativeMeta)<=TOTAL_SHARED-META_OFFSET); +__device__ __forceinline__ char *native_base(){extern __shared__ __align__(16) char arena[];return arena;} +__device__ __forceinline__ NativeMeta *native_meta(){return reinterpret_cast(native_base()+META_OFFSET);} +__device__ __forceinline__ int native_subid(){const NativeMeta *m=native_meta();return (int)((threadIdx.x>>m->tshift)&(unsigned)m->submask);} +__device__ __forceinline__ char *native_arena(){const NativeMeta *m=native_meta();return native_base()+native_subid()*m->slice;} +__device__ __forceinline__ char *native_aux(){const NativeMeta *m=native_meta();return native_base()+native_subid()*m->slice+(m->slice-AUX_BYTES);} +__device__ __forceinline__ uint3 native_tid(){uint3 r;r.x=threadIdx.x&(unsigned)native_meta()->tmask;r.y=0u;r.z=0u;return r;} +__device__ __forceinline__ dim3 native_block(){const uint3 b=native_meta()->block[native_subid()];return dim3(b.x,b.y,b.z);} +__device__ __forceinline__ dim3 native_grid(){return native_meta()->grid;} +__device__ __forceinline__ dim3 native_threads(){return native_meta()->threads;} +__device__ __forceinline__ int native_resident(){return native_meta()->resident;} +__device__ __forceinline__ void native_local_sync(){ + const NativeMeta *m=native_meta();unsigned n=m->threads.x,id=15u-(unsigned)native_subid(); + asm volatile("barrier.sync %0, %1;"::"r"(id),"r"(n):"memory"); +} +__device__ __forceinline__ dim3 &native_config_grid(cudaLaunchConfig_t &c){return c.gridDim;} +__device__ __forceinline__ dim3 &native_config_threads(cudaLaunchConfig_t &c){return c.blockDim;} +__device__ __forceinline__ int native_true_tokens(int row,int padded_tokens){ + const NativeMeta *m=native_meta(); + return (padded_tokens/m->padded_frames)*m->lengths[row]; +} +__device__ __forceinline__ void native_warm(const char *p,size_t n,int k,int nk,int me){ + if(p==nullptr||n<16||nk<=0)return; + const size_t a=(size_t)p&15u;p-=a;n=(n+a)&~(size_t)15; + const uint4 *q=reinterpret_cast(p); + const size_t cnt=n>>4,step=(size_t)nk*PHYSICAL_THREADS; + uint32_t acc=0u; + for(size_t i=(size_t)k*PHYSICAL_THREADS+(size_t)me;i __device__ __forceinline__ void native_invoke_pf(F fn,dim3 g,dim3 t,size_t shared,const void *pf,size_t pfb,A&&... args){ + const int me=threadIdx.x,owner=blockIdx.x,owners=gridDim.x; + const cg::grid_group all=cg::this_grid();NativeMeta *meta=native_meta(); + if(t.x>512||t.y!=1||t.z!=1||t.x%32||shared>MAIN_ARENA){ + if(me==0&&owner==0)meta->status->code=10; + all.sync();return; + } + int sub=1,tshift=9; + const unsigned tx=t.x; + const int vtotal=(int)(g.x*g.y*g.z); + if(NATIVE_CODEC_PACK_SUBCTAS && (tx&(tx-1))==0u){ + int s=0;while((1u<MAX_SUBCTA)cand=MAX_SUBCTA; + const int want=(vtotal+owners-1)/owners; + while(cand>1&&cand>want)cand>>=1; + while(cand>1&&shared>(size_t)(SLICE_BYTES/cand-AUX_BYTES))cand>>=1; + sub=cand; + } + const int slice=(SLICE_BYTES/sub)&~127; + if(me==0){meta->grid=g;meta->threads=t;meta->submask=sub-1;meta->tshift=tshift; + meta->tmask=(int)((1u<slice=slice;} + const int total=vtotal,grp=me>>tshift; + const bool holder=me=busy&&pf!=nullptr)native_warm((const char*)pf,pfb,owner-busy,owners-busy,me); + for(int base=owner*sub;baseblock[grp]=make_uint3(i%g.x,(i/g.x)%g.y,i/(g.x*g.y)); + __syncthreads(); + if(holder&&i(args)...); + asm volatile("cp.async.wait_all;":::"memory"); + if(t.x>=128)asm volatile("wgmma.wait_group.sync.aligned 0;":::"memory"); + } + __syncthreads(); + } + all.sync(); +} +template __device__ __forceinline__ void native_invoke(F fn,dim3 g,dim3 t,size_t shared,A&&... args){ + native_invoke_pf(fn,g,t,shared,nullptr,(size_t)0,static_cast(args)...); +} +template __device__ __forceinline__ void native_cudaLaunchKernelEx(cudaLaunchConfig_t *c,F fn,A... args){ + native_invoke(fn,c->gridDim,c->blockDim,c->dynamicSmemBytes,args...); +} +template __device__ __forceinline__ void native_cudaLaunchKernelEx_pf(cudaLaunchConfig_t *c,F fn,const void *pf,size_t pfb,A... args){ + native_invoke_pf(fn,c->gridDim,c->blockDim,c->dynamicSmemBytes,pf,pfb,args...); +} +#define blockIdx native_block() +#define gridDim native_grid() +#define blockDim native_threads() +#define threadIdx native_tid() +#define __syncthreads() native_local_sync() +#include "legacy_codec.cuh" +#undef __syncthreads +#undef threadIdx +#undef blockDim +#undef gridDim +#undef blockIdx +namespace { +struct FastArgs { + CodecArgs codec;Plan plan;RepackSet repack;int repack_grid;int input_frames;int blocks; + const int64_t *codes;const int32_t *lengths,*slots; + const __half *state[26];__half *next[26];float *pcm;int32_t *samples; + NativeStage1Status *status;int nslots; + int channels[26],histories[26]; +}; +static_assert(sizeof(FastArgs)<=4096,"precomputed kernel argument must fit the conservative parameter budget"); +struct Handle {FastArgs args{};Plan plans[17]{};RepackSet repacks[17]{};int repack_grids[17]{};void *scratch=nullptr;size_t bytes=0;void *immutable=nullptr;size_t immutable_bytes=0;int setup_launches=0;std::vector packed;int max_frames=0,blocks=0,device=-1;}; +thread_local std::string error; +bool check(cudaError_t e,const char *where){if(e==cudaSuccess)return true;error=std::string(where)+": "+cudaGetErrorString(e);return false;} +__device__ void reject(NativeStage1Status *s,int code,int row,int frame,int detail){ + if(atomicCAS(&s->code,0,code)==0){s->row=row;s->frame=frame;s->detail=detail;} +} +__global__ __launch_bounds__(512,1) void stage1_fast_persistent(const __grid_constant__ FastArgs a){ + const auto all=cg::this_grid();size_t me=(size_t)blockIdx.x*512+threadIdx.x,step=(size_t)gridDim.x*512; + if(threadIdx.x==0){native_meta()->status=a.status;native_meta()->lengths=a.lengths;native_meta()->padded_frames=a.codec.F;native_meta()->resident=a.blocks;} + if(me==0)*a.status=NativeStage1Status{0,-1,-1,0}; + all.sync();const int B=a.codec.B,F=a.codec.F,inputF=a.input_frames; + for(int row=me;rowinputF)reject(a.status,8,row,-1,n); + if(slot<0||slot>=a.nslots)reject(a.status,2,row,-1,slot); + for(int j=0;j=0&&n<=inputF)?n*1764:0; + } + for(size_t i=me;i<(size_t)B*inputF;i+=step){ + int row=i/inputF,t=i%inputF;bool any=false; + for(int c=0;c<16;++c){int64_t v=a.codes[i*16+c];if(v<0||v>=1024)reject(a.status,4,row,t,c);any|=v!=0;} + if(t=a.lengths[row]&&any)reject(a.status,6,row,t,0); + } + if(a.nslots!=B) + for(int p=0;p<26;++p) + for(size_t i=me;i<(size_t)a.nslots*a.channels[p]*a.histories[p];i+=step) + a.next[p][i]=a.state[p][i]; + for(size_t i=me;i<(size_t)B*inputF*1764;i+=step)a.pcm[i]=0.f; + all.sync(); + if(a.status->code){ + if(a.nslots==B) + for(int p=0;p<26;++p) + for(size_t i=me;i<(size_t)B*a.channels[p]*a.histories[p];i+=step) + a.next[p][i]=a.state[p][i]; + return; + } + for(int p=0;p<26;++p){ + const size_t stride=(size_t)a.channels[p]*a.histories[p]; + __half *dest=(__half*)a.codec.st[p]; + for(size_t i=me;i<(size_t)B*stride;i+=step)dest[i]=a.state[p][(size_t)a.slots[i/stride]*stride+i%stride]; + } + __half *latent=(__half*)a.codec.inp[0];int T=2*F; + for(size_t i=me;i<(size_t)B*40*T;i+=step){ + int t=i%T,c=(i/T)%40,row=i/((size_t)40*T);float v=0.f; + if(t/2>(2*(c%5)))&3)-2)*.5f;} + latent[i]=__float2half_rn(v); + } + all.sync(); + codec_program(a.codec,a.plan,a.repack,a.repack_grid); + if(a.status->code){ + if(a.nslots==B) + for(int p=0;p<26;++p) + for(size_t i=me;i<(size_t)B*a.channels[p]*a.histories[p];i+=step) + a.next[p][i]=a.state[p][i]; + return; + } + for(int p=0;p<26;++p){ + const size_t stride=(size_t)a.channels[p]*a.histories[p];const __half *src=(const __half*)a.codec.ns[p]; + for(size_t i=me;i<(size_t)B*stride;i+=step){ + const int row=i/stride;const size_t dst=(size_t)a.slots[row]*stride+i%stride; + if(a.lengths[row])a.next[p][dst]=src[i]; + else if(a.nslots==B)a.next[p][dst]=a.state[p][dst]; + } + } + const __half *wave=(const __half*)a.codec.waveform; + for(size_t i=me;i<(size_t)B*inputF*1764;i+=step){ + const int row=i/(inputF*1764),t=i%(inputF*1764); + if(tstatus=a.status;native_meta()->resident=a.blocks;} + if(blockIdx.x==0&&threadIdx.x==0)*a.status=NativeStage1Status{0,-1,-1,0}; + all.sync(); + native_invoke(repack_rs_kernel<7>,dim3(a.repack_grid),dim3(128),MAIN_ARENA, + a.repack,H16(a.codec.inp[0]),H16(a.codec.st[0]),a.plan.LB, + 40,a.plan.s0.T,a.plan.s0.L,a.plan.s0.SC,2,a.codec.B); + native_invoke(repack3_rs_kernel,dim3(1152,12),dim3(128),MAIN_ARENA, + H16(a.codec.inp[3]),H16(a.codec.inp[6]),a.plan.W3i,a.plan.W3s, + a.plan.A,a.plan.Ac,a.plan.s1.SC,a.plan.s1.SC,DIVUP(a.plan.s1.SC,128)); +} +bool stack(Handle *h,int slot,const NativeStage1Weights *w,int first,int field,size_t elements){ + const int ids[7]={first,first+3,first+5,first+7,first+9,first+11,first+13}; + const size_t stride=elements*(field==0?sizeof(__half):sizeof(float));void *ptr=nullptr; + if(!check(cudaMalloc(&ptr,7*stride),"packed-weight allocation"))return false; + h->packed.push_back(ptr);h->args.codec.inp[slot]=ptr; + for(int i=0;i<7;++i){const void *src=field==0?w[ids[i]].weight:field==1?(const void*)w[ids[i]].bias:(const void*)w[ids[i]].alpha; + if(!check(cudaMemcpy((char*)ptr+i*stride,src,stride,cudaMemcpyDeviceToDevice),"immutable weight stack"))return false;} + return true; +} +void release(Handle *h){if(h->immutable)cudaFree(h->immutable);if(h->scratch)cudaFree(h->scratch);for(void *p:h->packed)cudaFree(p);delete h;} +} +extern "C" void *native_stage1_create(int b,int maxf,int slots,int blocks,const NativeStage1Weights *w){ + error.clear();if((b!=1&&b!=16&&b!=32&&b!=64)||maxf<1||maxf>16||slots4096||blocks<0||!w){error="invalid create arguments";return nullptr;} + for(int p=0;p<26;++p)if(!w[p].weight||!w[p].bias||(p!=0&&p!=25&&!w[p].alpha)){error="missing weight pointer";return nullptr;} + Handle *h=new(std::nothrow)Handle;if(!h){error="host allocation";return nullptr;} + h->max_frames=maxf;h->args.nslots=slots;h->args.codec.B=b; + const int cs[26]={40,768,1536,768,768,1536,768,1536,768,1536,768,1536,768,1536,768,1536,768,384,768,384,128,256,128,32,64,32}; + const int hs[26]={2,2,2,1,2,2,2,2,2,2,2,2,2,2,2,2,1,6,6,1,6,6,1,6,6,6}; + cudaDeviceProp prop{};int occupancy=0,coop=0; + if(!check(cudaGetDevice(&h->device),"device")||!check(cudaGetDeviceProperties(&prop,h->device),"properties")|| + !check(cudaDeviceGetAttribute(&coop,cudaDevAttrCooperativeLaunch,h->device),"cooperative support")|| + !check(cudaFuncSetAttribute(stage1_fast_persistent,cudaFuncAttributeMaxDynamicSharedMemorySize,TOTAL_SHARED),"shared-memory optin")|| + !check(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&occupancy,stage1_fast_persistent,512,TOTAL_SHARED),"occupancy")) {release(h);return nullptr;} + if(prop.major!=9||prop.minor!=0||!coop||occupancy<1){error="requires resident SM90 cooperative kernel";release(h);return nullptr;} + h->blocks=blocks?blocks:occupancy*prop.multiProcessorCount; + if(h->blocks>occupancy*prop.multiProcessorCount){error="grid exceeds cooperative occupancy";release(h);return nullptr;} + h->args.blocks=h->blocks; + const int capacityF=maxf<=3?3:(maxf<=5?5:std::max(8,maxf)); + size_t decoder=codec_scratch_bytes(b,capacityF),halfs=(size_t)b*40*2*capacityF+(size_t)b*1764*capacityF; + for(int p=0;p<26;++p){h->args.channels[p]=cs[p];h->args.histories[p]=hs[p];halfs+=2*(size_t)b*cs[p]*hs[p];} + h->bytes=decoder+halfs*sizeof(__half); + if(!check(cudaMalloc(&h->scratch,h->bytes),"scratch allocation")){release(h);return nullptr;} + if(!check(cudaMemset(h->scratch,0,h->bytes),"initial scratch padding")){release(h);return nullptr;} + h->args.codec.scratch=h->scratch;__half *q=(__half*)((char*)h->scratch+decoder); + h->args.codec.inp[0]=q;q+=(size_t)b*40*2*capacityF;h->args.codec.waveform=q;q+=(size_t)b*1764*capacityF; + for(int p=0;p<26;++p){h->args.codec.st[p]=q;q+=(size_t)b*cs[p]*hs[p];h->args.codec.ns[p]=q;q+=(size_t)b*cs[p]*hs[p];} + const void **in=h->args.codec.inp; + in[1]=w[0].weight;in[2]=w[0].bias; + if(!stack(h,3,w,1,0,(size_t)1536*768*3)||!stack(h,4,w,1,1,1536)||!stack(h,5,w,1,2,768)|| + !stack(h,6,w,2,0,(size_t)768*1536*3)||!stack(h,7,w,2,1,768)||!stack(h,8,w,2,2,384)){release(h);return nullptr;} + const int upops[4]={3,16,19,22};for(int i=0;i<4;++i){in[9+3*i]=w[upops[i]].weight;in[10+3*i]=w[upops[i]].bias;in[11+3*i]=w[upops[i]].alpha;} + for(int i=0;i<3;++i){int p=17+3*i;in[21+6*i]=w[p].weight;in[22+6*i]=w[p].bias;in[23+6*i]=w[p].alpha; + in[24+6*i]=w[p+1].weight;in[25+6*i]=w[p+1].bias;in[26+6*i]=w[p+1].alpha;} + in[39]=w[25].weight;in[40]=w[25].bias; + __half *immutable_w7[6]{};__half *immutable_w3i=nullptr,*immutable_w3s=nullptr; + if(capacityF>=3){ + size_t wh=0;for(int j=0;j<6;++j)wh+=(size_t)native_RP_COUTP(j)*native_RP_CIN(j)*7; + const size_t IW=(size_t)3*1536*768;wh+=14*IW; + h->immutable_bytes=wh*sizeof(__half); + if(!check(cudaMalloc(&h->immutable,h->immutable_bytes+sizeof(NativeStage1Status)),"immutable weight allocation")){release(h);return nullptr;} + if(!check(cudaMemset(h->immutable,0,h->immutable_bytes+sizeof(NativeStage1Status)),"immutable padding initialization")){release(h);return nullptr;} + __half *q=(__half*)h->immutable; + for(int j=0;j<6;++j){immutable_w7[j]=q;q+=(size_t)native_RP_COUTP(j)*native_RP_CIN(j)*7;} + immutable_w3i=q+IW;q+=7*IW;immutable_w3s=q+IW;q+=7*IW; + if((char*)q!=(char*)h->immutable+h->immutable_bytes){error="immutable layout size mismatch";release(h);return nullptr;} + } + for(int f:{3,5,8}) if(f<=capacityF){ + Plan p=build_plan(b,f,static_cast<__half*>(h->scratch)); + for(int j=0;j<6;++j)p.W7[j]=immutable_w7[j];p.W3i=immutable_w3i;p.W3s=immutable_w3s; + const void *const *W=h->args.codec.inp; + RepackSet js{}; + const void *SRC[RP_JOBS] = {W[21], W[24], W[27], W[30], W[33], W[36]}; + for (int j = 0; j < RP_JOBS; ++j) { + js.src[j] = H16(SRC[j]); + js.dst[j] = p.W7[j]; + js.Cout[j] = native_RP_COUT(j); + js.CoutP[j] = native_RP_COUTP(j); + js.Cin[j] = native_RP_CIN(j); + } + int acc = 0; + for (int j = 0; j < RP_JOBS; ++j) { + js.base[j] = acc; + js.nch[j] = native_RP_CIN(j) / 16; + js.rs[j] = (j < 4) ? 1 : 0; + acc += (native_RP_COUTP(j) / 16) * DIVUP(native_RP_CIN(j), 64); + } + js.k3base = acc; + js.k3src[0] = H16(W[3]); + js.k3src[1] = H16(W[6]); + js.k3dst[0] = p.W3i ? p.W3i - (size_t)3 * 1536 * 768 : nullptr; + js.k3dst[1] = p.W3s ? p.W3s - (size_t)3 * 768 * 1536 : nullptr; + const int k3n = p.W3i ? 2304 : 0; + h->plans[f]=p;h->repacks[f]=js;h->repack_grids[f]=acc+k3n; + } + if(capacityF>=3){ + FastArgs a=h->args;a.codec.F=3;a.plan=h->plans[3];a.repack=h->repacks[3];a.repack_grid=h->repack_grids[3]; + a.status=(NativeStage1Status*)((char*)h->immutable+h->immutable_bytes); + int setup_occupancy=0; + if(!check(cudaFuncSetAttribute(stage1_prepare_immutable,cudaFuncAttributeMaxDynamicSharedMemorySize,TOTAL_SHARED),"setup shared-memory optin")|| + !check(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&setup_occupancy,stage1_prepare_immutable,512,TOTAL_SHARED),"setup occupancy")|| + setup_occupancy<1||h->blocks>setup_occupancy*prop.multiProcessorCount){ + if(error.empty())error="setup grid exceeds cooperative occupancy";release(h);return nullptr;} + void *params[]={&a}; + if(!check(cudaLaunchCooperativeKernel((void*)stage1_prepare_immutable,dim3(h->blocks),dim3(512),params,TOTAL_SHARED,nullptr),"immutable setup launch")|| + !check(cudaStreamSynchronize(nullptr),"immutable setup completion")){release(h);return nullptr;} + NativeStage1Status status{}; + if(!check(cudaMemcpy(&status,a.status,sizeof(status),cudaMemcpyDeviceToHost),"immutable setup status")){release(h);return nullptr;} + if(status.code){error="immutable setup device status "+std::to_string(status.code);release(h);return nullptr;} + h->setup_launches=1; + } + return h; +} +extern "C" int native_stage1_run(void *handle,const int64_t *codes,const int32_t *lengths,const int32_t *slots, + const void *const *states,void *const *next,float *pcm,int32_t *samples,NativeStage1Status *status,int f,void *stream){ + error.clear();Handle *h=(Handle*)handle; + if(!h||!codes||!lengths||!slots||!states||!next||!pcm||!samples||!status){error="null run argument";return 1;} + if(f<1||f>8||f>h->max_frames){error="true-state variant supports public F1..8";return 1;} + const int paddedF=f<=3?3:(f<=5?5:8); + int device=-1;if(!check(cudaGetDevice(&device),"device")||device!=h->device){error="device mismatch";return 1;} + FastArgs a=h->args;a.input_frames=f;a.plan=h->plans[paddedF];a.repack=h->repacks[paddedF];a.repack_grid=h->repack_grids[paddedF];a.codec.F=paddedF;a.codes=codes;a.lengths=lengths;a.slots=slots;a.pcm=pcm;a.samples=samples;a.status=status; + for(int p=0;p<26;++p){if(!states[p]||!next[p]||states[p]==next[p]){error="invalid/aliased state";return 1;}a.state[p]=(const __half*)states[p];a.next[p]=(__half*)next[p];} + void *params[]={&a};if(!check(cudaLaunchCooperativeKernel((void*)stage1_fast_persistent,dim3(h->blocks),dim3(512),params,TOTAL_SHARED,(cudaStream_t)stream),"fast Stage1 launch"))return 1; + return check(cudaPeekAtLastError(),"immediate launch status")?0:1; +} +extern "C" int native_stage1_destroy(void *p){error.clear();if(!p)return 0;release((Handle*)p);return check(cudaPeekAtLastError(),"destroy status")?0:1;} +extern "C" const char *native_stage1_last_error(){return error.c_str();} +extern "C" size_t native_stage1_scratch_bytes(void *p){return p?((Handle*)p)->bytes:0;} +extern "C" int native_stage1_resident_blocks(void *p){return p?((Handle*)p)->blocks:0;} +extern "C" int native_stage1_setup_kernel_launches(void *p){return p?((Handle*)p)->setup_launches:0;} +extern "C" size_t native_stage1_immutable_weight_bytes(void *p){return p?((Handle*)p)->immutable_bytes:0;} diff --git a/deploy/nemotron_native/kernels/b32/native/codec/native_stage1.h b/deploy/nemotron_native/kernels/b32/native/codec/native_stage1.h new file mode 100644 index 0000000..a95b5f0 --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/codec/native_stage1.h @@ -0,0 +1,26 @@ +#pragma once +#include +#include +struct NativeStage1Weights { const void *weight; const float *bias; const float *alpha; }; +struct NativeStage1Status { int32_t code, row, frame, detail; }; +#ifdef __cplusplus +extern "C" { +#endif +void *native_stage1_create(int batch, int max_frames, int state_slots, + int resident_blocks, const NativeStage1Weights *weights26); +int native_stage1_run(void *handle, const int64_t *codes, const int32_t *lengths, + const int32_t *slots, const void *const *states26, + void *const *next_states26, float *pcm, int32_t *samples, + NativeStage1Status *status, int frames, void *stream); +int native_stage1_run_profile(void *handle, const int64_t *codes, const int32_t *lengths, + const int32_t *slots, const void *const *states26, + void *const *next_states26, float *pcm, int32_t *samples, + NativeStage1Status *status, int frames, void *stream, + uint64_t *phase_timestamps30); +int native_stage1_destroy(void *handle); +const char *native_stage1_last_error(); +size_t native_stage1_scratch_bytes(void *handle); +int native_stage1_resident_blocks(void *handle); +#ifdef __cplusplus +} +#endif diff --git a/deploy/nemotron_native/kernels/b32/native/codec/repack3.inc b/deploy/nemotron_native/kernels/b32/native/codec/repack3.inc new file mode 100644 index 0000000..2f61c93 --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/codec/repack3.inc @@ -0,0 +1,66 @@ +__device__ __noinline__ void repack3_kernel( + const __half *__restrict__ win, const __half *__restrict__ wsk, + __half *__restrict__ din, __half *__restrict__ dsk) { + int4 (*sm)[49]=reinterpret_cast(native_arena()); + const int job = blockIdx.y; + const int Cout = job ? 768 : 1536; + const int Cin = job ? 1536 : 768; + const __half *src = job ? wsk : win; + __half *dst = job ? dsk : din; + const int GPB = Cin >> 3; + const int Kg = Cin * 3; + const int GT = GPB >> 4; + const int RT = Cout >> 5; + const int t = threadIdx.x; + for (int tile = blockIdx.x; tile < 6 * RT * GT; tile += gridDim.x) { + const int gblk = tile % GT; + const int rest = tile / GT; + const int corow = rest % RT; + const int b = rest / RT; + const int co0 = corow * 32; + const int g0 = gblk * 16; + const __half *sbase = src + (size_t)(b + 1) * Cout * Kg + + (size_t)co0 * Kg + g0 * 24; +#pragma unroll + for (int j = 0; j < 6; ++j) { + const int idx = t + j * 256; + const int r = idx / 48, c = idx - (idx / 48) * 48; + sm[r][c] = *reinterpret_cast(sbase + (size_t)r * Kg + + c * 8); + } + __syncthreads(); + __half *dbase = dst + (size_t)b * (3 * (size_t)Cout * Cin); +#pragma unroll + for (int p = 0; p < 2; ++p) { + const int u = t + p * 256; + const int col = (u & 7) + 8 * (u >> 7); + const int g = g0 + ((u >> 3) & 15); + const int co = co0 + col; + const int c = g >> 1, v = g & 1; + const int4 a = sm[col][((u >> 3) & 15) * 3]; + const int4 e = sm[col][((u >> 3) & 15) * 3 + 1]; + const int4 f = sm[col][((u >> 3) & 15) * 3 + 2]; + uint32_t h[12]; + h[0] = a.x; h[1] = a.y; h[2] = a.z; h[3] = a.w; + h[4] = e.x; h[5] = e.y; h[6] = e.z; h[7] = e.w; + h[8] = f.x; h[9] = f.y; h[10] = f.z; h[11] = f.w; + __half *d = dbase + (size_t)c * (48 * Cout) + + (size_t)(co >> 3) * 128 + v * 64 + (co & 7) * 8; +#pragma unroll + for (int tp = 0; tp < 3; ++tp) { + uint32_t o[4]; +#pragma unroll + for (int jj = 0; jj < 4; ++jj) { + const uint32_t g0w = h[3 * jj], g1w = h[3 * jj + 1], + g2w = h[3 * jj + 2]; + o[jj] = (tp == 0) ? __byte_perm(g0w, g1w, 0x7610) + : (tp == 1) ? __byte_perm(g0w, g2w, 0x5432) + : __byte_perm(g1w, g2w, 0x7610); + } + *reinterpret_cast(d + (size_t)tp * (16 * Cout)) = + make_int4(o[0], o[1], o[2], o[3]); + } + } + __syncthreads(); + } +} diff --git a/deploy/nemotron_native/kernels/b32/native/codec/tail_kernel.inc b/deploy/nemotron_native/kernels/b32/native/codec/tail_kernel.inc new file mode 100644 index 0000000..8794f42 --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/codec/tail_kernel.inc @@ -0,0 +1,492 @@ +#define TL_NT 512 +#define TL_MW 512 +#define TL_MLD TL_MW +#define TL_VW 504 +#define TL_UW 520 +#define TL_BN 496 +#define TL_VLD 520 +#define TL_LBO_U (TL_UW * 8) +#define TL_LBO_M (TL_MLD * 8) +#define TL_LBO_V (TL_VLD * 8) +#define TL_SBO_A 128 +#define TL_PU_E (4 * TL_UW * 8) +#define TL_PM_E (8 * TL_MLD * 8) +#define TL_PV_E (32 * TL_VLD) +#define TL_AI_E (14 * 1024) +#define TL_AS_E (32 * 512 + 512) +#define TL_SHARED ((2 * TL_PU_E + TL_PM_E + TL_PV_E + TL_AI_E + TL_AS_E) * 2) +__device__ __forceinline__ void tl_stmatrix4(void *dst, const uint32_t r[4], + bool trans) { + const uint32_t a = static_cast(__cvta_generic_to_shared(dst)); + if (trans) + asm volatile( + "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1,%2,%3,%4};\n" ::"r"( + a), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3])); + else + asm volatile( + "stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1,%2,%3,%4};\n" ::"r"( + a), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3])); +} +__device__ __forceinline__ void tl_ldmatrix4t(uint32_t r[4], const void *src) { + const uint32_t a = static_cast(__cvta_generic_to_shared(src)); + asm volatile( + "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(a)); +} +__device__ __forceinline__ void tl_stmatrix1t(void *dst, uint32_t r) { + const uint32_t a = static_cast(__cvta_generic_to_shared(dst)); + asm volatile( + "stmatrix.sync.aligned.m8n8.x1.trans.shared.b16 [%0], {%1};\n" :: + "r"(a), "r"(r)); +} +template +__device__ __forceinline__ void tl_skip_band( + const __half *PU, __half *PM, const __half *AS, + __half *PV, int c0, int warp, int lane, int er0, const float *sh_ask, + const float *sh_ari_v, float b0v, float b1v) { + constexpr int NB8 = BW / 8; + float acc[BW / 2]; +#pragma unroll + for (int i = 0; i < BW / 2; ++i) acc[i] = ((i & 2) == 0) ? b0v : b1v; + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int c = 0; c < 4; ++c) { + const uint64_t da0 = + wg_desc(AS + c * 8 * 512, 128u, TL_SBO_A * 2u); + const uint64_t db0 = + wg_desc(PM + (2 * c) * TL_LBO_M + (c0 + 2) * 8, + TL_LBO_M * 2u, 128u); +#pragma unroll + for (int p = 0; p < 4; ++p) { + const uint64_t da = da0 + 64 * (2 * p); + const uint64_t db = db0 + 2 * p; + if constexpr (BW == 120) { + WGMMA_M64N120K16_T(acc, da, db, 0, 0); + } else { + WGMMA_M64N128K16_T(acc, da, db, 0, 0); + } + } + } + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + const int m_addr = lane >> 3, row = lane & 7; + const int cg = (warp & 1) * 2; + if (er0 >= 32) { +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = __floats2half2_rn( + acc[jj * 4 + h * 2 + 0], acc[jj * 4 + h * 2 + 1]); + r[m] = *reinterpret_cast(&hv); + } else { + r[m] = 0; + } + } + const int jj = j + (m_addr >> 1), h = m_addr & 1; + tl_stmatrix4(PV + (cg + h) * TL_LBO_V + (c0 + jj * 8 + row) * 8, + r, true); + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + if (er0 < 32) { +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; + const int jja = j + (m_addr >> 1), ha = m_addr & 1; + tl_ldmatrix4t(r, PV + (cg + ha) * TL_LBO_V + + (c0 + jja * 8 + row + 1) * 8); +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = *reinterpret_cast<__half2 *>(&r[m]); + acc[jj * 4 + h * 2 + 0] += __half2float(hv.x); + acc[jj * 4 + h * 2 + 1] += __half2float(hv.y); + } + } + } +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; + const int jja = j + (m_addr >> 1), ha = m_addr & 1; + tl_ldmatrix4t(r, PU + (cg + ha) * TL_LBO_U + + (c0 + jja * 8 + row + 16) * 8); +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = *reinterpret_cast<__half2 *>(&r[m]); + acc[jj * 4 + h * 2 + 0] += __half2float(hv.x); + acc[jj * 4 + h * 2 + 1] += __half2float(hv.y); + } + } + } + } + if (er0 >= 32) return; + if (er0 < 16) { + const float a0 = sh_ask[er0], r0 = sh_ari_v[er0]; + const float a1 = sh_ask[er0 + 8], r1 = sh_ari_v[er0 + 8]; +#pragma unroll + for (int i = 0; i < BW / 2; ++i) { + const int h = (i & 2) >> 1; + const float v = acc[i]; + const float sn = __sinf((h ? a1 : a0) * v); + acc[i] = __fmaf_rn(sn * sn, h ? r1 : r0, v); + } + } else { +#pragma unroll + for (int i = 0; i < BW / 2; ++i) { + const float v = acc[i]; + acc[i] = v > 0.f ? v : 0.01f * v; + } + } +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = __floats2half2_rn( + acc[jj * 4 + h * 2 + 0], acc[jj * 4 + h * 2 + 1]); + r[m] = *reinterpret_cast(&hv); + } else { + r[m] = 0; + } + } + const int jj = j + (m_addr >> 1), h = m_addr & 1; + tl_stmatrix4(PM + (cg + h) * TL_LBO_V + + (c0 + jj * 8 + row) * 8, + r, true); + } +} +__device__ __noinline__ void tail_fused_kernel( + const __half *__restrict__ U, const __half *__restrict__ win, + const __half *__restrict__ wsk, const __half *__restrict__ wpo, + const float *__restrict__ bin, const float *__restrict__ ain, + const float *__restrict__ bsk, const float *__restrict__ ask, + const float *__restrict__ bpo, const __half *__restrict__ st_m4, + const __half *__restrict__ st_v, __half *__restrict__ wave, + __half *__restrict__ ns_u, __half *__restrict__ ns_m4, + __half *__restrict__ ns_v, int B, int T4, int L4, int SC4, int N4, + int NTILE) { + char *tl_arena=reinterpret_cast(native_arena()); + __half *const PU = reinterpret_cast<__half *>(tl_arena); + __half *const PM = PU + 2 * TL_PU_E; + __half *const PV = PM + TL_PM_E; + __half *const AI = PV + TL_PV_E; + __half *const AS = AI + TL_AI_E; + __half *const PB = AS + 32 * 512; + float *sh_ain=(float*)native_aux(), *sh_ari=sh_ain+32, *sh_ask=sh_ari+32, *sh_ari_v=sh_ask+16; + const int tid = threadIdx.x; + const int wg = tid >> 7; + const int warp = (tid & 127) >> 5; + const int lane = tid & 31; + const int er0 = warp * 16 + (lane >> 2); + const int cb2 = (lane & 3) * 2; + for (int i = tid; i < 32; i += TL_NT) { + const float a = ain[i]; + sh_ain[i] = a; + sh_ari[i] = __frcp_rn(a + 1e-9f); + } + for (int i = tid; i < 16; i += TL_NT) { + const float a = ask[i]; + sh_ask[i] = a; + sh_ari_v[i] = __frcp_rn(a + 1e-9f); + } + for (int i = tid; i < TL_AI_E / 8; i += TL_NT) { + const int blk = i / 128, off = (i - blk * 128) * 8; + const int c = blk / 7, k = blk - c * 7; + *reinterpret_cast(AI + blk * 1024 + off) = + *reinterpret_cast(win + (size_t)c * 14336 + + (size_t)k * 2048 + off); + } + for (int i = tid; i < 32 * 64; i += TL_NT) { + const int blk = i / 64, off = (i - blk * 64) * 8; + const int c = blk >> 3, k = blk & 7; + *reinterpret_cast(AS + blk * 512 + off) = + (k < 7) ? *reinterpret_cast( + wsk + (size_t)c * 14336 + (size_t)k * 2048 + off) + : make_int4(0, 0, 0, 0); + } + for (int i = tid; i < 4 * (TL_VLD - TL_VW) * 8; i += TL_NT) { + const int g = i / ((TL_VLD - TL_VW) * 8); + PV[g * TL_LBO_V + TL_VW * 8 + (i - g * ((TL_VLD - TL_VW) * 8))] = + __float2half(0.f); + } + for (int i = tid; i < 512; i += TL_NT) PB[i] = __float2half(0.f); + for (int i = tid; i < 2 * 16 * 8; i += TL_NT) { + const int c = i / (16 * 8); + const int r = i - c * 16 * 8; + const int ki = r / 8, k = r & 7; + PB[c * 128 + (ki >> 3) * 64 + (ki & 7) * 8 + k] = + (k < 7) ? wpo[(c * 16 + ki) * 7 + k] : __float2half(0.f); + } + const float bias_i0 = bin[er0], bias_i1 = bin[er0 + 8]; + const float bias_s0 = (er0 < 32) ? bsk[er0] : 0.f; + const float bias_s1 = (er0 + 8 < 32) ? bsk[er0 + 8] : 0.f; + const float bpo0 = bpo[0]; + constexpr int UCOL = TL_UW / 32; + constexpr int UUNIT = 4 * UCOL; + constexpr int UWARP = TL_NT / 32; + constexpr int UITER = (UUNIT + UWARP - 1) / UWARP; + const int uwarp = tid >> 5; + const int uq = lane & 3; + int ugc[UITER], uoc[UITER]; +#pragma unroll + for (int t = 0; t < UITER; ++t) { + const int bi = uwarp + t * UWARP; + const int gc = bi / UCOL; + ugc[t] = gc; + uoc[t] = bi - gc * UCOL; + } + const __half *const urow = + U + (size_t)(lane >> 2) * SC4 + 8 * uq; + __syncthreads(); + PDL_ENTER(); +#define TL_STAGE_U(PUD,UB) { _Pragma("unroll") for (int t2 = 0; t2 < UITER; t2 += 2) { uint32_t u[2][4]; _Pragma("unroll") for (int d = 0; d < 2; ++d) { const int t = t2 + d; const int n = (UB) + uoc[t] * 32; const int gn = n + 8 * uq; const bool ok = gn >= 0 && gn < N4; int4 v = make_int4(0, 0, 0, 0); if (ok) v = *reinterpret_cast( urow + (size_t)ugc[t] * (8 * (size_t)SC4) + n); u[d][0] = v.x; u[d][1] = v.y; u[d][2] = v.z; u[d][3] = v.w; } _Pragma("unroll") for (int d = 0; d < 2; ++d) { const int t = t2 + d; uint32_t *q = u[d]; uint32_t x, y; x = (uq & 1) ? q[0] : q[1]; y = (uq & 1) ? q[2] : q[3]; x = __shfl_xor_sync(0xffffffffu, x, 1); y = __shfl_xor_sync(0xffffffffu, y, 1); if (uq & 1) { q[0] = x; q[2] = y; } else { q[1] = x; q[3] = y; } x = (uq & 2) ? q[0] : q[2]; y = (uq & 2) ? q[1] : q[3]; x = __shfl_xor_sync(0xffffffffu, x, 2); y = __shfl_xor_sync(0xffffffffu, y, 2); if (uq & 2) { q[0] = x; q[1] = y; } else { q[2] = x; q[3] = y; } tl_stmatrix4((PUD) + ugc[t] * TL_LBO_U + (uoc[t] * 32 + lane) * 8, q, true); } } if (uwarp < 4) { const int n = (UB) + 512; const int gn = n + 2 * (lane & 3); uint32_t q = 0; if (gn >= 0 && gn + 1 < N4) q = *reinterpret_cast( U + (size_t)(uwarp * 8 + (lane >> 2)) * SC4 + gn); tl_stmatrix1t((PUD) + uwarp * TL_LBO_U + (512 + (lane & 7)) * 8, q); } } + TL_STAGE_U(PU, (int)blockIdx.x * TL_BN - 16) + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + int ubuf = 0; + for (int tile = blockIdx.x; tile < NTILE; tile += gridDim.x, ubuf ^= 1) { + const int g0 = tile * TL_BN; + __half *const PUc = PU + ubuf * TL_PU_E; + __half *const PUn = PU + (ubuf ^ 1) * TL_PU_E; + { + const int b0 = g0 / L4; + for (int s = 0; s < 2; ++s) { + const int sb = b0 + s; + if (sb >= B) break; + const int cv = (int)((long long)sb * L4 + native_true_tokens(sb,T4) - g0); + if (cv < 0 || cv >= TL_BN) continue; + for (int i = tid; i < 32 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + ns_u[((size_t)sb * 32 + co) * 6 + h] = + PUc[(co >> 3) * TL_LBO_U + + (cv + h + 16) * 8 + (co & 7)]; + } + } + } + { + float acc[64]; +#pragma unroll + for (int i = 0; i < 64; ++i) + acc[i] = ((i & 2) == 0) ? bias_i0 : bias_i1; + const int c0 = wg * 128; + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int c = 0; c < 2; ++c) { + const uint64_t da0 = + wg_desc(AI + c * 7 * 1024, 128u, TL_SBO_A * 2u); + const uint64_t db0 = + wg_desc(PUc + (2 * c) * TL_LBO_U + (c0 + 2) * 8, + TL_LBO_U * 2u, 128u); +#pragma unroll + for (int k = 0; k < 7; ++k) { + const uint64_t da = da0 + 128 * k; + const uint64_t db = db0 + k; + WGMMA_M64N128K16_T(acc, da, db, 0, 0); + } + } + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); + if (tile + (int)gridDim.x < NTILE) + TL_STAGE_U(PUn, (tile + (int)gridDim.x) * TL_BN - 16) + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + if (er0 < 32) { + const float a0 = sh_ain[er0], r0 = sh_ari[er0]; + const float a1 = sh_ain[er0 + 8], r1 = sh_ari[er0 + 8]; +#pragma unroll + for (int i = 0; i < 64; ++i) { + const int h = (i & 2) >> 1; + const float v = acc[i]; + const float sn = __sinf((h ? a1 : a0) * v); + acc[i] = __fmaf_rn(sn * sn, h ? r1 : r0, v); + } + } else { +#pragma unroll + for (int i = 0; i < 64; ++i) { + const float v = acc[i]; + acc[i] = v > 0.f ? v : 0.01f * v; + } + } +#pragma unroll + for (int j = 0; j < 16; j += 2) { + uint32_t r[4]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + const __half2 hv = __floats2half2_rn( + acc[jj * 4 + h * 2 + 0], acc[jj * 4 + h * 2 + 1]); + r[m] = *reinterpret_cast(&hv); + } + const int m_addr = lane >> 3, row = lane & 7; + const int jj = j + (m_addr >> 1), h = m_addr & 1; + tl_stmatrix4(PM + (warp * 2 + h) * TL_LBO_M + + (c0 + jj * 8 + row) * 8, + r, true); + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + { + const int base = g0 - 8; + const int b0 = (base >= 0) ? base / L4 : 0; + const int cm0 = (int)((long long)b0 * L4 - base); + const int cm1 = cm0 + L4; + const bool hit0 = (b0 < B) && (cm0 + 6 > 0) && (cm0 < TL_MW); + const bool hit1 = (b0 + 1 < B) && (cm1 + 6 > 0) && (cm1 < TL_MW); + if (hit0 || hit1) { + for (int s = 0; s < 2; ++s) { + if (!(s ? hit1 : hit0)) continue; + const int sb = b0 + s, cm = s ? cm1 : cm0; + for (int i = tid; i < 64 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + const int c = cm + h; + if (c >= 0 && c < TL_MW) + PM[(co >> 3) * TL_LBO_M + c * 8 + (co & 7)] = + st_m4[((size_t)sb * 64 + co) * 6 + h]; + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + } + } + { + const int b0 = g0 / L4; + for (int s = 0; s < 2; ++s) { + const int sb = b0 + s; + if (sb >= B) break; + const int cv = (int)((long long)sb * L4 + native_true_tokens(sb,T4) - g0); + if (cv < 0 || cv >= TL_BN) continue; + for (int i = tid; i < 64 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + ns_m4[((size_t)sb * 64 + co) * 6 + h] = + PM[(co >> 3) * TL_LBO_M + + (cv + h + 8) * 8 + (co & 7)]; + } + } + } + { + const int c0 = wg * 128; + if (wg == 3) + tl_skip_band<120>(PUc, PM, AS, PV, c0, warp, lane, er0, sh_ask, + sh_ari_v, bias_s0, bias_s1); + else + tl_skip_band<128>(PUc, PM, AS, PV, c0, warp, lane, er0, sh_ask, + sh_ari_v, bias_s0, bias_s1); + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + { + const int b0 = g0 / L4; + const int cv0 = (int)((long long)b0 * L4 - g0), cv1 = cv0 + L4; + const bool hv0 = (b0 < B) && (cv0 + 6 > 0) && (cv0 < TL_VW); + const bool hv1 = (b0 + 1 < B) && (cv1 + 6 > 0) && (cv1 < TL_VW); + if (hv0 || hv1) { + for (int s = 0; s < 2; ++s) { + if (!(s ? hv1 : hv0)) continue; + const int sb = b0 + s, cv = s ? cv1 : cv0; + for (int i = tid; i < 32 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + const int c = cv + h; + if (c >= 0 && c < TL_VW) + PM[(co >> 3) * TL_LBO_V + c * 8 + (co & 7)] = + st_v[((size_t)sb * 32 + co) * 6 + h]; + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + } + } + { + const int b0 = g0 / L4; + for (int s = 0; s < 2; ++s) { + const int sb = b0 + s; + if (sb >= B) break; + const int cv = (int)((long long)sb * L4 + native_true_tokens(sb,T4) - g0); + if (cv < 0 || cv >= TL_BN) continue; + for (int i = tid; i < 32 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + ns_v[((size_t)sb * 32 + co) * 6 + h] = + PM[(co >> 3) * TL_LBO_V + (cv + h) * 8 + (co & 7)]; + } + } + } + __syncthreads(); + constexpr int PC_LD = 10; + { + float *const PC = reinterpret_cast(PV); +#pragma unroll + for (int round = 0; round < 2; ++round) { + float acc[4]; + const int p0 = round * 256 + wg * 64; + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); + const uint64_t da0 = + wg_desc(PM + p0 * 8, TL_LBO_V * 2u, 128u); + const uint64_t db0 = wg_desc(PB, 128u, 128u); + const uint64_t da1 = + wg_desc(PM + 2 * TL_LBO_V + p0 * 8, + TL_LBO_V * 2u, 128u); + const uint64_t db1 = wg_desc(PB + 128, 128u, 128u); + WGMMA_M64N8K16_ZERO_T(acc, da0, db0, 0, 1); + WGMMA_M64N8K16_T(acc, da1, db1, 0, 1); + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + const int row0 = p0 + warp * 16 + (lane >> 2); + const int col = (lane & 3) * 2; + *reinterpret_cast(PC + row0 * PC_LD + col) = + make_float2(acc[0], acc[1]); + *reinterpret_cast(PC + (row0 + 8) * PC_LD + col) = + make_float2(acc[2], acc[3]); + } + } + __syncthreads(); + { + const float *const PC = reinterpret_cast(PV); + const int p0 = tid; + const bool live = p0 < TL_BN; + const unsigned live_mask = __ballot_sync(0xffffffffu, live); + if (live) { + const float s0 = __fadd_rn(PC[(p0 + 0) * PC_LD + 0], + PC[(p0 + 1) * PC_LD + 1]); + const float s1 = __fadd_rn(PC[(p0 + 2) * PC_LD + 2], + PC[(p0 + 3) * PC_LD + 3]); + const float s2 = __fadd_rn(PC[(p0 + 4) * PC_LD + 4], + PC[(p0 + 5) * PC_LD + 5]); + const float s3 = PC[(p0 + 6) * PC_LD + 6]; + float y0 = __fadd_rn(__fadd_rn(s0, s1), + __fadd_rn(s2, s3)); + y0 = fminf(fmaxf(y0 + bpo0, -1.f), 1.f); + const float y1 = __shfl_down_sync(live_mask, y0, 1); + if ((p0 & 1) == 0) { + const int b0 = g0 / L4; + const int rr = g0 - b0 * L4; + int t0 = rr + p0, bo0 = b0; + int t1 = t0 + 1, bo1 = bo0; + if (t0 >= L4) { t0 -= L4; ++bo0; } + if (t1 >= L4) { t1 -= L4; ++bo1; } + if (bo0 == bo1 && bo0 < B && t1 < T4) { + *reinterpret_cast<__half2 *>( + wave + (size_t)bo0 * T4 + t0) = + __floats2half2_rn(y0, y1); + } else { + if (bo0 < B && t0 < T4) + wave[(size_t)bo0 * T4 + t0] = __float2half(y0); + if (bo1 < B && t1 < T4) + wave[(size_t)bo1 * T4 + t1] = __float2half(y1); + } + } + } + } + } +#undef TL_STAGE_U +} diff --git a/deploy/nemotron_native/kernels/b32/native/codec/wg2_kernel.inc b/deploy/nemotron_native/kernels/b32/native/codec/wg2_kernel.inc new file mode 100644 index 0000000..ebe059e --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/codec/wg2_kernel.inc @@ -0,0 +1,280 @@ +__device__ __forceinline__ void cpa16(void *dsm, const void *src, bool pred) { + unsigned a = static_cast(__cvta_generic_to_shared(dsm)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(a), + "l"(src), "r"(pred ? 16 : 0) + : "memory"); +} +template +__device__ __forceinline__ void cpa16o(unsigned a, const void *src, bool pred) { + asm volatile("cp.async.cg.shared.global [%0+%3], [%1], 16, %2;\n" ::"r"(a), + "l"(src), "r"(pred ? 16 : 0), "n"(OFF) + : "memory"); +} +template +__device__ __noinline__ +void wg2_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int BK = CIB * KT; + constexpr int WG_BM = 64 * MSPLIT; + constexpr int WG_BNT = WG_BN * NSPLIT; + constexpr int NTS = 128 * MSPLIT * NSPLIT; + constexpr int WG_NT = NTS * KSPLIT; + constexpr int NCRAW = WG_BNT / 4; + constexpr int EPN = + (WG_NT % NCRAW == 0) ? WG_BNT : (((WG_BNT + 127) / 128) * 128); + constexpr int WG_CLD = EPN + 8; + constexpr bool NB8_P2 = ((WG_BN / 8) & (WG_BN / 8 - 1)) == 0; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int ST_E = SA * A_E + SB * NSPLIT * B_E; + constexpr int SBO_AE = KG * WG_SLOT; + constexpr int LBO_BE = (WG_BN / 8) * WG_SLOT; + constexpr int NB8 = WG_BN / 8; + constexpr bool ALDG = (WG_BN <= 64); + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *smem = reinterpret_cast<__half *>(arena); + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int kgrp = tid / NTS; + const int stid = tid - kgrp * NTS; + const int wgl = stid >> 7; + const int mg = wgl % MSPLIT; + const int ngid = wgl / MSPLIT; + const int warp = (stid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * WG_BM; + const int nb = blockIdx.y * WG_BNT; + const int n0 = nb + ngid * WG_BN; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + const int NCH = (Kg + BK - 1) / BK; + if (tid < WG_BM) { + const int co = m0 + tid; + if (do_snake && co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + for (int i = tid; i < EPN; i += WG_NT) { + const int n = nb + i; + colok[i] = (i < WG_BNT && n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mg * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (kgrp == 0 && eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (kgrp == 0 && eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + constexpr int KGX = ALDG ? (KG <= 8 ? 8 : (KG <= 16 ? 16 : 32)) : KG; + constexpr int ATOT = WG_BM * KGX; + constexpr int AREG = (ATOT + NTS - 1) / NTS; + constexpr int BTOT = NSPLIT * CIB * NB8; + constexpr int BREG = (BTOT + NTS - 1) / NTS; + const __half *aptr[AREG]; + int aoff[AREG]; + int akmax[AREG]; +#pragma unroll + for (int r = 0; r < AREG; ++r) { + const int idx = stid + r * NTS; + const int m = idx / KGX, v = idx - (idx / KGX) * KGX; + const int co = m0 + m; + const bool live = (ATOT % NTS == 0 || idx < ATOT) && (co < Cout) && + v < KG; + aoff[r] = (m >> 3) * SBO_AE + (v < KG ? v : 0) * WG_SLOT + (m & 7) * 8; + akmax[r] = live ? (Kg - v * 8) : 0; + aptr[r] = w + (size_t)(live ? co : 0) * Kg + v * 8; + } + const __half *bptr[BREG]; + int bcimax[BREG], bnc[BREG], bcil[BREG], bgrp[BREG]; +#pragma unroll + for (int r = 0; r < BREG; ++r) { + const int idx = stid + r * NTS; + const int gg = idx / (CIB * NB8); + const int rr = idx - gg * (CIB * NB8); + const int cil = rr / NB8, nc = rr - (rr / NB8) * NB8; + bcil[r] = cil; + bnc[r] = nc; + bgrp[r] = gg; + const bool live = (BTOT % NTS == 0 || idx < BTOT); + bcimax[r] = live ? (Cin - cil) : 0; + bptr[r] = in + (size_t)cil * SCin + nb + gg * WG_BN + nc * 8; + } + int4 br0[BREG], br1[BREG]; + int4 ar[ALDG ? AREG : 1]; + uint32_t edge[BREG]; + __half *const kaB = smem + kgrp * ST_E; + __half *const kbB = kaB + SA * A_E; + const unsigned barid = 1u + kgrp + (unsigned)(KSPLIT * native_subid()); +#define WG2_BAR() asm volatile("barrier.sync %0, %1;" ::"r"(barid), "r"(NTS) : "memory") +#define WG2_LOADB(c) { const int _ci0 = (c) * CIB; _Pragma("unroll") for (int r = 0; r < BREG; ++r) { br0[r] = make_int4(0, 0, 0, 0); if constexpr (KT == 3) edge[r] = 0; else br1[r] = make_int4(0, 0, 0, 0); if (_ci0 < bcimax[r]) { const int4 *src = reinterpret_cast( bptr[r] + (size_t)_ci0 * SCin); br0[r] = src[0]; if constexpr (KT == 3) { if (!NB8_P2 || bnc[r] == NB8 - 1) edge[r] = src[1].x; } else { br1[r] = src[1]; } } } } +#define WG2_LOADA(c) if constexpr (ALDG) { const int _kg0 = (c) * BK; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { ar[r] = make_int4(0, 0, 0, 0); if (_kg0 < akmax[r]) ar[r] = *reinterpret_cast(aptr[r] + _kg0); } } +#define WG2_PUTA(c,AD) { _Pragma("unroll") for (int r = 0; r < AREG; ++r) { if constexpr (ALDG) { if (akmax[r] > 0) *reinterpret_cast((AD) + aoff[r]) = ar[r]; } else { const int _kg0 = (c) * BK; if (akmax[r] > 0) cpa16((AD) + aoff[r], aptr[r] + _kg0, _kg0 < akmax[r]); } } if constexpr (!ALDG) asm volatile("cp.async.commit_group;\n" ::: "memory"); } +#define WG2_STSB(BD) { _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int cil = bcil[r], nc = bnc[r]; unsigned int u[8]; u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; if constexpr (KT == 3) { if constexpr (NB8_P2) { u[4] = __shfl_down_sync(0xffffffffu, u[0], 1, NB8); if (nc == NB8 - 1) u[4] = edge[r]; } else { u[4] = edge[r]; } u[5] = u[6] = u[7] = 0; } else { u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; } if (BTOT % NTS == 0 || (stid + r * NTS) < BTOT) { _Pragma("unroll") for (int k = 0; k < KT; ++k) { const int kh = k >> 1; int4 o; if (k & 1) { o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); } else { o.x = u[kh + 0]; o.y = u[kh + 1]; o.z = u[kh + 2]; o.w = u[kh + 3]; } const int row = cil * KT + k; *reinterpret_cast((BD) + bgrp[r] * B_E + nc * WG_SLOT + (row >> 3) * LBO_BE + (row & 7) * 8) = o; } } } } +#define WG2_MMA(AD,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t da0 = wg_desc((AD) + mg * 8 * SBO_AE, WG_SLOT * 2, SBO_AE * 2); const uint64_t db0 = wg_desc((BD) + ngid * B_E, LBO_BE * 2, WG_SLOT * 2); _Pragma("unroll") for (int kk = 0; kk < BK; kk += 16) { const uint64_t da = da0 + (kk >> 3) * (WG_SLOT / 8); const uint64_t db = db0 + (kk >> 3) * (LBO_BE / 8); if constexpr (WG_BN == 64) { WGMMA_M64N64K16(acc, da, db); } else if constexpr (WG_BN == 256) { WGMMA_M64N256K16(acc, da, db); } else if constexpr (WG_BN == 104) { WGMMA_M64N104K16(acc, da, db); } else if constexpr (WG_BN == 112) { WGMMA_M64N112K16(acc, da, db); } else if constexpr (WG_BN == 120) { WGMMA_M64N120K16(acc, da, db); } else if constexpr (WG_BN == 224) { WGMMA_M64N224K16(acc, da, db); } else { static_assert(WG_BN == 128, "no wgmma macro for WG_BN"); WGMMA_M64N128K16(acc, da, db); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + constexpr int SMAX = SA > SB ? SA : SB; + WG2_LOADA(kgrp) + WG2_LOADB(kgrp) +#pragma unroll 1 + for (int s = 0; s < SMAX; ++s) { + const int c = s * KSPLIT + kgrp; + if (s < SA) { + if (c < NCH) WG2_PUTA(c, kaB + s * A_E) + if (c + KSPLIT < NCH) WG2_LOADA(c + KSPLIT) + } + if (s < SB) { + if (c < NCH) WG2_STSB(kbB + s * NSPLIT * B_E) + if (c + KSPLIT < NCH) WG2_LOADB(c + KSPLIT) + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + const int NIT = (NCH - kgrp + KSPLIT - 1) / KSPLIT; + for (int i = 0; i < NIT; ++i) { + if constexpr (!ALDG) + asm volatile("cp.async.wait_group %0;\n" ::"n"(SA - 2) : "memory"); + WG2_BAR(); + WG2_MMA(kaB + (i % SA) * A_E, kbB + (i % SB) * NSPLIT * B_E) + asm volatile("wgmma.wait_group.sync.aligned 1;\n" ::: "memory"); + if constexpr (NTS > 128) WG2_BAR(); + if (i >= 1) { + const int ca = (i + SA - 1) * KSPLIT + kgrp; + if (ca < NCH) WG2_PUTA(ca, kaB + ((i - 1) % SA) * A_E) + if (ALDG && ca + KSPLIT < NCH) WG2_LOADA(ca + KSPLIT) + const int cb = (i + SB - 1) * KSPLIT + kgrp; + if (cb < NCH) { + WG2_STSB(kbB + ((i - 1) % SB) * NSPLIT * B_E) + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if (cb + KSPLIT < NCH) WG2_LOADB(cb + KSPLIT) + } else if constexpr (ALDG) { + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + } + } + } + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); +#undef WG2_LOADB +#undef WG2_LOADA +#undef WG2_PUTA +#undef WG2_STSB +#undef WG2_MMA +#undef WG2_BAR + __syncthreads(); + float *Ck = Cs + kgrp * WG_BM * WG_CLD; + const int cb = (lane & 3) * 2; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Ck + (er0 + h * 8) * WG_CLD + + ngid * WG_BN + j * 8 + cb) = v; + } + } + __syncthreads(); + constexpr int NC = EPN / 4; + constexpr int MSTEP = WG_NT / NC; + constexpr int MITER = WG_BM / MSTEP; + constexpr int UNR = (MITER % 4 == 0) ? 4 : ((MITER % 2 == 0) ? 2 : 1); + static_assert(MITER % UNR == 0, "row batch must divide the row count"); + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + const int n = nb + lc; +#pragma unroll 1 + for (int t0 = 0; t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + v[u] = *reinterpret_cast(Cs + m * WG_CLD + lc); + if (resid != nullptr) { + const int co = m0 + m; + const int cc = co < Cout ? co : 0; + const __half2 *rp = reinterpret_cast( + resid + (size_t)cc * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + const int co = m0 + m; + if (co >= Cout) continue; + float4 vv = v[u]; +#pragma unroll + for (int g = 1; g < KSPLIT; ++g) { + float4 p = *reinterpret_cast( + Cs + g * WG_BM * WG_CLD + m * WG_CLD + lc); + vv.x += p.x; vv.y += p.y; vv.z += p.z; vv.w += p.w; + } + if (resid != nullptr) { + vv.x += __half2float(r0[u].x); vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (do_snake) { + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + } + __half *op = out + (size_t)co * SCout + n + H; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B * H; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int total = Cout * B * pad_len; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} diff --git a/deploy/nemotron_native/kernels/b32/native/codec/wg3_kernel.inc b/deploy/nemotron_native/kernels/b32/native/codec/wg3_kernel.inc new file mode 100644 index 0000000..fa55b32 --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/codec/wg3_kernel.inc @@ -0,0 +1,177 @@ +template +__device__ __noinline__ void wg3_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int WG_BNT = WG_BN * NSG; + constexpr int WG_CLD = WG_BNT + 8; + constexpr int WG_NT = 128 * NWG * NSG; + constexpr int WG_BM = 64 * NWG; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int SBO_AE = KG * WG_SLOT; + constexpr int LBO_BE = (WG_BN / 8) * WG_SLOT; + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + constexpr int ST_E = A_E + NSG * B_E; + __half *As0 = reinterpret_cast<__half *>(arena); + __half *Bs0 = As0 + A_E; + __half *As1 = (STAGES == 2) ? As0 + ST_E : As0; + __half *Bs1 = (STAGES == 2) ? Bs0 + ST_E : Bs0; + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int wgid = tid >> 7; + const int mgid = wgid % NWG, ngid = wgid / NWG; + const int warp = (tid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * WG_BM; + constexpr int WGROWS = 8; + const int nb = blockIdx.y * WG_BNT; + const int n0 = nb + ngid * WG_BN; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + if (tid < WG_BM) { + const int co = m0 + tid; + if (do_snake && co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + if (tid < WG_BNT) { + const int n = nb + tid; + colok[tid] = (n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mgid * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + constexpr int ATOT = WG_BM * KG; + constexpr int AREG = (ATOT + WG_NT - 1) / WG_NT; + constexpr int CIB = BK / KT; + constexpr int BTOT = NSG * CIB * (WG_BN / 8); + constexpr int BREG = (BTOT + WG_NT - 1) / WG_NT; + int4 ar[AREG]; + int4 br0[BREG], br1[BREG]; + uint32_t edge[BREG]; +#define WG_LOADG(kg0) { const int _ci0 = (kg0) / KT; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; ar[r] = make_int4(0, 0, 0, 0); if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; const int co = m0 + m, kg = (kg0) + v * 8; if (co < Cout && kg < Kg) ar[r] = *reinterpret_cast( w + (size_t)co * Kg + kg); } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; br0[r] = make_int4(0, 0, 0, 0); if constexpr (KT == 3) edge[r] = 0; else br1[r] = make_int4(0, 0, 0, 0); if (BTOT % WG_NT == 0 || idx < BTOT) { const int _g = idx / (CIB * (WG_BN / 8)); const int _r = idx - _g * (CIB * (WG_BN / 8)); const int ci = _ci0 + _r / (WG_BN / 8); const int nc = _r % (WG_BN / 8); if (ci < Cin) { const int4 *src = reinterpret_cast( in + (size_t)ci * SCin + nb + _g * WG_BN + nc * 8); br0[r] = src[0]; if constexpr (KT == 3) { if (nc == WG_BN / 8 - 1) edge[r] = src[1].x; } else { br1[r] = src[1]; } } } } } + WG_LOADG(0) +#define WG_STS(AD,BD) { _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; *reinterpret_cast( (AD) + (m >> 3) * SBO_AE + v * WG_SLOT + (m & 7) * 8) = ar[r]; } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; if (BTOT % WG_NT == 0 || idx < BTOT) { const int _g = idx / (CIB * (WG_BN / 8)); const int _r = idx - _g * (CIB * (WG_BN / 8)); const int cil = _r / (WG_BN / 8), nc = _r % (WG_BN / 8); unsigned int u[8]; u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; if constexpr (KT == 3) { u[4] = __shfl_down_sync(0xffffffffu, u[0], 1, WG_BN / 8); if (nc == WG_BN / 8 - 1) u[4] = edge[r]; u[5] = u[6] = u[7] = 0; } else { u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; } _Pragma("unroll") for (int k = 0; k < KT; ++k) { const int kh = k >> 1; int4 o; if (k & 1) { o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); } else { o.x = u[kh + 0]; o.y = u[kh + 1]; o.z = u[kh + 2]; o.w = u[kh + 3]; } const int row = cil * KT + k; *reinterpret_cast((BD) + _g * B_E + nc * WG_SLOT + (row >> 3) * LBO_BE + (row & 7) * 8) = o; } } } if constexpr (WG_BN == 64) { asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); } } +#define WG_MMA(AD,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t da0 = wg_desc((AD) + mgid * WGROWS * SBO_AE, WG_SLOT * 2, SBO_AE * 2); const uint64_t db0 = wg_desc((BD) + ngid * B_E, LBO_BE * 2, WG_SLOT * 2); _Pragma("unroll") for (int kk = 0; kk < BK; kk += 16) { const uint64_t da = da0 + (kk >> 3) * (WG_SLOT / 8); const uint64_t db = db0 + (kk >> 3) * (LBO_BE / 8); if constexpr (WG_BN == 64) { WGMMA_M64N64K16(acc, da, db); } else { WGMMA_M64N128K16(acc, da, db); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + WG_STS(As0, Bs0) + if (BK < Kg) WG_LOADG(BK) + __syncthreads(); + int cur = 0; + for (int kg0 = 0; kg0 < Kg; kg0 += BK) { + if (cur == 0) { WG_MMA(As0, Bs0) } else { WG_MMA(As1, Bs1) } + if (kg0 + BK < Kg) { + if (STAGES == 2) + asm volatile("wgmma.wait_group.sync.aligned 1;\n" ::: "memory"); + else + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + __syncthreads(); + if (cur == 0) { WG_STS(As1, Bs1) } else { WG_STS(As0, Bs0) } + if (kg0 + 2 * BK < Kg) WG_LOADG(kg0 + 2 * BK) + __syncthreads(); + } else { + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + } + if (STAGES == 2) cur ^= 1; + } +#undef WG_LOADG +#undef WG_STS +#undef WG_MMA + __syncthreads(); + const int r0 = er0; + const int cb = (lane & 3) * 2; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Cs + (r0 + h * 8) * WG_CLD + + ngid * WG_BN + j * 8 + cb) = v; + } + } + __syncthreads(); + constexpr int NC = WG_BNT / 4; + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + for (int m = lr; m < WG_BM; m += WG_NT / NC) { + const int co = m0 + m; + if (co >= Cout) continue; + const int n = nb + lc; + float4 v = *reinterpret_cast(Cs + m * WG_CLD + lc); + __half *op = out + (size_t)co * SCout + n + H; + if (resid != nullptr) { + const __half2 *rp = reinterpret_cast( + resid + (size_t)co * SCout + n + H); + __half2 r0h = rp[0], r1h = rp[1]; + v.x += __half2float(r0h.x); v.y += __half2float(r0h.y); + v.z += __half2float(r1h.x); v.w += __half2float(r1h.y); + } + if (do_snake) { + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + } + if (ok4) { + __half2 o0 = __floats2half2_rn(v.x, v.y); + __half2 o1 = __floats2half2_rn(v.z, v.w); + reinterpret_cast<__half2 *>(op)[0] = o0; + reinterpret_cast<__half2 *>(op)[1] = o1; + } else { + const float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B * H; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int total = Cout * B * pad_len; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} diff --git a/deploy/nemotron_native/kernels/b32/native/codec/wg4_step_widths.inc b/deploy/nemotron_native/kernels/b32/native/codec/wg4_step_widths.inc new file mode 100644 index 0000000..9fe1878 --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/codec/wg4_step_widths.inc @@ -0,0 +1,459 @@ +template +__device__ __noinline__ void wg4_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int CoutP, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + if constexpr (PDL && !LATE_PDL) PDL_ENTER(); + constexpr int NT = 128 * MSPLIT * NSPLIT; + constexpr int BM = 64 * MSPLIT; + constexpr int BNT = WG_BN * NSPLIT; + constexpr int NCRAW = BNT / 4; + constexpr int EPN = (NT % NCRAW == 0) ? BNT : (((BNT + 127) / 128) * 128); + constexpr int WG_CLD = EPN + 8; + constexpr int CG = CIB / 8; + constexpr int NBR = BNT + KT - 1; + constexpr int NBRP = BC8 ? ((NBR + 7) & ~7) : ((NBR + 31) & ~31); + constexpr int UCOL = NBRP / 32; + constexpr int NUNIT = CG * UCOL; + constexpr int NWARP = NT / 32; + constexpr int BITER = (NUNIT + NWARP - 1) / NWARP; + constexpr int B_E = CIB * NBRP; + constexpr int ST_E = B_E; + constexpr int LBO_B = NBRP * 8; + constexpr int ABUF = ABUF_; + static_assert(!OC8 || PREACT, + "K=7 C8 output is only used for residual-free pair heads"); + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *smem = reinterpret_cast<__half *>(arena); + __half *Hs = smem; + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int wgid = tid >> 7; + const int mg = wgid % MSPLIT, ng = wgid / MSPLIT; + const int warp = (tid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * BM; + const int nb = blockIdx.y * BNT; + const int H = KT - 1; + const int half = Cout >> 1; + const int NCH = (Cin + CIB - 1) / CIB; + for (int i = tid; i < BM; i += NT) { + const int co = m0 + i; + if (co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[i] = a; + row_ri[i] = __frcp_rn(a + 1e-9f); + } + } + for (int i = tid; i < EPN; i += NT) { + const int n = nb + i; + colok[i] = (i < BNT && n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mg * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + const int mtile = blockIdx.x * MSPLIT + mg; + const __half *const apt = w + ((size_t)mtile * NCH * KT) * 1024 + + (size_t)(tid & 127) * 8; + uint32_t areg[ABUF][KT][4]; + const int wi = tid >> 5; + int bdst[BITER], boff[BITER]; + bool blive[BITER]; + const __half *const bbase = in + (size_t)(lane >> 2) * SCin + nb + + 8 * (lane & 3); +#pragma unroll + for (int t = 0; t < BITER; ++t) { + const int u = wi + t * NWARP; + const int gc = (u < NUNIT) ? u / UCOL : 0; + const int cu = (u < NUNIT) ? u - gc * UCOL : 0; + blive[t] = u < NUNIT; + bdst[t] = gc * LBO_B + (cu * 32 + lane) * 8; + boff[t] = gc * 8 * SCin + cu * 32; + } + const size_t SC8i = (size_t)SCin * 8; + constexpr int BQ = B_E / 8; + constexpr int BCREG = (BQ + NT - 1) / NT; + int cdst[BCREG]; + const __half *csrc[BCREG]; + bool clive[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) { + const int q = tid + r * NT; + const int gc = q / NBRP, rem = q - gc * NBRP; + clive[r] = ((BQ % NT == 0) || q < BQ) && gc < CG; + cdst[r] = clive[r] ? gc * LBO_B + rem * 8 : 0; + csrc[r] = in + (clive[r] ? (size_t)gc * SC8i + + (size_t)nb * 8 + rem * 8 : 0); + } + uint32_t br[BITER][4]; + constexpr int CPA_STAGES = 4; + constexpr int BSTAGES = BC8 ? CPA_STAGES : STAGES; + static_assert(STAGES >= CPA_STAGES, "C8 B-panel pipeline allocation"); + static_assert(ABUF == 2 && BSTAGES == 4, "wg4 unrolled pipeline shape"); + uint64_t dslot[BSTAGES]; +#pragma unroll + for (int s = 0; s < BSTAGES; ++s) + dslot[s] = wg_desc(smem + s * ST_E + ng * WG_BN * 8, LBO_B * 2u, 128u); + uint32_t csm0[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) + csm0[r] = static_cast(__cvta_generic_to_shared(smem + cdst[r])); + uint32_t bsm0[BITER]; +#pragma unroll + for (int t = 0; t < BITER; ++t) + bsm0[t] = static_cast(__cvta_generic_to_shared(smem + bdst[t])); + const size_t CSTEP = (size_t)CG * SC8i; + const __half *crol[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) crol[r] = csrc[r] + CPA_STAGES * CSTEP; + const __half *arol = apt + (size_t)ABUF * (KT * 1024); +#define WG4_CPB(c,BD) { const size_t _o = (size_t)(c) * CG * SC8i; _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16((BD) + cdst[r], csrc[r] + _o, true); } +#define WG4_CPB_R(SLOT) { _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16o<(SLOT) * ST_E * 2>(csm0[r], crol[r], true); _Pragma("unroll") for (int r = 0; r < BCREG; ++r) crol[r] += CSTEP; } +#define WG4_LOADB(c) { const size_t _o = (size_t)(c) * CIB * SCin; _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { const int4 _v = *reinterpret_cast( bbase + _o + boff[t]); br[t][0] = _v.x; br[t][1] = _v.y; br[t][2] = _v.z; br[t][3] = _v.w; } } } +#define WG4_XPOSE(u) { const int _q = lane & 3; uint32_t _x, _y; _x = (_q & 1) ? u[0] : u[1]; _y = (_q & 1) ? u[2] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 1); _y = __shfl_xor_sync(0xffffffffu, _y, 1); if (_q & 1) { u[0] = _x; u[2] = _y; } else { u[1] = _x; u[3] = _y; } _x = (_q & 2) ? u[0] : u[2]; _y = (_q & 2) ? u[1] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 2); _y = __shfl_xor_sync(0xffffffffu, _y, 2); if (_q & 2) { u[0] = _x; u[1] = _y; } else { u[2] = _x; u[3] = _y; } } +#define WG4_STOREB(BD) { _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { uint32_t _u[4]; _u[0] = br[t][0]; _u[1] = br[t][1]; _u[2] = br[t][2]; _u[3] = br[t][3]; WG4_XPOSE(_u) const uint32_t _a = static_cast( __cvta_generic_to_shared((BD) + bdst[t])); asm volatile( "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 " "[%0], {%1,%2,%3,%4};\n" ::"r"(_a), "r"(_u[0]), "r"(_u[1]), "r"(_u[2]), "r"(_u[3])); } } } +#define WG4_STOREB_R(SLOT) { _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { uint32_t _u[4]; _u[0] = br[t][0]; _u[1] = br[t][1]; _u[2] = br[t][2]; _u[3] = br[t][3]; WG4_XPOSE(_u) asm volatile( "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 " "[%0+%5], {%1,%2,%3,%4};\n" ::"r"(bsm0[t]), "r"(_u[0]), "r"(_u[1]), "r"(_u[2]), "r"(_u[3]), "n"((SLOT) * ST_E * 2)); } } } +#define WG4_LOADA_R(BUF) { _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { const int4 _v = *reinterpret_cast(arol + kk * 1024); areg[BUF][kk][0] = _v.x; areg[BUF][kk][1] = _v.y; areg[BUF][kk][2] = _v.z; areg[BUF][kk][3] = _v.w; } arol += KT * 1024; } +#define WG4_MMA_R(BUF,SLOT) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { WGMMA_RS_ANY_T(WG_BN, acc, areg[BUF][kk], dslot[SLOT] + kk, 0); } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } +#define WG4_LOADA(BUF,c) { const __half *_p = apt + (size_t)(c) * (KT * 1024); _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { const int4 _v = *reinterpret_cast(_p + kk * 1024); areg[BUF][kk][0] = _v.x; areg[BUF][kk][1] = _v.y; areg[BUF][kk][2] = _v.z; areg[BUF][kk][3] = _v.w; } } +#define WG4_KEEPA(BUF) { _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) asm volatile("" : "+r"(areg[BUF][kk][0]), "+r"(areg[BUF][kk][1]), "+r"(areg[BUF][kk][2]), "+r"(areg[BUF][kk][3])); } +#define WG4_MMA(BUF,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t db0 = wg_desc((BD) + ng * WG_BN * 8, LBO_B * 2u, 128u); _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { WGMMA_RS_ANY_T(WG_BN, acc, areg[BUF][kk], db0 + kk, 0); } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + if constexpr (PDL && LATE_PDL) PDL_ENTER(); + constexpr int WGD = 1; + if constexpr (BC8) { +#pragma unroll 1 + for (int st = 0; st < CPA_STAGES; ++st) { + if (st < NCH) WG4_CPB(st, smem + st * ST_E) + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } + WG4_LOADA(0, 0) + if (ABUF > 1 && 1 < NCH) WG4_LOADA(1, 1) + if (ABUF > 2 && 2 < NCH) WG4_LOADA(2, 2) + } else { + WG4_LOADB(0) + WG4_LOADA(0, 0) + if (ABUF > 1 && 1 < NCH) WG4_LOADA(1, 1) + if (ABUF > 2 && 2 < NCH) WG4_LOADA(2, 2) + WG4_STOREB(smem) + if (1 < NCH) WG4_LOADB(1) + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + } +#define WG4_BODY(BUF,i,SLOT) { if constexpr (BC8) { asm volatile("cp.async.wait_group %0;\n" : : "n"(CPA_STAGES - 3) : "memory"); __syncthreads(); if ((i) >= 2) { if ((i) + 2 < NCH) { WG4_CPB_R(((SLOT) + 2) % CPA_STAGES) } asm volatile("cp.async.commit_group;\n" ::: "memory"); } } else { if (i) __syncthreads(); } WG4_MMA_R(BUF, SLOT) asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(WGD) : "memory"); if constexpr (NSPLIT > 1) __syncthreads(); if constexpr (!BC8) { const int cb = (i) + 1; if (cb < NCH) { WG4_STOREB_R(((SLOT) + 1) % BSTAGES) asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); if (cb + 1 < NCH) WG4_LOADB(cb + 1) } } if ((i) >= 1) { const int nb2 = ((BUF) + ABUF - 1) % ABUF; if ((i) + ABUF - 1 < NCH) { WG4_KEEPA(nb2) WG4_LOADA_R(nb2) } } } + for (int i = 0; i < NCH; i += 4) { + WG4_BODY(0, i, 0) + if (i + 1 < NCH) WG4_BODY(1, i + 1, 1) + if (i + 2 < NCH) WG4_BODY(0, i + 2, 2) + if (i + 3 < NCH) WG4_BODY(1, i + 3, 3) + } + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + if constexpr (BC8) + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); +#undef WG4_CPB +#undef WG4_CPB_R +#undef WG4_STOREB_R +#undef WG4_LOADA_R +#undef WG4_MMA_R +#undef WG4_LOADB +#undef WG4_XPOSE +#undef WG4_STOREB +#undef WG4_LOADA +#undef WG4_KEEPA +#undef WG4_MMA +#undef WG4_BODY +#define TL4_SW(c,r) ((c) * BM + (((r) & ~7) ^ (((c) & 7) << 3)) + ((r) & 7)) + __syncthreads(); + const int cb2 = (lane & 3) * 2; + const bool snake = PREACT && (m0 + er0) < half; + const float sn_a0 = snake ? row_alpha[er0] : 0.f; + const float sn_r0 = snake ? row_ri[er0] : 0.f; + const float sn_a1 = snake ? row_alpha[er0 + 8] : 0.f; + const float sn_r1 = snake ? row_ri[er0 + 8] : 0.f; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + const int row = er0 + h * 8; + if constexpr (PREACT) { + if (snake) { + const float a = h ? sn_a1 : sn_a0; + const float ri = h ? sn_r1 : sn_r0; + const float s0 = __sinf(a * v.x); + const float s1 = __sinf(a * v.y); + v.x = __fmaf_rn(s0 * s0, ri, v.x); + v.y = __fmaf_rn(s1 * s1, ri, v.y); + } else { + v.x = v.x > 0.f ? v.x : 0.01f * v.x; + v.y = v.y > 0.f ? v.y : 0.01f * v.y; + } + if constexpr (OC8) { + const int c0 = ng * WG_BN + j * 8 + cb2; + Hs[TL4_SW(c0, row)] = __float2half(v.x); + Hs[TL4_SW(c0 + 1, row)] = __float2half(v.y); + } else { + *reinterpret_cast<__half2 *>( + Hs + row * WG_CLD + ng * WG_BN + j * 8 + cb2) = + __floats2half2_rn(v.x, v.y); + } + } else { + *reinterpret_cast( + Cs + row * WG_CLD + ng * WG_BN + j * 8 + cb2) = v; + } + } + } + __syncthreads(); + if constexpr (PREACT) { + if constexpr (OC8) { + constexpr int GRP = BM / 8; + constexpr int TOTG = GRP * BNT; + constexpr int GITER = (TOTG + NT - 1) / NT; + const size_t SC8o = (size_t)SCout * 8; +#pragma unroll + for (int t = 0; t < GITER; ++t) { + const int idx = tid + t * NT; + if (TOTG % NT != 0 && idx >= TOTG) continue; + const int g = idx / BNT; + const int j = idx - g * BNT; + if (!colok[j]) continue; + const int m0g = g * 8; + if (m0 + m0g >= Cout) continue; + *reinterpret_cast( + out + (size_t)((m0 + m0g) >> 3) * SC8o + + (size_t)(nb + j + H) * 8) = + *reinterpret_cast(Hs + TL4_SW(j, m0g)); + } + } else { + static_assert((BNT & 3) == 0, "four-column preactivated epilogue"); + constexpr int NC = BNT / 4; + constexpr int TOT = BM * NC; + for (int q = tid; q < TOT; q += NT) { + const int m = q / NC; + const int lc = (q - m * NC) * 4; + const int co = m0 + m; + const bool any4 = (colok[lc] | colok[lc + 1] | + colok[lc + 2] | colok[lc + 3]) != 0; + if (co >= Cout || !any4) continue; + const __half2 *sp = reinterpret_cast( + Hs + m * WG_CLD + lc); + const __half2 v0 = sp[0], v1 = sp[1]; + __half *op = out + (size_t)co * SCout + nb + lc + H; + const bool ok4 = (colok[lc] & colok[lc + 1] & + colok[lc + 2] & colok[lc + 3]) != 0; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = v0; + reinterpret_cast<__half2 *>(op)[1] = v1; + } else { + if (colok[lc + 0]) op[0] = v0.x; + if (colok[lc + 1]) op[1] = v0.y; + if (colok[lc + 2]) op[2] = v1.x; + if (colok[lc + 3]) op[3] = v1.y; + } + } + } + } else { + constexpr bool COMPACT_EPI = + (WG_BN == 136 || WG_BN == 160 || WG_BN == 208); + if constexpr (COMPACT_EPI) { + static_assert((BNT & 3) == 0, "4-column epilogue groups"); + constexpr int NC = BNT / 4; + constexpr int TOT = BM * NC; + constexpr int MITER = (TOT + NT - 1) / NT; + constexpr int UNR = (MITER >= 8) ? 8 : ((MITER >= 4) ? 4 : + ((MITER >= 2) ? 2 : 1)); +#pragma unroll 1 + for (int t0 = 0; t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int q = tid + (t0 + u) * NT; + if (q >= TOT) continue; + const int m = q / NC; + const int lc = (q - m * NC) * 4; + const int co = m0 + m; + const bool any4 = (colok[lc] | colok[lc + 1] | + colok[lc + 2] | colok[lc + 3]) != 0; + if (co >= Cout || !any4) continue; + v[u] = *reinterpret_cast(Cs + m * WG_CLD + lc); + if constexpr (!PREACT) { + const int n = nb + lc; + const __half2 *rp = reinterpret_cast( + resid + (size_t)co * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int q = tid + (t0 + u) * NT; + if (q >= TOT) continue; + const int m = q / NC; + const int lc = (q - m * NC) * 4; + const int co = m0 + m; + const bool any4 = (colok[lc] | colok[lc + 1] | + colok[lc + 2] | colok[lc + 3]) != 0; + if (co >= Cout || !any4) continue; + const int n = nb + lc; + float4 vv = v[u]; + if constexpr (!PREACT) { + vv.x += __half2float(r0[u].x); vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (co < half) { + const float a = row_alpha[m], ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + __half *op = out + (size_t)co * SCout + n + H; + const bool ok4 = (colok[lc] & colok[lc + 1] & + colok[lc + 2] & colok[lc + 3]) != 0; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + } else { + constexpr int NC = EPN / 4; + constexpr int MSTEP = NT / NC; + constexpr int MITER = BM / MSTEP; + constexpr int UNR = (MITER % 8 == 0) ? 8 + : ((MITER % 4 == 0) ? 4 + : ((MITER % 2 == 0) ? 2 : 1)); + static_assert(MITER % UNR == 0, "row batch must divide row count"); + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && + colok[lc + 2] && colok[lc + 3]; + const bool live4 = colok[lc] || colok[lc + 1] || + colok[lc + 2] || colok[lc + 3]; + const int n = nb + lc; +#pragma unroll 1 + for (int t0 = 0; live4 && t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + v[u] = *reinterpret_cast( + Cs + m * WG_CLD + lc); + if constexpr (!PREACT) { + const int co = m0 + m; + const int cc = co < Cout ? co : 0; + const __half2 *rp = reinterpret_cast( + resid + (size_t)cc * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + const int co = m0 + m; + if (co >= Cout) continue; + float4 vv = v[u]; + if constexpr (!PREACT) { + vv.x += __half2float(r0[u].x); + vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); + vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (co < half) { + const float a = row_alpha[m], ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + __half *op = out + (size_t)co * SCout + n + H; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B; + for (int idx = bxy * NT + tid; idx < total; idx += nbx * NT) { + const int b = idx / Cin, ci = idx - b * Cin; + const int x = b * L + native_true_tokens(b,T); + __half *const dst = ns + ((size_t)b * Cin + ci) * H; + if constexpr (BC8) { + const __half *const src = + in + (size_t)(ci >> 3) * ((size_t)SCin * 8) + + (size_t)x * 8 + (ci & 7); +#pragma unroll + for (int j = 0; j < H; ++j) dst[j] = src[j * 8]; + } else { + const __half *const src = in + (size_t)ci * SCin + x; +#pragma unroll + for (int j = 0; j < H; j += 2) + *reinterpret_cast(dst + j) = + *reinterpret_cast(src + j); + } + } + } + if (pad_src != nullptr) { + const int total = Cout * B; + for (int idx = bxy * NT + tid; idx < total; idx += nbx * NT) { + const int b = idx / Cout, co = idx - b * Cout; + const int x = b * L + pad_off; + const __half *const src = + pad_src + ((size_t)b * Cout + co) * pad_len; + if constexpr (OC8) { + __half *const dst = + out + (size_t)(co >> 3) * ((size_t)SCout * 8) + + (size_t)x * 8 + (co & 7); + for (int j = 0; j < pad_len; ++j) dst[j * 8] = src[j]; + } else { + __half *const dst = out + (size_t)co * SCout + x; + for (int j = 0; j < pad_len; ++j) dst[j] = src[j]; + } + } + } + } +} diff --git a/deploy/nemotron_native/kernels/b32/native/codec/wg5_step_widths.inc b/deploy/nemotron_native/kernels/b32/native/codec/wg5_step_widths.inc new file mode 100644 index 0000000..5a1a0ed --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/codec/wg5_step_widths.inc @@ -0,0 +1,469 @@ +struct alignas(8) Wg5BulkBarrier { unsigned long long opaque; }; +__device__ __forceinline__ uint32_t wg5_bulk_smem(const void *p) { + return static_cast(__cvta_generic_to_shared(p)); +} +__device__ __forceinline__ void wg5_bulk_init(Wg5BulkBarrier &bar) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(wg5_bulk_smem(&bar)), "r"(1u) : "memory"); +} +__device__ __forceinline__ void wg5_bulk_expect( + Wg5BulkBarrier &bar, uint32_t bytes) { + asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 " + "_, [%0], %1;" + :: "r"(wg5_bulk_smem(&bar)), "r"(bytes) : "memory"); +} +__device__ __forceinline__ void wg5_bulk_load( + void *dst, const void *src, uint32_t bytes, Wg5BulkBarrier &bar) { + asm volatile( + "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes " + "[%0], [%1], %2, [%3];" + :: "r"(wg5_bulk_smem(dst)), "l"(src), "r"(bytes), + "r"(wg5_bulk_smem(&bar)) : "memory"); +} +__device__ __forceinline__ void wg5_bulk_wait( + Wg5BulkBarrier &bar, int phase) { + asm volatile( + "{ .reg .pred done; WG5_BULK_WAIT: " + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 " + "done, [%0], %1; @!done bra.uni WG5_BULK_WAIT; }" + :: "r"(wg5_bulk_smem(&bar)), "r"(phase) : "memory"); +} +__device__ __forceinline__ void wg5_bulk_inval(Wg5BulkBarrier &bar) { + asm volatile("mbarrier.inval.shared::cta.b64 [%0];" + :: "r"(wg5_bulk_smem(&bar)) : "memory"); +} +template +__device__ __noinline__ +void wg5_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int KT = 3; + constexpr int BM = 64 * MSPLIT; + constexpr int BNT = WG_BN; + constexpr int NTS = 128 * MSPLIT; + constexpr int WG_NT = NTS * KSPLIT; + constexpr int NCRAW = BNT / 4; + constexpr int EPN = + (WG_NT % NCRAW == 0) ? BNT : (((BNT + 127) / 128) * 128); + constexpr int WG_CLD_R = (OC8 ? WG_BN : EPN) + 8; + constexpr int WG_CLD = (WG_CLD_R % 16 == 0) ? WG_CLD_R + 4 : WG_CLD_R; + constexpr int CG = CIB / 8; + constexpr int CSUB = CIB / 16; + constexpr int AK = KT * CSUB; + constexpr int NBRP = BC8 ? (((BNT + KT - 1) + 7) & ~7) + : (((BNT + KT - 1) + 31) & ~31); + constexpr int UCOL = NBRP / 32; + constexpr int NUNIT = CG * UCOL; + constexpr int NWARP = NTS / 32; + constexpr int BITER = (NUNIT + NWARP - 1) / NWARP; + constexpr int B_E = CIB * NBRP; + constexpr int ST_E = SB * B_E; + constexpr int LBO_B = NBRP * 8; + constexpr int ABUF = 3; + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *smem = reinterpret_cast<__half *>(arena); + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int kgrp = tid / NTS; + const int stid = tid - kgrp * NTS; + const int mg = stid >> 7; + const int warp = (stid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * BM; + const int nb = blockIdx.y * BNT; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + constexpr int NCH = (CIN + CIB - 1) / CIB; + static_assert(NCH % KSPLIT == 0, "wg5 K-split must divide the chunk count"); + if (tid < BM) { + const int co = m0 + tid; + if (co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + for (int i = tid; i < EPN; i += WG_NT) { + const int n = nb + i; + colok[i] = (i < BNT && n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mg * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (kgrp == 0 && eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (kgrp == 0 && eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + const int mtile = blockIdx.x * MSPLIT + mg; + const __half *const apt = w + ((size_t)mtile * NCH * AK) * 1024 + + (size_t)(stid & 127) * 8; + uint32_t areg[ABUF][AK][4]; + const int wi = stid >> 5; + int bdst[BITER], boff[BITER]; + bool blive[BITER]; + const __half *const bbase = in + (size_t)(lane >> 2) * SCin + nb + + 8 * (lane & 3); +#pragma unroll + for (int t = 0; t < BITER; ++t) { + const int u = wi + t * NWARP; + const int gc = (u < NUNIT) ? u / UCOL : 0; + const int cu = (u < NUNIT) ? u - gc * UCOL : 0; + blive[t] = u < NUNIT; + bdst[t] = gc * LBO_B + (cu * 32 + lane) * 8; + boff[t] = gc * 8 * SCin + cu * 32; + } + const size_t SC8i = (size_t)SCin * 8; + constexpr int BQ = B_E / 8; + constexpr int BCREG = (BQ + NTS - 1) / NTS; + int cdst[BCREG]; + const __half *csrc[BCREG]; + bool clive[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) { + const int q = stid + r * NTS; + const int gc = q / NBRP, rem = q - gc * NBRP; + clive[r] = ((BQ % NTS == 0) || q < BQ) && gc < CG; + cdst[r] = clive[r] ? gc * LBO_B + rem * 8 : 0; + csrc[r] = in + (clive[r] ? (size_t)gc * SC8i + + (size_t)nb * 8 + rem * 8 : 0); + } + uint32_t br[BITER][4]; + __half *const kbB = smem + kgrp * ST_E; + const unsigned barid = 1u + kgrp + (unsigned)(KSPLIT * native_subid()); + Wg5BulkBarrier *const bulk_ready = reinterpret_cast( + native_aux() + 3072) + kgrp * ABUF; + const __half *bulk_rol = in + + (size_t)(ABUF * KSPLIT + kgrp) * CG * SC8i + (size_t)nb * 8; + uint64_t dbuf[ABUF]; +#pragma unroll + for (int s = 0; s < ABUF; ++s) + dbuf[s] = wg_desc(kbB + s * B_E, LBO_B * 2u, 128u); + uint32_t csm0[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) + csm0[r] = static_cast( + __cvta_generic_to_shared(kbB + cdst[r])); + const size_t CSTEP = (size_t)KSPLIT * CG * SC8i; + constexpr int ASTEP = KSPLIT * AK * 1024; + const __half *crol[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) + crol[r] = csrc[r] + (size_t)(ABUF * KSPLIT + kgrp) * CG * SC8i; + const __half *arol = apt + (size_t)(ABUF * KSPLIT + kgrp) * (AK * 1024); + int crolc = ABUF * KSPLIT + kgrp; +#define WG5_BAR() asm volatile("barrier.sync %0, %1;" ::"r"(barid), "r"(NTS) : "memory") +#define WG5_BULK(SLOT,SRC) \ + { \ + if (stid == 0) { \ + Wg5BulkBarrier &_bar = bulk_ready[(SLOT)]; \ + wg5_bulk_expect(_bar, B_E * (int)sizeof(__half)); \ + _Pragma("unroll") for (int _gc = 0; _gc < CG; ++_gc) \ + wg5_bulk_load(kbB + (SLOT) * B_E + _gc * LBO_B, \ + (SRC) + (size_t)_gc * SC8i, \ + LBO_B * (int)sizeof(__half), _bar); \ + } \ + } + if constexpr (BC8) { + if (stid == 0) { +#pragma unroll + for (int s = 0; s < ABUF; ++s) + wg5_bulk_init(bulk_ready[s]); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + WG5_BAR(); + } +#define WG5_CPB(c,BD) { const size_t _o = (size_t)(c) * CG * SC8i; _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16((BD) + cdst[r], csrc[r] + _o, true); } +#define WG5_LOADA(BUF,c) { const __half *_p = apt + (size_t)(c) * (AK * 1024); _Pragma("unroll") for (int kk = 0; kk < AK; ++kk) { const int4 _v = *reinterpret_cast(_p + kk * 1024); areg[BUF][kk][0] = _v.x; areg[BUF][kk][1] = _v.y; areg[BUF][kk][2] = _v.z; areg[BUF][kk][3] = _v.w; } } +#define WG5_KEEPA(BUF) { _Pragma("unroll") for (int kk = 0; kk < AK; ++kk) asm volatile("" : "+r"(areg[BUF][kk][0]), "+r"(areg[BUF][kk][1]), "+r"(areg[BUF][kk][2]), "+r"(areg[BUF][kk][3])); } +#define WG5_LOADB(c) { const size_t _o = (size_t)(c) * CIB * SCin; _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { const int4 _v = *reinterpret_cast( bbase + _o + boff[t]); br[t][0] = _v.x; br[t][1] = _v.y; br[t][2] = _v.z; br[t][3] = _v.w; } } } +#define WG5_BXPOSE(u) { const int _q = lane & 3; uint32_t _x, _y; _x = (_q & 1) ? u[0] : u[1]; _y = (_q & 1) ? u[2] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 1); _y = __shfl_xor_sync(0xffffffffu, _y, 1); if (_q & 1) { u[0] = _x; u[2] = _y; } else { u[1] = _x; u[3] = _y; } _x = (_q & 2) ? u[0] : u[2]; _y = (_q & 2) ? u[1] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 2); _y = __shfl_xor_sync(0xffffffffu, _y, 2); if (_q & 2) { u[0] = _x; u[1] = _y; } else { u[2] = _x; u[3] = _y; } } +#define WG5_STOREB(BD) { _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { uint32_t _u[4]; _u[0] = br[t][0]; _u[1] = br[t][1]; _u[2] = br[t][2]; _u[3] = br[t][3]; WG5_BXPOSE(_u) const uint32_t _a = static_cast( __cvta_generic_to_shared((BD) + bdst[t])); asm volatile( "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 " "[%0], {%1,%2,%3,%4};\n" ::"r"(_a), "r"(_u[0]), "r"(_u[1]), "r"(_u[2]), "r"(_u[3])); } } } +#define WG5_MMA(BUF) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); _Pragma("unroll") for (int h = 0; h < CSUB; ++h) { _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { WGMMA_RS_ANY_T(WG_BN, acc, areg[BUF][h * KT + kk], dbuf[BUF] + kk + h * (LBO_B / 4), 0); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } +#define WG5_CPB_R(SLOT) { _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16o<(SLOT) * B_E * 2>(csm0[r], crol[r], true); } +#define WG5_LOADA_R(BUF) { _Pragma("unroll") for (int kk = 0; kk < AK; ++kk) { const int4 _v = *reinterpret_cast(arol + kk * 1024); areg[BUF][kk][0] = _v.x; areg[BUF][kk][1] = _v.y; areg[BUF][kk][2] = _v.z; areg[BUF][kk][3] = _v.w; } } +#define WG5_ROLL() { crolc += KSPLIT; arol += ASTEP; if constexpr (BC8) bulk_rol += CSTEP; } + if constexpr (PDL) { + (void)0; + if (tid == 0) + (void)0; + } + static_assert(SB == 3 && WGD == 1 && ABUF == 3, "wg5 pipeline shape"); + if constexpr (BC8) { +#pragma unroll 1 + for (int s = 0; s < SB; ++s) { + const int c = s * KSPLIT + kgrp; + if (c < NCH) { + const __half *src = in + (size_t)c * CG * SC8i + + (size_t)nb * 8; + WG5_BULK(s, src) + } + } + } else { + WG5_LOADB(kgrp) +#pragma unroll 1 + for (int s = 0; s < SB; ++s) { + const int c = s * KSPLIT + kgrp; + if (c < NCH) WG5_STOREB(kbB + s * B_E) + if (c + KSPLIT < NCH) WG5_LOADB(c + KSPLIT) + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + } + WG5_LOADA(0, kgrp) + if (KSPLIT + kgrp < NCH) WG5_LOADA(1, KSPLIT + kgrp) + if (2 * KSPLIT + kgrp < NCH) WG5_LOADA(2, 2 * KSPLIT + kgrp) + constexpr int NIT = NCH / KSPLIT; + static_assert(NIT % ABUF == 0, "wg5 mainloop unroll shape"); +#define WG5_BODY(BUF,i) { if constexpr (BC8) wg5_bulk_wait(bulk_ready[(BUF)], ((i) / ABUF) & 1); WG5_BAR(); WG5_MMA(BUF) asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(WGD) : "memory"); if constexpr (NTS > 128) WG5_BAR(); if ((i) >= 1) { const int nb2 = ((BUF) + 2) % ABUF; if (crolc < NCH) { if constexpr (BC8) { WG5_BULK(nb2, bulk_rol) } else { WG5_STOREB(kbB + nb2 * B_E) asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); if (crolc + KSPLIT < NCH) WG5_LOADB(crolc + KSPLIT) } WG5_KEEPA(nb2) WG5_LOADA_R(nb2) } WG5_ROLL() } } +#pragma unroll + for (int i = 0; i < NIT; i += 3) { + WG5_BODY(0, i) + WG5_BODY(1, i + 1) + WG5_BODY(2, i + 2) + } + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + if constexpr (BC8) { + WG5_BAR(); + if (stid == 0) { +#pragma unroll + for (int s = 0; s < ABUF; ++s) + wg5_bulk_inval(bulk_ready[s]); + } + WG5_BAR(); + } +#undef WG5_BXPOSE +#undef WG5_BULK +#undef WG5_CPB +#undef WG5_CPB_R +#undef WG5_LOADA_R +#undef WG5_ROLL +#undef WG5_LOADB +#undef WG5_STOREB +#undef WG5_MMA +#undef WG5_LOADA +#undef WG5_KEEPA +#undef WG5_BODY +#undef WG5_BAR + __syncthreads(); + const int cb2 = (lane & 3) * 2; + if constexpr (OC8) { + __half *Ck = smem + kgrp * BM * WG_CLD; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + *reinterpret_cast<__half2 *>(Ck + (er0 + h * 8) * WG_CLD + + j * 8 + cb2) = + __floats2half2_rn(acc[j * 4 + h * 2], + acc[j * 4 + h * 2 + 1]); + } + } + } else { + float *Ck = Cs + kgrp * BM * WG_CLD; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], + acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Ck + (er0 + h * 8) * WG_CLD + + j * 8 + cb2) = v; + } + } + } + __syncthreads(); + if constexpr (OC8) { + constexpr int GRP = BM / 8; + constexpr int TOTG = GRP * WG_BN; + constexpr int GITER = (TOTG + WG_NT - 1) / WG_NT; + const size_t SC8o = (size_t)SCout * 8; +#pragma unroll + for (int t = 0; t < GITER; ++t) { + const int idx = tid + t * WG_NT; + if (TOTG % WG_NT != 0 && idx >= TOTG) continue; + const int j = idx - (idx / WG_BN) * WG_BN, g = idx / WG_BN; + if (!colok[j]) continue; + const int m0g = g * 8; + if (m0 + m0g >= Cout) continue; + const int n = nb + j; + __half rr[8]; + if constexpr (CIN == 1536) { + if constexpr (RC8) { + *reinterpret_cast(rr) = + *reinterpret_cast( + resid + (size_t)((m0 + m0g) >> 3) * SC8o + + (size_t)(n + H) * 8); + } else { +#pragma unroll + for (int h = 0; h < 8; ++h) + rr[h] = resid[(size_t)(m0 + m0g + h) * SCout + n + H]; + } + } + uint32_t pk[4]; + const bool snake = (m0 + m0g) < half; +#pragma unroll + for (int h = 0; h < 8; h += 2) { + float f[2]; +#pragma unroll + for (int e = 0; e < 2; ++e) { + const int m = m0g + h + e; + float a = __half2float(smem[m * WG_CLD + j]); +#pragma unroll + for (int kg = 1; kg < KSPLIT; ++kg) + a += __half2float( + smem[kg * BM * WG_CLD + m * WG_CLD + j]); + if constexpr (CIN == 1536) + a += __half2float(rr[h + e]); + if (snake) { + const float sn = __sinf(row_alpha[m] * a); + a = __fmaf_rn(sn * sn, row_ri[m], a); + } else { + a = a > 0.f ? a : 0.01f * a; + } + f[e] = a; + } + const __half2 p = __floats2half2_rn(f[0], f[1]); + pk[h >> 1] = *reinterpret_cast(&p); + } + *reinterpret_cast( + out + (size_t)((m0 + m0g) >> 3) * SC8o + + (size_t)(n + H) * 8) = make_int4(pk[0], pk[1], pk[2], pk[3]); + } + } else { + constexpr int NC = EPN / 4; + constexpr int MSTEP = WG_NT / NC; + constexpr int MITER = BM / MSTEP; + constexpr int UNR = (MITER % 4 == 0) ? 4 : ((MITER % 2 == 0) ? 2 : 1); + static_assert(MITER % UNR == 0, "row batch must divide the row count"); + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + const bool live4 = colok[lc] || colok[lc + 1] || colok[lc + 2] || + colok[lc + 3]; + const int n = nb + lc; +#pragma unroll 1 + for (int t0 = 0; live4 && t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + v[u] = *reinterpret_cast(Cs + m * WG_CLD + lc); + if constexpr (CIN == 1536) { + const int co = m0 + m; + const int cc = co < Cout ? co : 0; + if constexpr (RC8) { + const __half *rp = + resid + (size_t)(cc >> 3) * ((size_t)SCout * 8) + + (size_t)(n + H) * 8 + (cc & 7); + r0[u].x = rp[0]; r0[u].y = rp[8]; + r1[u].x = rp[16]; r1[u].y = rp[24]; + } else { + const __half2 *rp = reinterpret_cast( + resid + (size_t)cc * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + const int co = m0 + m; + if (co >= Cout) continue; + float4 vv = v[u]; +#pragma unroll + for (int g = 1; g < KSPLIT; ++g) { + float4 p = *reinterpret_cast( + Cs + g * BM * WG_CLD + m * WG_CLD + lc); + vv.x += p.x; vv.y += p.y; vv.z += p.z; vv.w += p.w; + } + if constexpr (CIN == 1536) { + vv.x += __half2float(r0[u].x); vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + __half *op = out + (size_t)co * SCout + n + H; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + static_assert(CIN == 768 || CIN == 1536); + constexpr int COUT = (CIN == 768) ? 1536 : 768; + if (ns != nullptr) { + const int total = CIN * B; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int b = idx / CIN, ci = idx - b * CIN; + const int x = b * L + native_true_tokens(b,T); + __half *const dst = ns + ((size_t)b * CIN + ci) * H; + if constexpr (BC8) { + const __half *const src = + in + (size_t)(ci >> 3) * ((size_t)SCin * 8) + + (size_t)x * 8 + (ci & 7); +#pragma unroll + for (int j = 0; j < H; ++j) dst[j] = src[j * 8]; + } else { + const __half *const src = in + (size_t)ci * SCin + x; +#pragma unroll + for (int j = 0; j < H; j += 2) + *reinterpret_cast(dst + j) = + *reinterpret_cast(src + j); + } + } + } + if (pad_src != nullptr) { + const int total = COUT * B; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int b = idx / COUT, co = idx - b * COUT; + const int x = b * L + pad_off; + const __half *const src = + pad_src + ((size_t)b * COUT + co) * pad_len; + if constexpr (OC8) { + __half *const dst = + out + (size_t)(co >> 3) * ((size_t)SCout * 8) + + (size_t)x * 8 + (co & 7); + for (int j = 0; j < pad_len; ++j) dst[j * 8] = src[j]; + } else { + __half *const dst = out + (size_t)co * SCout + x; + for (int j = 0; j < pad_len; ++j) dst[j] = src[j]; + } + } + } + } +} diff --git a/deploy/nemotron_native/kernels/b32/native/codec/wg_kernel.inc b/deploy/nemotron_native/kernels/b32/native/codec/wg_kernel.inc new file mode 100644 index 0000000..2d5c766 --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/codec/wg_kernel.inc @@ -0,0 +1,182 @@ +__device__ __forceinline__ uint64_t wg_desc(const void *p, uint32_t lbo_bytes, + uint32_t sbo_bytes) { + uint32_t a = static_cast(__cvta_generic_to_shared(p)); + uint64_t d = static_cast((a & 0x3FFFFu) >> 4); + d |= static_cast((lbo_bytes & 0x3FFFFu) >> 4) << 16; + d |= static_cast((sbo_bytes & 0x3FFFFu) >> 4) << 32; + return d; +} +#include "wgmma_macro.inc" +constexpr int WG_SLOT = 72; +template +__device__ __noinline__ void wg_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int WG_CLD = WG_BN + 8; + constexpr int WG_NT = 128 * NWG; + constexpr int WG_BM = 64 * NWG; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int SBO_AE = KG * WG_SLOT; + constexpr int LBO_BE = (WG_BN / 8) * WG_SLOT; + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *As0 = reinterpret_cast<__half *>(arena); + __half *Bs0 = As0 + A_E; + __half *As1 = (STAGES == 2) ? Bs0 + B_E : As0; + __half *Bs1 = (STAGES == 2) ? Bs0 + B_E + A_E : Bs0; + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int wgid = tid >> 7; + const int warp = (tid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * WG_BM; + constexpr int WGROWS = 8; + const int n0 = blockIdx.y * WG_BN; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + if (tid < WG_BM) { + const int co = m0 + tid; + if (do_snake && co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + if (tid < WG_BN) { + const int n = n0 + tid; + colok[tid] = (n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = wgid * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + constexpr int ATOT = WG_BM * KG; + constexpr int AREG = (ATOT + WG_NT - 1) / WG_NT; + constexpr int CIB = BK / KT; + constexpr int BTOT = CIB * (WG_BN / 8); + constexpr int BREG = (BTOT + WG_NT - 1) / WG_NT; + int4 ar[AREG]; + int4 br0[BREG], br1[BREG]; + uint32_t edge[BREG]; +#define WG_LOADG(kg0) { const int _ci0 = (kg0) / KT; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; ar[r] = make_int4(0, 0, 0, 0); if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; const int co = m0 + m, kg = (kg0) + v * 8; if (co < Cout && kg < Kg) ar[r] = *reinterpret_cast( w + (size_t)co * Kg + kg); } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; br0[r] = make_int4(0, 0, 0, 0); if constexpr (KT == 3) edge[r] = 0; else br1[r] = make_int4(0, 0, 0, 0); if (BTOT % WG_NT == 0 || idx < BTOT) { const int ci = _ci0 + idx / (WG_BN / 8); const int nc = idx % (WG_BN / 8); if (ci < Cin) { const int4 *src = reinterpret_cast( in + (size_t)ci * SCin + n0 + nc * 8); br0[r] = src[0]; if constexpr (KT == 3) { if (nc == WG_BN / 8 - 1) edge[r] = src[1].x; } else { br1[r] = src[1]; } } } } } + WG_LOADG(0) +#define WG_STS(AD,BD) { _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; *reinterpret_cast( (AD) + (m >> 3) * SBO_AE + v * WG_SLOT + (m & 7) * 8) = ar[r]; } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; if (BTOT % WG_NT == 0 || idx < BTOT) { const int cil = idx / (WG_BN / 8), nc = idx % (WG_BN / 8); unsigned int u[8]; u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; if constexpr (KT == 3) { u[4] = __shfl_down_sync(0xffffffffu, u[0], 1, WG_BN / 8); if (nc == WG_BN / 8 - 1) u[4] = edge[r]; u[5] = u[6] = u[7] = 0; } else { u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; } _Pragma("unroll") for (int k = 0; k < KT; ++k) { const int kh = k >> 1; int4 o; if (k & 1) { o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); } else { o.x = u[kh + 0]; o.y = u[kh + 1]; o.z = u[kh + 2]; o.w = u[kh + 3]; } const int row = cil * KT + k; *reinterpret_cast((BD) + nc * WG_SLOT + (row >> 3) * LBO_BE + (row & 7) * 8) = o; } } } if constexpr (WG_BN == 64) { asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); } } +#define WG_MMA(AD,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t da0 = wg_desc((AD) + wgid * WGROWS * SBO_AE, WG_SLOT * 2, SBO_AE * 2); const uint64_t db0 = wg_desc((BD), LBO_BE * 2, WG_SLOT * 2); _Pragma("unroll") for (int kk = 0; kk < BK; kk += 16) { const uint64_t da = da0 + (kk >> 3) * (WG_SLOT / 8); const uint64_t db = db0 + (kk >> 3) * (LBO_BE / 8); if constexpr (WG_BN == 64) { WGMMA_M64N64K16(acc, da, db); } else { WGMMA_M64N128K16(acc, da, db); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + WG_STS(As0, Bs0) + if (BK < Kg) WG_LOADG(BK) + __syncthreads(); + int cur = 0; + for (int kg0 = 0; kg0 < Kg; kg0 += BK) { + if (cur == 0) { WG_MMA(As0, Bs0) } else { WG_MMA(As1, Bs1) } + if (kg0 + BK < Kg) { + if (STAGES == 2) + asm volatile("wgmma.wait_group.sync.aligned 1;\n" ::: "memory"); + else + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + __syncthreads(); + if (cur == 0) { WG_STS(As1, Bs1) } else { WG_STS(As0, Bs0) } + if (kg0 + 2 * BK < Kg) WG_LOADG(kg0 + 2 * BK) + __syncthreads(); + } else { + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + } + cur ^= 1; + } +#undef WG_LOADG +#undef WG_STS +#undef WG_MMA + __syncthreads(); + const int r0 = er0; + const int cb = (lane & 3) * 2; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Cs + (r0 + h * 8) * WG_CLD + j * 8 + cb) = v; + } + } + __syncthreads(); + constexpr int NC = WG_BN / 4; + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + for (int m = lr; m < WG_BM; m += WG_NT / NC) { + const int co = m0 + m; + if (co >= Cout) continue; + const int n = n0 + lc; + float4 v = *reinterpret_cast(Cs + m * WG_CLD + lc); + __half *op = out + (size_t)co * SCout + n + H; + if (resid != nullptr) { + const __half2 *rp = reinterpret_cast( + resid + (size_t)co * SCout + n + H); + __half2 r0h = rp[0], r1h = rp[1]; + v.x += __half2float(r0h.x); v.y += __half2float(r0h.y); + v.z += __half2float(r1h.x); v.w += __half2float(r1h.y); + } + if (do_snake) { + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + } + if (ok4) { + __half2 o0 = __floats2half2_rn(v.x, v.y); + __half2 o1 = __floats2half2_rn(v.z, v.w); + reinterpret_cast<__half2 *>(op)[0] = o0; + reinterpret_cast<__half2 *>(op)[1] = o1; + } else { + const float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B * H; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int total = Cout * B * pad_len; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} diff --git a/deploy/nemotron_native/kernels/b32/native/codec/wgmma_macro.inc b/deploy/nemotron_native/kernels/b32/native/codec/wgmma_macro.inc new file mode 100644 index 0000000..db624e9 --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/codec/wgmma_macro.inc @@ -0,0 +1,102 @@ +#define WGMMA_M64N8K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " "{" "%0,%1,%2,%3}" ", %4, %5, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N8K16(D,DA,DB) WGMMA_M64N8K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N8K16_ZERO_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " "{" "%0,%1,%2,%3}" ", %4, %5, p, 1, 1, " #TA ", " #TB ";\n}\n" : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) : "l"(DA), "l"(DB), "r"(0)) +#define WGMMA_M64N16K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %10, 0;\n" "wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7}" ", %8, %9, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N16K16(D,DA,DB) WGMMA_M64N16K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N24K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %14, 0;\n" "wgmma.mma_async.sync.aligned.m64n24k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11}" ", %12, %13, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N24K16(D,DA,DB) WGMMA_M64N24K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N32K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %18, 0;\n" "wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15}" ", %16, %17, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N32K16(D,DA,DB) WGMMA_M64N32K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N40K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %22, 0;\n" "wgmma.mma_async.sync.aligned.m64n40k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19}" ", %20, %21, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N40K16(D,DA,DB) WGMMA_M64N40K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N48K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %26, 0;\n" "wgmma.mma_async.sync.aligned.m64n48k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N48K16(D,DA,DB) WGMMA_M64N48K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N56K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %30, 0;\n" "wgmma.mma_async.sync.aligned.m64n56k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N56K16(D,DA,DB) WGMMA_M64N56K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N64K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %34, 0;\n" "wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31}" ", %32, %33, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N64K16(D,DA,DB) WGMMA_M64N64K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N72K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %38, 0;\n" "wgmma.mma_async.sync.aligned.m64n72k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35}" ", %36, %37, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N72K16(D,DA,DB) WGMMA_M64N72K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N80K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %42, 0;\n" "wgmma.mma_async.sync.aligned.m64n80k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39}" ", %40, %41, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N80K16(D,DA,DB) WGMMA_M64N80K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N88K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %46, 0;\n" "wgmma.mma_async.sync.aligned.m64n88k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43}" ", %44, %45, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N88K16(D,DA,DB) WGMMA_M64N88K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N96K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %50, 0;\n" "wgmma.mma_async.sync.aligned.m64n96k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47}" ", %48, %49, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N96K16(D,DA,DB) WGMMA_M64N96K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N104K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %54, 0;\n" "wgmma.mma_async.sync.aligned.m64n104k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51}" ", %52, %53, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N104K16(D,DA,DB) WGMMA_M64N104K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N112K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %58, 0;\n" "wgmma.mma_async.sync.aligned.m64n112k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55}" ", %56, %57, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N112K16(D,DA,DB) WGMMA_M64N112K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N120K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %62, 0;\n" "wgmma.mma_async.sync.aligned.m64n120k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59}" ", %60, %61, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N120K16(D,DA,DB) WGMMA_M64N120K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N128K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %66, 0;\n" "wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63}" ", %64, %65, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N128K16(D,DA,DB) WGMMA_M64N128K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N136K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %70, 0;\n" "wgmma.mma_async.sync.aligned.m64n136k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67}" ", %68, %69, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N136K16(D,DA,DB) WGMMA_M64N136K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N144K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %74, 0;\n" "wgmma.mma_async.sync.aligned.m64n144k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71}" ", %72, %73, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N144K16(D,DA,DB) WGMMA_M64N144K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N152K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %78, 0;\n" "wgmma.mma_async.sync.aligned.m64n152k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75}" ", %76, %77, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N152K16(D,DA,DB) WGMMA_M64N152K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N160K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %82, 0;\n" "wgmma.mma_async.sync.aligned.m64n160k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79}" ", %80, %81, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N160K16(D,DA,DB) WGMMA_M64N160K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N168K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %86, 0;\n" "wgmma.mma_async.sync.aligned.m64n168k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83}" ", %84, %85, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N168K16(D,DA,DB) WGMMA_M64N168K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N176K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %90, 0;\n" "wgmma.mma_async.sync.aligned.m64n176k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87}" ", %88, %89, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N176K16(D,DA,DB) WGMMA_M64N176K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N184K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %94, 0;\n" "wgmma.mma_async.sync.aligned.m64n184k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91}" ", %92, %93, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N184K16(D,DA,DB) WGMMA_M64N184K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N192K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %98, 0;\n" "wgmma.mma_async.sync.aligned.m64n192k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95}" ", %96, %97, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N192K16(D,DA,DB) WGMMA_M64N192K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N200K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %102, 0;\n" "wgmma.mma_async.sync.aligned.m64n200k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99}" ", %100, %101, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N200K16(D,DA,DB) WGMMA_M64N200K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N208K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %106, 0;\n" "wgmma.mma_async.sync.aligned.m64n208k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103}" ", %104, %105, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N208K16(D,DA,DB) WGMMA_M64N208K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N216K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %110, 0;\n" "wgmma.mma_async.sync.aligned.m64n216k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107}" ", %108, %109, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N216K16(D,DA,DB) WGMMA_M64N216K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N224K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %114, 0;\n" "wgmma.mma_async.sync.aligned.m64n224k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111}" ", %112, %113, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N224K16(D,DA,DB) WGMMA_M64N224K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N232K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %118, 0;\n" "wgmma.mma_async.sync.aligned.m64n232k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115}" ", %116, %117, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N232K16(D,DA,DB) WGMMA_M64N232K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N240K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %122, 0;\n" "wgmma.mma_async.sync.aligned.m64n240k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119}" ", %120, %121, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N240K16(D,DA,DB) WGMMA_M64N240K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N248K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %126, 0;\n" "wgmma.mma_async.sync.aligned.m64n248k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123}" ", %124, %125, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N248K16(D,DA,DB) WGMMA_M64N248K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N256K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %130, 0;\n" "wgmma.mma_async.sync.aligned.m64n256k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123,%124,%125,%126,%127}" ", %128, %129, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]), "+f"(D[124]), "+f"(D[125]), "+f"(D[126]), "+f"(D[127]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N256K16(D,DA,DB) WGMMA_M64N256K16_T(D, DA, DB, 0, 1) +#define WGMMA_ANY_T(NW,D,DA,DB,TA,TB) do { if constexpr (NW == 8) { WGMMA_M64N8K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 16) { WGMMA_M64N16K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 24) { WGMMA_M64N24K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 32) { WGMMA_M64N32K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 40) { WGMMA_M64N40K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 48) { WGMMA_M64N48K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 56) { WGMMA_M64N56K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 64) { WGMMA_M64N64K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 72) { WGMMA_M64N72K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 80) { WGMMA_M64N80K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 88) { WGMMA_M64N88K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 96) { WGMMA_M64N96K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 104) { WGMMA_M64N104K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 112) { WGMMA_M64N112K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 120) { WGMMA_M64N120K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 128) { WGMMA_M64N128K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 136) { WGMMA_M64N136K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 144) { WGMMA_M64N144K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 152) { WGMMA_M64N152K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 160) { WGMMA_M64N160K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 168) { WGMMA_M64N168K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 176) { WGMMA_M64N176K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 184) { WGMMA_M64N184K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 192) { WGMMA_M64N192K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 200) { WGMMA_M64N200K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 208) { WGMMA_M64N208K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 216) { WGMMA_M64N216K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 224) { WGMMA_M64N224K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 232) { WGMMA_M64N232K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 240) { WGMMA_M64N240K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 248) { WGMMA_M64N248K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 256) { WGMMA_M64N256K16_T(D, DA, DB, TA, TB); } } while (0) +#define WGMMA_ANY(NW,D,DA,DB) WGMMA_ANY_T(NW, D, DA, DB, 0, 1) +#define WGMMA_RS_M64N96K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %53, 0;\n" "wgmma.mma_async.sync.aligned.m64n96k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47}" ", {%48,%49,%50,%51}, %52, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N104K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %57, 0;\n" "wgmma.mma_async.sync.aligned.m64n104k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51}" ", {%52,%53,%54,%55}, %56, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N112K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %61, 0;\n" "wgmma.mma_async.sync.aligned.m64n112k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55}" ", {%56,%57,%58,%59}, %60, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N120K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %65, 0;\n" "wgmma.mma_async.sync.aligned.m64n120k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59}" ", {%60,%61,%62,%63}, %64, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N128K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %69, 0;\n" "wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63}" ", {%64,%65,%66,%67}, %68, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_ANY_T(NW,D,A,DB,TB) do { if constexpr (NW == 96) { WGMMA_RS_M64N96K16_T(D, A, DB, TB); } else if constexpr (NW == 104) { WGMMA_RS_M64N104K16_T(D, A, DB, TB); } else if constexpr (NW == 112) { WGMMA_RS_M64N112K16_T(D, A, DB, TB); } else if constexpr (NW == 120) { WGMMA_RS_M64N120K16_T(D, A, DB, TB); } else if constexpr (NW == 128) { WGMMA_RS_M64N128K16_T(D, A, DB, TB); } } while (0) +#define WGMMA_RS_M64N8K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %9, 0;\n" "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " "{" "%0,%1,%2,%3}" ", {%4,%5,%6,%7}, %8, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N16K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %13, 0;\n" "wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7}" ", {%8,%9,%10,%11}, %12, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N24K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %17, 0;\n" "wgmma.mma_async.sync.aligned.m64n24k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11}" ", {%12,%13,%14,%15}, %16, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N32K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %21, 0;\n" "wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15}" ", {%16,%17,%18,%19}, %20, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N40K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %25, 0;\n" "wgmma.mma_async.sync.aligned.m64n40k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N48K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %29, 0;\n" "wgmma.mma_async.sync.aligned.m64n48k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N56K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %33, 0;\n" "wgmma.mma_async.sync.aligned.m64n56k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N64K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %37, 0;\n" "wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31}" ", {%32,%33,%34,%35}, %36, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N72K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %41, 0;\n" "wgmma.mma_async.sync.aligned.m64n72k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35}" ", {%36,%37,%38,%39}, %40, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N80K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %45, 0;\n" "wgmma.mma_async.sync.aligned.m64n80k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39}" ", {%40,%41,%42,%43}, %44, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N88K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %49, 0;\n" "wgmma.mma_async.sync.aligned.m64n88k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43}" ", {%44,%45,%46,%47}, %48, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N136K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %73, 0;\n" "wgmma.mma_async.sync.aligned.m64n136k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67}" ", {%68,%69,%70,%71}, %72, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N144K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %77, 0;\n" "wgmma.mma_async.sync.aligned.m64n144k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71}" ", {%72,%73,%74,%75}, %76, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N152K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %81, 0;\n" "wgmma.mma_async.sync.aligned.m64n152k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75}" ", {%76,%77,%78,%79}, %80, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N160K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %85, 0;\n" "wgmma.mma_async.sync.aligned.m64n160k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79}" ", {%80,%81,%82,%83}, %84, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N168K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %89, 0;\n" "wgmma.mma_async.sync.aligned.m64n168k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83}" ", {%84,%85,%86,%87}, %88, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N176K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %93, 0;\n" "wgmma.mma_async.sync.aligned.m64n176k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87}" ", {%88,%89,%90,%91}, %92, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N184K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %97, 0;\n" "wgmma.mma_async.sync.aligned.m64n184k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91}" ", {%92,%93,%94,%95}, %96, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N192K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %101, 0;\n" "wgmma.mma_async.sync.aligned.m64n192k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95}" ", {%96,%97,%98,%99}, %100, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N200K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %105, 0;\n" "wgmma.mma_async.sync.aligned.m64n200k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99}" ", {%100,%101,%102,%103}, %104, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N208K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %109, 0;\n" "wgmma.mma_async.sync.aligned.m64n208k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103}" ", {%104,%105,%106,%107}, %108, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N216K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %113, 0;\n" "wgmma.mma_async.sync.aligned.m64n216k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107}" ", {%108,%109,%110,%111}, %112, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N224K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %117, 0;\n" "wgmma.mma_async.sync.aligned.m64n224k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111}" ", {%112,%113,%114,%115}, %116, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N232K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %121, 0;\n" "wgmma.mma_async.sync.aligned.m64n232k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115}" ", {%116,%117,%118,%119}, %120, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N240K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %125, 0;\n" "wgmma.mma_async.sync.aligned.m64n240k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119}" ", {%120,%121,%122,%123}, %124, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N248K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %129, 0;\n" "wgmma.mma_async.sync.aligned.m64n248k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123}" ", {%124,%125,%126,%127}, %128, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N256K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %133, 0;\n" "wgmma.mma_async.sync.aligned.m64n256k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123,%124,%125,%126,%127}" ", {%128,%129,%130,%131}, %132, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]), "+f"(D[124]), "+f"(D[125]), "+f"(D[126]), "+f"(D[127]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#undef WGMMA_RS_ANY_T +#define WGMMA_RS_ANY_T(NW,D,A,DB,TB) do { if constexpr (NW == 8) { WGMMA_RS_M64N8K16_T(D, A, DB, TB); } else if constexpr (NW == 16) { WGMMA_RS_M64N16K16_T(D, A, DB, TB); } else if constexpr (NW == 24) { WGMMA_RS_M64N24K16_T(D, A, DB, TB); } else if constexpr (NW == 32) { WGMMA_RS_M64N32K16_T(D, A, DB, TB); } else if constexpr (NW == 40) { WGMMA_RS_M64N40K16_T(D, A, DB, TB); } else if constexpr (NW == 48) { WGMMA_RS_M64N48K16_T(D, A, DB, TB); } else if constexpr (NW == 56) { WGMMA_RS_M64N56K16_T(D, A, DB, TB); } else if constexpr (NW == 64) { WGMMA_RS_M64N64K16_T(D, A, DB, TB); } else if constexpr (NW == 72) { WGMMA_RS_M64N72K16_T(D, A, DB, TB); } else if constexpr (NW == 80) { WGMMA_RS_M64N80K16_T(D, A, DB, TB); } else if constexpr (NW == 88) { WGMMA_RS_M64N88K16_T(D, A, DB, TB); } else if constexpr (NW == 96) { WGMMA_RS_M64N96K16_T(D, A, DB, TB); } else if constexpr (NW == 104) { WGMMA_RS_M64N104K16_T(D, A, DB, TB); } else if constexpr (NW == 112) { WGMMA_RS_M64N112K16_T(D, A, DB, TB); } else if constexpr (NW == 120) { WGMMA_RS_M64N120K16_T(D, A, DB, TB); } else if constexpr (NW == 128) { WGMMA_RS_M64N128K16_T(D, A, DB, TB); } else if constexpr (NW == 136) { WGMMA_RS_M64N136K16_T(D, A, DB, TB); } else if constexpr (NW == 144) { WGMMA_RS_M64N144K16_T(D, A, DB, TB); } else if constexpr (NW == 152) { WGMMA_RS_M64N152K16_T(D, A, DB, TB); } else if constexpr (NW == 160) { WGMMA_RS_M64N160K16_T(D, A, DB, TB); } else if constexpr (NW == 168) { WGMMA_RS_M64N168K16_T(D, A, DB, TB); } else if constexpr (NW == 176) { WGMMA_RS_M64N176K16_T(D, A, DB, TB); } else if constexpr (NW == 184) { WGMMA_RS_M64N184K16_T(D, A, DB, TB); } else if constexpr (NW == 192) { WGMMA_RS_M64N192K16_T(D, A, DB, TB); } else if constexpr (NW == 200) { WGMMA_RS_M64N200K16_T(D, A, DB, TB); } else if constexpr (NW == 208) { WGMMA_RS_M64N208K16_T(D, A, DB, TB); } else if constexpr (NW == 216) { WGMMA_RS_M64N216K16_T(D, A, DB, TB); } else if constexpr (NW == 224) { WGMMA_RS_M64N224K16_T(D, A, DB, TB); } else if constexpr (NW == 232) { WGMMA_RS_M64N232K16_T(D, A, DB, TB); } else if constexpr (NW == 240) { WGMMA_RS_M64N240K16_T(D, A, DB, TB); } else if constexpr (NW == 248) { WGMMA_RS_M64N248K16_T(D, A, DB, TB); } else if constexpr (NW == 256) { WGMMA_RS_M64N256K16_T(D, A, DB, TB); } } while (0) diff --git a/deploy/nemotron_native/kernels/b32/native/prefill/dedup_native_prefill_impl.inc b/deploy/nemotron_native/kernels/b32/native/prefill/dedup_native_prefill_impl.inc new file mode 100644 index 0000000..5d09d57 --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/prefill/dedup_native_prefill_impl.inc @@ -0,0 +1,249 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 16 || BATCH == 32 || BATCH == 64)) + throw std::runtime_error("native prefill requires matching batch 1/16/32/64"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (padded > INT_MAX / ((BB_EXPERTS + 1) * D_MODEL)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + float *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(float))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(float))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + bf16 *private_state = nullptr; + mk2_check(owned.alloc(&private_state, size_t(BATCH) * ssm_elements * sizeof(bf16))); + mk2_check(cudaMemset(private_state, 0, size_t(BATCH) * ssm_elements * sizeof(bf16))); + st->mb_host[i].b32_private = private_state; + binding.ssm_state[i] = reinterpret_cast(private_state); + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + mk2_check(owned.alloc(&st->mb_host[i].r8_canonical_ssm, ssm_elements * sizeof(float))); + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_dtda, 2 * MB_HEADS); + { + const size_t estride = padded; + const size_t grows = size_t(BB_EXPERTS) * estride; + mk2_check(owned.alloc(&hb.pf_moe_gath, grows * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_gath, 0, grows * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_hc, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_hc, 0, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_gidx, grows * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_gidx, 0, grows * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ecnt, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ecnt, 0, BB_EXPERTS * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ntiles, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ntiles, 0, BB_EXPERTS * sizeof(int))); + hb.pf_moe_estride = int(estride); + hb.pf_moe_gath_v = act_gl{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_v = act_gl{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + hb.pf_moe_gath_vn = act_gl_n{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_vn = act_gl_n{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + scratch_bytes += grows * (D_MODEL + BB_MOE_INTER) * sizeof(bf16); + } + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(owned.alloc(&hb.native_pf_cohort_check, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(hb.native_pf_cohort_check, 0, 2 * sizeof(unsigned int))); +#if PF_ACTIVE_TOKEN_EXPERTS + mk2_check(owned.alloc(&hb.native_pf_expert_masks, (padded/BATCH_PAD)*sizeof(unsigned int))); +#endif + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} +extern "C" int mk2_native_prefill_effective_sequences(MK2NativePrefill *p, int *out) { + try { + if (!p || !out) throw std::runtime_error("null effective-sequence query"); + auto *st = reinterpret_cast(p->base); + mk2_check(cudaMemcpy(out, st->host_bb.native_pf_cohort_check + 1, + sizeof(int), cudaMemcpyDeviceToHost)); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" int mk2_r8_initial_state(MK2NativePrefill *p, void **initial, void **ready) { + if (!p || !initial || !ready) return 1; + auto *st = reinterpret_cast(p->base); + for (int i = 0; i < 14; ++i) initial[i] = st->mb_host[i].r8_canonical_ssm; + *ready = st->host_bb.native_pf_cohort_check + 1; + return 0; +} diff --git a/deploy/nemotron_native/kernels/b32/native/prefill/dedup_predictor_mk2.cuh b/deploy/nemotron_native/kernels/b32/native/prefill/dedup_predictor_mk2.cuh new file mode 100644 index 0000000..27b90af --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/prefill/dedup_predictor_mk2.cuh @@ -0,0 +1,4600 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef PF_ACTIVE_TOKEN_EXPERTS +#define PF_ACTIVE_TOKEN_EXPERTS 0 +#endif +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +#ifndef MK2_PREFILL_ONLY +#define MK2_PREFILL_ONLY 0 +#endif +#ifndef MK2_BATCH_TILE +#define MK2_BATCH_TILE (((MK2_BATCH + 15) / 16) * 16) +#endif +constexpr int BATCH_PAD = MK2_BATCH_TILE; +static_assert(BATCH >= 1 && BATCH <= BATCH_PAD); +static_assert(BATCH_PAD == 16 || BATCH_PAD == 32 || BATCH_PAD == 64 || BATCH_PAD == 128, + "native prefill tile must be 16, 32, 64 or 128"); +static_assert(MK2_PREFILL_ONLY || BATCH_PAD == ((BATCH + 15) / 16) * 16, + "a wider token tile requires a prefill-only library"); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_PW_MULTI +#define PF_SSM_PW_MULTI 4 +#endif +#ifndef PF_SSM_PW_SINGLE +#define PF_SSM_PW_SINGLE 2 +#endif +static_assert(MB_HEAD_DIM % PF_SSM_PW_MULTI == 0 && + MB_HEAD_DIM % PF_SSM_PW_SINGLE == 0, + "head_dim must split evenly into scan channel groups"); +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef MK2_SMEM_SLACK +#define MK2_SMEM_SLACK 1024 +#endif +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +#ifndef MK2_THREADS +#define MK2_THREADS 256 +#endif +constexpr int THREADS = MK2_THREADS; +constexpr int NORM_THREADS = 256; +constexpr int NORM_WARPS = NORM_THREADS / 32; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +#ifndef MK2_GEMM_CONSUMERS +#define MK2_GEMM_CONSUMERS 1 +#endif +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 +#define MK2_HEAD_STAGE_BYTES 0 +#else +#define MK2_HEAD_STAGE_BYTES sizeof(out_tile_f) +#endif +constexpr int NCONS = MK2_GEMM_CONSUMERS; +static_assert(NCONS >= 1 && NCONS <= 3, "consumers occupy warpgroups 0, 2 and 3"); +static_assert(THREADS >= (NCONS + 1) * 128, "one producer plus NCONS consumers"); +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +constexpr int PAD_N = 32; +using act_tile_n = st_bf; +using out_tile_n = st_bf; +using act_gl_n = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_MOE_ACC_C, EPI_ACC_F32 }; +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; + float *r8_canonical_ssm = nullptr; + bf16 *b32_private = nullptr; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + unsigned int *native_pf_expert_masks = nullptr; + float *pf_dtda = nullptr; + bf16 *pf_moe_gath = nullptr; + bf16 *pf_moe_hc = nullptr; + int *pf_moe_gidx = nullptr; + int *pf_moe_ecnt = nullptr; + int *pf_moe_ntiles = nullptr; + int pf_moe_estride = 0; + act_gl pf_moe_gath_v; + act_gl pf_moe_hc_v; + act_gl_n pf_moe_gath_vn; + act_gl_n pf_moe_hc_vn; + unsigned int *native_pf_cohort_check = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + volatile unsigned int *generation = state + 1; + const unsigned int seen = *generation; + if (atomicAdd(state, 1u) == gridDim.x - 1) { + atomicExch(state, 0u); + atomicAdd(state + 1, 1u); + } else { + while (*generation == seen) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const bf16 *in = src + (size_t)row * D_MODEL; + float v[PER_THREAD]; + float sum = 0.f, sq = 0.f; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (threadIdx.x < NORM_THREADS) { + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + v[j] = __bfloat162float(in[threadIdx.x + j * NORM_THREADS]); + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + if (lane == 0) { scratch[warp] = sum; scratch[NORM_WARPS + warp] = sq; } + } + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f, q = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) { s += scratch[w]; q += scratch[NORM_WARPS + w]; } + scratch[2 * NORM_WARPS] = s / D_MODEL; + scratch[2 * NORM_WARPS + 1] = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + } + __syncthreads(); + const float mean = scratch[2 * NORM_WARPS], inv = scratch[2 * NORM_WARPS + 1]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + if (threadIdx.x >= NORM_THREADS) break; + const int i = threadIdx.x + j * NORM_THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * weight[i]); + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, + int idx, unsigned mask); +__device__ __forceinline__ void norm_sub_sync(int sub) { + if (THREADS == NORM_THREADS) { __syncthreads(); return; } + asm volatile("bar.sync %0, %1;" :: "r"(6 + sub), "n"(NORM_THREADS) : "memory"); +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, const float *fold_src = nullptr, + const float *fold_routes = nullptr) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + constexpr int NSUB = THREADS / NORM_THREADS; + if (threadIdx.x >= NSUB * NORM_THREADS) return; + const int sub = threadIdx.x / NORM_THREADS; + const int tid = threadIdx.x - sub * NORM_THREADS; + float *const sc = scratch + sub * (NORM_WARPS + 1); + for (int row = blockIdx.x * NSUB + sub; row < n_rows; + row += gridDim.x * NSUB) { + const size_t base = (size_t)row * D_MODEL; + unsigned fold_mask = (1u << BB_EXPERTS) - 1u; + if (fold_src && fold_routes) { + fold_mask = 0u; + const float *rw = fold_routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) + if (rw[e] != 0.f) fold_mask |= 1u << e; + } + float v[PER_THREAD]; + float sq = 0.f; + const int warp = tid >> 5, lane = tid & 31; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + bf16 hv; + if (fold_src) { + hv = bb_moe_fold_deterministic(fold_src, (int)(base + i), fold_mask); + h[base + i] = hv; + } else { + hv = h[base + i]; + } + v[j] = __bfloat162float(hv) + __bfloat162float(r[base + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + if (lane == 0) sc[warp] = sq; + norm_sub_sync(sub); + if (tid == 0) { + float s = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) s += sc[w]; + sc[NORM_WARPS] = rsqrtf(s / D_MODEL + RMS_EPS); + } + norm_sub_sync(sub); + const float inv = sc[NORM_WARPS]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + out[base + i] = __float2bfloat16( + __bfloat162float(weight[i]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + } + norm_sub_sync(sub); + } +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + const int c = idx - row * MB_CONV_DIM; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + const float raw = __bfloat162float( + zxbcdt[(size_t)row * MB_IN_PROJ + MB_INNER + c]); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + xbc_act[(size_t)row * MB_CONV_DIM + c] = __float2bfloat16(acc); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + const int pairs = BATCH * MB_HEADS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_HEADS * split); + const int head = (pair / split) % MB_HEADS; + const int first_channel = (pair % split) * channels; + const int group = head / MB_HEADS_PER_GROUP; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + L.dt_bias[head]); + const float dA = __expf(L.A[head] * dt); + const float Dh = L.D[head]; + float *st = L.ssm_state + (size_t)slot_map[row] * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + for (int p = first_channel; p < first_channel + channels; ++p) { + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + const size_t at = (size_t)p * MB_STATE + n; + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + static_assert(MB_CONV_K == 4, "packed conv taps assume a 4-tap kernel"); + constexpr int CPT = 4; + static_assert(MB_CONV_DIM % CPT == 0, "channel vector width"); + constexpr int CGROUPS = MB_CONV_DIM / CPT; + const int total = n_tokens * CGROUPS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / CGROUPS; + const int c = (idx - t * CGROUPS) * CPT; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + float wv[CPT][MB_CONV_K]; + { + const float4 *wsrc = reinterpret_cast( + L.conv_w + (size_t)c * MB_CONV_K); + #pragma unroll + for (int j = 0; j < CPT / 2; ++j) { + const float4 w4 = wsrc[j]; + const unsigned uu0 = __float_as_uint(w4.x); + const unsigned uu1 = __float_as_uint(w4.y); + const unsigned uu2 = __float_as_uint(w4.z); + const unsigned uu3 = __float_as_uint(w4.w); + wv[2 * j][0] = __uint_as_float(uu0 << 16); + wv[2 * j][1] = __uint_as_float(uu0 & 0xffff0000u); + wv[2 * j][2] = __uint_as_float(uu1 << 16); + wv[2 * j][3] = __uint_as_float(uu1 & 0xffff0000u); + wv[2 * j + 1][0] = __uint_as_float(uu2 << 16); + wv[2 * j + 1][1] = __uint_as_float(uu2 & 0xffff0000u); + wv[2 * j + 1][2] = __uint_as_float(uu3 << 16); + wv[2 * j + 1][3] = __uint_as_float(uu3 & 0xffff0000u); + } + } + float acc[CPT]; + { + const uint2 b2 = *reinterpret_cast(L.conv_b + c); + acc[0] = __uint_as_float(b2.x << 16); + acc[1] = __uint_as_float(b2.x & 0xffff0000u); + acc[2] = __uint_as_float(b2.y << 16); + acc[3] = __uint_as_float(b2.y & 0xffff0000u); + } + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v[CPT]; + if (src < 0) { + #pragma unroll + for (int q = 0; q < CPT; ++q) + v[q] = __bfloat162float( + state[(size_t)q * (MB_CONV_K - 1) + pos + k]); + } else { + const uint2 p = *reinterpret_cast( + zxbcdt + (size_t)src * MB_IN_PROJ + MB_INNER + c); + v[0] = __uint_as_float(p.x << 16); + v[1] = __uint_as_float(p.x & 0xffff0000u); + v[2] = __uint_as_float(p.y << 16); + v[3] = __uint_as_float(p.y & 0xffff0000u); + } + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] += wv[q][k] * v[q]; + } + uint2 out; + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] = acc[q] / (1.f + __expf(-acc[q])); + out.x = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[0])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[1])) << 16); + out.y = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[2])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[3])) << 16); + *reinterpret_cast(xbc_act + (size_t)t * MB_CONV_DIM + c) = out; + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_dtda_prefill( + const MambaLayer &L, const bf16 *zxbcdt, float *dtda, int n_tokens) { + const int total = n_tokens * MB_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_HEADS; + const int head = idx - t * MB_HEADS; + const float dt = mamba_softplus( + __bfloat162float(zxbcdt[(size_t)t * MB_IN_PROJ + MB_INNER + + MB_CONV_DIM + head]) + L.dt_bias[head]); + dtda[(size_t)idx * 2] = dt; + dtda[(size_t)idx * 2 + 1] = __expf(L.A[head] * dt); + } +} +template +__device__ __forceinline__ void mamba_ssm_prefill_impl( + const MambaLayer &L, const bf16 *xbc_act, const float *dtda, bf16 *gated, + int n_pos, int n_seq, int slot, const int *seq_slot) { + constexpr int NS = MB_STATE / 32; + constexpr int PG = MB_HEAD_DIM / PW; + constexpr bool RS = (PW % 4) == 0; + constexpr int H1 = RS ? PW / 2 : 1; + constexpr int H2 = RS ? PW / 4 : 1; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const bool keep_hi = (lane & 16) != 0; + const bool keep_hi2 = (lane & 8) != 0; + const int warps_total = gridDim.x * WARPS; + const int units = n_seq * MB_HEADS * PG; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int pg = unit % PG; + const int rest = unit / PG; + const int head = rest % MB_HEADS; + const int seq = rest / MB_HEADS; + const int p0 = pg * PW; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Dh = L.D[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p0 * MB_STATE; + float sreg[PW][NS]; + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + sreg[q][j] = st[(size_t)q * MB_STATE + lane + j * 32]; + } + auto fetch = [&](int p, float (&oB)[NS], float (&oC)[NS], + float (&oX)[PW], float2 &od) { + const size_t tt = (size_t)p * n_seq + seq; + const bf16 *xa2 = xbc_act + tt * MB_CONV_DIM; + const bf16 *Bs = xa2 + MB_INNER + group * MB_STATE; + const bf16 *Cs = xa2 + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + od = *reinterpret_cast( + dtda + (tt * MB_HEADS + head) * 2); + #pragma unroll + for (int j = 0; j < NS; ++j) { + oB[j] = __bfloat162float(Bs[lane + j * 32]); + oC[j] = __bfloat162float(Cs[lane + j * 32]); + } + if constexpr (RS) { + const uint2 *xp = reinterpret_cast( + xa2 + head * MB_HEAD_DIM + p0); + #pragma unroll + for (int q = 0; q < PW; q += 4) { + const uint2 w = xp[q >> 2]; + oX[q] = __uint_as_float(w.x << 16); + oX[q + 1] = __uint_as_float(w.x & 0xffff0000u); + oX[q + 2] = __uint_as_float(w.y << 16); + oX[q + 3] = __uint_as_float(w.y & 0xffff0000u); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) + oX[q] = __bfloat162float(xa2[head * MB_HEAD_DIM + p0 + q]); + } + }; + float pB[NS], pC[NS], pX[PW]; + float2 pd; + fetch(0, pB, pC, pX, pd); + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const float dt = pd.x; + const float dA = pd.y; + float Bv[NS], Cv[NS]; + #pragma unroll + for (int j = 0; j < NS; ++j) { + Bv[j] = pB[j] * dt; + Cv[j] = pC[j]; + } + float acc[PW], xvs[PW]; + #pragma unroll + for (int q = 0; q < PW; ++q) xvs[q] = pX[q]; + if (pos + 1 < n_pos) fetch(pos + 1, pB, pC, pX, pd); + #pragma unroll + for (int q = 0; q < PW; ++q) { + const float xv = xvs[q]; + float a = 0.f; + #pragma unroll + for (int j = 0; j < NS; ++j) { + sreg[q][j] = sreg[q][j] * dA + Bv[j] * xv; + a += sreg[q][j] * Cv[j]; + } + acc[q] = a; + } + if constexpr (RS) { + float m1[H1], t1[H1], x1[H1]; + #pragma unroll + for (int q = 0; q < H1; ++q) { + m1[q] = keep_hi ? acc[q + H1] : acc[q]; + t1[q] = keep_hi ? acc[q] : acc[q + H1]; + x1[q] = keep_hi ? xvs[q + H1] : xvs[q]; + } + #pragma unroll + for (int q = 0; q < H1; ++q) + m1[q] += __shfl_xor_sync(0xffffffff, t1[q], 16); + float m2[H2], t2[H2], x2[H2]; + #pragma unroll + for (int q = 0; q < H2; ++q) { + m2[q] = keep_hi2 ? m1[q + H2] : m1[q]; + t2[q] = keep_hi2 ? m1[q] : m1[q + H2]; + x2[q] = keep_hi2 ? x1[q + H2] : x1[q]; + } + #pragma unroll + for (int q = 0; q < H2; ++q) + m2[q] += __shfl_xor_sync(0xffffffff, t2[q], 8); + #pragma unroll + for (int q = 0; q < H2; ++q) { + #pragma unroll + for (int off = 4; off > 0; off >>= 1) + m2[q] += __shfl_xor_sync(0xffffffff, m2[q], off); + } + if ((lane & 7) == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + + p0 + (keep_hi ? H1 : 0) + (keep_hi2 ? H2 : 0); + #pragma unroll + for (int q = 0; q < H2; ++q) + dst[q] = __float2bfloat16(m2[q] + x2[q] * Dh); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc[q] += __shfl_xor_sync(0xffffffff, acc[q], off); + } + if (lane == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + p0; + #pragma unroll + for (int q = 0; q < PW; ++q) + dst[q] = __float2bfloat16(acc[q] + xvs[q] * Dh); + } + } + } + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + st[(size_t)q * MB_STATE + lane + j * 32] = sreg[q][j]; + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, + const float *dtda, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + (void)zxbcdt; + if (n_seq > 1) + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens / n_seq, n_seq, slot, seq_slot); + else + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens, 1, slot, nullptr); +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16( + held[j] * inv * __bfloat162float(L.gnorm_w[idx])); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = BATCH * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / BB_ATT_HEADS; + const int head = pair % BB_ATT_HEADS; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const int kv = head / kv_per_q; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *q = base + head * BB_HEAD_DIM; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (head % kv_per_q == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + for (int tpos = 0; tpos < T; ++tpos) { +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new + : kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#endif + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m, s); + const float rescale = __expf(m - m_new); + const float w = __expf(s - m_new); + denom = denom * rescale + w; +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *vt = kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#else + const bf16 *vt = cur_token + ? v_new + : kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + w * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + constexpr int CHUNKS = BB_HEAD_DIM / 8; + const int total = n_tokens * BB_KV_HEADS * CHUNKS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int c = idx % CHUNKS; + const int rest = idx / CHUNKS; + const int t = rest / BB_KV_HEADS; + const int kv = rest - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM + c * 8; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM + c * 8; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv) + c * 8; + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv) + c * 8; + *reinterpret_cast(kc) = *reinterpret_cast(k_new); + *reinterpret_cast(vc) = *reinterpret_cast(v_new); + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float qv[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qv[j] = __bfloat162float(q[lane + j * 32]); + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + const int *const krow = + L.block_table + ((size_t)seq * 2 + 0) * L.max_blocks_per_seq; + const int *const vrow = + L.block_table + ((size_t)seq * 2 + 1) * L.max_blocks_per_seq; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + const int head_bias = kv * L.tokens_per_block * BB_HEAD_DIM; + int page = -1; + const bf16 *kpage = nullptr, *vpage = nullptr; + float kn[PER_LANE], vn[PER_LANE]; + { + page = 0; + kpage = L.pool_base + (size_t)(krow[0] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + (size_t)(vrow[0] & 0x7FFFFFFF) * elems_per_block; + const int at = head_bias + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + for (int s = 0; s <= pos; ++s) { + float kc[PER_LANE], vc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { kc[j] = kn[j]; vc[j] = vn[j]; } + if (s < pos) { + const int np = (s + 1) / L.tokens_per_block; + if (np != page) { + page = np; + kpage = L.pool_base + + (size_t)(krow[np] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + + (size_t)(vrow[np] & 0x7FFFFFFF) * elems_per_block; + } + const int at = head_bias + + ((s + 1) % L.tokens_per_block) * BB_HEAD_DIM + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) dot += qv[j] * kc[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[j] = acc[j] * rescale + wt * vc[j]; + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH, unsigned int *tile_masks = nullptr, + bf16 *gath = nullptr, int *gidx = nullptr, int *ecnt = nullptr, + int estride = 0, bf16 *row_stage = nullptr) { + __shared__ float score[BB_EXPERTS]; + __shared__ int s_pick[BB_TOP_K], s_slot[BB_TOP_K]; + constexpr int CH = D_MODEL / 8; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + { + const float4 *gsrc = + reinterpret_cast(normed + (size_t)row * D_MODEL); + float4 *gdst = reinterpret_cast(row_stage); + for (int i = threadIdx.x; i < CH; i += THREADS) gdst[i] = gsrc[i]; + } + if ((int)threadIdx.x < BB_EXPERTS) + moe_w[(size_t)row * BB_EXPERTS + threadIdx.x] = 0.f; + __syncthreads(); + for (int e = warp; e < BB_EXPERTS; e += WARPS) { + const bf16 *wr = gate_w + (size_t)e * D_MODEL; + float acc = 0.f; + for (int i = lane; i < D_MODEL; i += 32) + acc += __bfloat162float(row_stage[i]) * + __bfloat162float(wr[i]); + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) score[e] = 1.f / (1.f + __expf(-acc)); + } + __syncthreads(); + if (threadIdx.x == 0) { + float remaining[BB_EXPERTS]; + int pick[BB_TOP_K]; + for (int e = 0; e < BB_EXPERTS; ++e) + remaining[e] = score[e] + __bfloat162float(gate_bias[e]); + float total = 0.f; + for (int k = 0; k < BB_TOP_K; ++k) { + int best = 0; + for (int e = 1; e < BB_EXPERTS; ++e) + if (remaining[e] > remaining[best]) best = e; + pick[k] = best; + total += score[best]; + remaining[best] = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + unsigned int bits = 0u; + for (int k = 0; k < BB_TOP_K; ++k) { + moe_w[row * BB_EXPERTS + pick[k]] = score[pick[k]] * norm; + bits |= 1u << pick[k]; + s_pick[k] = pick[k]; + } + if (tile_masks) atomicOr(&tile_masks[row / BATCH_PAD], bits); + } + __syncthreads(); + if (gath) { + if ((int)threadIdx.x < BB_TOP_K) { + const int k = threadIdx.x; + const int slot = atomicAdd(&ecnt[s_pick[k]], 1); + s_slot[k] = slot; + gidx[(size_t)s_pick[k] * estride + slot] = row; + } + __syncthreads(); + const float4 *src = reinterpret_cast(row_stage); + for (int k = 0; k < BB_TOP_K; ++k) { + float4 *dst = reinterpret_cast( + gath + ((size_t)s_pick[k] * estride + s_slot[k]) * D_MODEL); + for (int i = threadIdx.x; i < CH; i += THREADS) dst[i] = src[i]; + } + } + __syncthreads(); + } +} +__device__ __forceinline__ unsigned moe_row_mask(const float *routes, int row) { + unsigned m = 0u; + const float *rw = routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) if (rw[e] != 0.f) m |= 1u << e; + return m; +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx, + unsigned mask) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + for (unsigned m = mask; m; m &= m - 1u) + value = __fadd_rn(value, parts[base + (size_t)(__ffs(m) - 1) * D_MODEL]); + return __float2bfloat16( + __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL])); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + uint64_t *keys, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int n = threadIdx.x; n < SORT_N; n += THREADS) { + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + keys[n] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + __syncthreads(); + for (int k = 2; k <= SORT_N; k <<= 1) { + if (k >= 64) __syncthreads(); + for (int j = k >> 1; j > 0; j >>= 1) { + for (int i = threadIdx.x; i < SORT_N; i += THREADS) { + const int partner = i ^ j; + if (partner > i) { + const bool descending = ((i & k) == 0); + if ((keys[i] < keys[partner]) == descending) { + const uint64_t t = keys[i]; + keys[i] = keys[partner]; + keys[partner] = t; + } + } + } + if (j >= 32) __syncthreads(); + else __syncwarp(); + } + } + __syncthreads(); + if (threadIdx.x == 0) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float v; + int n; + if (!eos_used && eos_v > av) { + v = eos_v; + n = AUDIO_EOS; + eos_used = true; + } else { + v = av; + n = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = + v + gumbel[((size_t)row * N_POS + pos) * TOP_K + j]; + if (noised > best_noised) { best_noised = noised; best_code = n; } + } + out_codes[row] = best_code; + } + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void prefill_expert_masks(const float *routes, unsigned int *masks, int rows) { + const int lane=threadIdx.x&31, warp=threadIdx.x>>5; + const int nt=(rows+BATCH_PAD-1)/BATCH_PAD; + for (int t=blockIdx.x*WARPS+warp;t(gath + g * D_MODEL + c * 8); + if (local < ecnt[e]) + *dst = *reinterpret_cast( + normed + (size_t)gidx[g] * D_MODEL + c * 8); + else + *dst = make_float4(0.f, 0.f, 0.f, 0.f); + } +} +__device__ __forceinline__ int gemm_eff(int tiles, int reps) { + if (NCONS <= 1) return 1; + const int g = (int)gridDim.x; + const int rounds1 = (tiles * reps + g - 1) / g; + if constexpr (NCONS == 3) { + const int max_eff = reps <= BB_EXPERTS ? 2 : 3; + const int rounds2 = (((tiles + 1) / 2) * reps + g - 1) / g; + if (max_eff == 2) return rounds2 < rounds1 ? 2 : 1; + const int rounds3 = (((tiles + 2) / 3) * reps + g - 1) / g; + if (rounds3 < rounds2 && rounds3 < rounds1) return 3; + return rounds2 < rounds1 ? 2 : 1; + } else { + const int roundsN = (((tiles + NCONS - 1) / NCONS) * reps + g - 1) / g; + return roundsN < rounds1 ? NCONS : 1; + } +} +__device__ __forceinline__ int gemm_groups(int tiles, int reps) { + const int eff = gemm_eff(tiles, reps); + return (tiles + eff - 1) / eff; +} +__device__ __forceinline__ int gemm_consumer_id() { + const int g = warpgroup::groupid(); + if (g == 1) return -1; + const int c = g == 0 ? 0 : g - 1; + return c < NCONS ? c : -1; +} +__device__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][BATCH_PAD]; + __shared__ int moe_gidx_stage[NCONS][BATCH_PAD]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + a_s[s], a, {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + a_s[s], a, {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / BATCH_PAD) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, BATCH_PAD> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], a_s[s]); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * BATCH_PAD; + if (epi == EPI_BIAS) { + warpgroup::store(o_sf, accum); + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + dest[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(o_sf[int2{col, row}] + + __bfloat162float(bias[base + col])); + } + warpgroup::sync(1 + cons); + continue; + } + warpgroup::store(o_s[cons], accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_s[cons][int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__device__ void gemm_narrow( + const weight_gl &w, const act_gl_n &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = PAD_N, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][PAD_N]; + __shared__ int moe_gidx_stage[NCONS][PAD_N]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + PAD_N - 1) / PAD_N; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + PAD_N - 1) / PAD_N; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / PAD_N) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, PAD_N> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], + *reinterpret_cast(&a_s[s])); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * PAD_N; + out_tile_n &o_sn = *reinterpret_cast(&o_s[cons]); + warpgroup::store(o_sn, accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_sn[int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&w_s)[STAGES * NCONS] = + alloc.allocate(); + act_tile (&a_s)[STAGES] = alloc.allocate(); + out_tile (&o_s)[NCONS] = alloc.allocate(); +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 + static_assert(NCONS * sizeof(out_tile) >= sizeof(out_tile_f), + "aliased head staging must fit the output tiles"); + out_tile_f &o_sf = *reinterpret_cast(&o_s[0]); +#else + out_tile_f (&o_sf) = alloc.allocate(); +#endif + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_WARPS + 2]; +#if MK2_PREFILL_ONLY + bf16 *const route_stage = reinterpret_cast(&o_s[0]); + static_assert(sizeof(o_s) >= D_MODEL * sizeof(bf16), + "router row stage must fit the output tiles"); +#else + __shared__ __align__(16) uint64_t sort_keys[SORT_N]; + bf16 *const route_stage = reinterpret_cast(&sort_keys[0]); + static_assert(sizeof(sort_keys) >= D_MODEL * sizeof(bf16), + "router row stage must fit the sampler workspace"); +#endif + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS * NCONS, 0); + for (int c = 0; c < CONSUMER_WARPS * NCONS; ++c) arrive(drained[s]); + } + } + int gemm_iter = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map); + grid_barrier(g.barrier); + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map); + grid_barrier(g.barrier); + mamba_gated_norm(L, bb->zxbcdt, bb->gated); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, nullptr, 1, BATCH, + bb->moe_w, BB_MOE_INTER / OUT_ATOM); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + 1, BATCH, bb->moe_w, D_MODEL / OUT_ATOM); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_deterministic(bb->moe_out, idx, + (1u << BB_EXPERTS) - 1u); + grid_barrier(g.barrier); + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int original_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + bool identical_cohort = false; + if constexpr (BATCH > 1) { + if (bb->native_pf_cohort_check && bb->native_prefill_reset && + bb->pf_n_seq == BATCH && !bb->pf_fuse_decode && + !bb->pf_layer_out && !bb->moe_w_inject && !bb->moe_w_audit) { + if (blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[0] = 0; + grid_barrier(g.barrier); + const auto *bits = reinterpret_cast(bb->prompt); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < original_total * D_MODEL; idx += gridDim.x * THREADS) { + const int canonical = (idx / (BATCH * D_MODEL)) * BATCH * D_MODEL + idx % D_MODEL; + if (bits[idx] != bits[canonical]) atomicOr(bb->native_pf_cohort_check, 1u); + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + if (bb->pf_seq_slot[row] != row || bb->slot_map[row] != row) + atomicOr(bb->native_pf_cohort_check, 1u); + grid_barrier(g.barrier); + identical_cohort = bb->native_pf_cohort_check[0] == 0; + } + } + const int n_total = identical_cohort ? original_total / BATCH : original_total; + if (bb->native_pf_cohort_check && blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[1] = identical_cohort ? 1 : BATCH; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int reset_slots = identical_cohort ? 1 : BATCH; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * reset_slots * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % reset_slots; + const MambaLayer &L = bb->mamba[row_layer / reset_slots]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * reset_slots * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % reset_slots; + const MambaLayer &L = bb->mamba[row_layer / reset_slots]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = original_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + const int src = identical_cohort ? (idx / D_MODEL) * BATCH * D_MODEL + idx % D_MODEL : idx; + bb->pf_h[idx] = bb->prompt[src]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = identical_cohort ? 1 : (bb->pf_n_seq > 0 ? bb->pf_n_seq : 1); + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T, + (!bb->pf_layer_out && layer > 0 && + bb->moe_of_layer[layer - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + if (bb->native_pf_expert_masks) + for (int i = blockIdx.x * THREADS + threadIdx.x; i < n_tok; + i += gridDim.x * THREADS) + bb->native_pf_expert_masks[i] = 0u; + if (bb->pf_moe_ecnt) + for (int e = blockIdx.x * THREADS + threadIdx.x; e < BB_EXPERTS; + e += gridDim.x * THREADS) + bb->pf_moe_ecnt[e] = 0; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + mamba_dtda_prefill(L, bb->pf_zxbcdt, bb->pf_dtda, T); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_dtda, + bb->pf_gated, T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, + T, bb->native_pf_expert_masks, bb->pf_moe_gath, + bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, route_stage); + } + grid_barrier(g.barrier); + int max_ecnt = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + max_ecnt = max(max_ecnt, bb->pf_moe_ecnt[e]); + const bool narrow = BATCH_PAD > PAD_N && max_ecnt <= PAD_N; + const int etile = narrow ? PAD_N : BATCH_PAD; + int nt_sum = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_sum += (bb->pf_moe_ecnt[e] + etile - 1) / etile; + const int up_steps = + gemm_groups(BB_MOE_INTER / OUT_ATOM, nt_sum) * + nt_sum * (D_MODEL / KCHUNK); + const int shup_steps = + gemm_groups(BB_SHARED_INTER / OUT_ATOM, n_tok) * + n_tok * (D_MODEL / KCHUNK); + const int down_steps = + gemm_groups(D_MODEL / OUT_ATOM, nt_sum) * + nt_sum * (BB_MOE_INTER / KCHUNK); + const int shdown_steps = + gemm_groups(D_MODEL / OUT_ATOM, n_tok) * + n_tok * (BB_SHARED_INTER / KCHUNK); + if (narrow) + gemm_narrow(L.exp_up, bb->pf_moe_gath_vn, + BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, PAD_N, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + else + gemm(L.exp_up, bb->pf_moe_gath_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, BATCH_PAD, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + up_steps, up_steps + shup_steps); + grid_barrier(g.barrier); + if (narrow) + gemm_narrow(L.exp_down, bb->pf_moe_hc_vn, + BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, PAD_N, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + else + gemm(L.exp_down, bb->pf_moe_hc_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + down_steps, down_steps + shdown_steps); + grid_barrier(g.barrier); + if (bb->pf_layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * D_MODEL; idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = + bb_moe_fold_deterministic(bb->pf_moe_out, idx, + moe_row_mask(bb->pf_moe_w, idx / D_MODEL)); + grid_barrier(g.barrier); + } + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, + identical_cohort ? bb->pf_normed : bb->pf_final, ln_scratch, + n_total, + (!bb->pf_layer_out && n_backbone_layers > 0 && + bb->moe_of_layer[n_backbone_layers - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + grid_barrier(g.barrier); + if (identical_cohort) { + const int worker = blockIdx.x * THREADS + threadIdx.x; + const int workers = gridDim.x * THREADS; + for (int i = worker; i < original_total * D_MODEL; i += workers) + bb->pf_final[i] = bb->pf_normed[(i / (BATCH * D_MODEL)) * D_MODEL + i % D_MODEL]; + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = worker; i < 14 * (BATCH - 1) * CV; i += workers) { + const int cell = i % CV, pair = i / CV; + const int row = pair % (BATCH - 1) + 1; + const MambaLayer &L = bb->mamba[pair / (BATCH - 1)]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = L.conv_state[cell]; + } + for (int i = worker; i < 14 * SS; i += workers) { + const int cell = i % SS; + const MambaLayer &L = bb->mamba[i / SS]; + L.r8_canonical_ssm[cell] = L.ssm_state[cell]; + } + const int kv_columns = BB_KV_HEADS * BB_HEAD_DIM; + for (int i = worker; i < 3 * 2 * (BATCH - 1) * n_total * kv_columns; i += workers) { + int q = i / kv_columns; + const int col = i % kv_columns; + const int pos = q % n_total; q /= n_total; + const int row = q % (BATCH - 1) + 1; q /= BATCH - 1; + const int kind = q % 2, layer = q / 2; + const AttnLayer &L = bb->attn[layer]; + const int offset = kv_offset(L, pos, col / BB_HEAD_DIM) + col % BB_HEAD_DIM; + kv_slot(L, row, kind, pos)[offset] = kv_slot(L, 0, kind, pos)[offset]; + } + grid_barrier(g.barrier); + } + if (!identical_cohort) { + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (size_t i = blockIdx.x * THREADS + threadIdx.x; i < size_t(14) * BATCH * SS; i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.b32_private[(size_t)row * L.ssm_slot_stride + cell] = __float2bfloat16(L.ssm_state[(size_t)row * L.ssm_slot_stride + cell]); + } + grid_barrier(g.barrier); + } + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } +#if MK2_PREFILL_ONLY + return; +#else + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, bb->phoneme_b); + grid_barrier(g.barrier); + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + if (threadIdx.x == 0) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = 0; n < BB_PHONEME; ++n) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + bb->phoneme_out[row] = arg; + } + __syncthreads(); + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + for (int i = lane; i < D_HEAD; i += 32) { + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + for (int t = 0; t <= pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + for (int i = lane; i < D_HEAD; i += 32) { + float acc = 0.f; + for (int t = 0; t <= pos; ++t) + acc += sc[t] * __bfloat162float( + g.vcache[cache_base + t * D_HEAD + i]); + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + i] = + __float2bfloat16(acc); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, g.sampled_pos, + sort_keys, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { +#ifdef ABLATE_SAMPLE + __syncthreads(); +#endif + } else { + grid_barrier(g.barrier); + } + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x < BATCH && blockIdx.x == 0) { + const int code = g.sampled_pos[threadIdx.x]; + g.sampled[threadIdx.x * N_POS + pos] = code; + if (g.native_codes) + g.native_codes[threadIdx.x * N_POS + pos] = (long long)code; + } + grid_barrier(g.barrier); + } + if (g.native_seq_len && blockIdx.x == 0 && threadIdx.x < BATCH) + g.native_seq_len[threadIdx.x] += 1; + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +#endif +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 2 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK), + 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = MK2_PREFILL_ONLY + ? out->sms * out->blocks_per_sm + : std::min(out->sms * out->blocks_per_sm, D_FFN / OUT_ATOM); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + 0, + nullptr, + nullptr, + nullptr, nullptr, nullptr, nullptr, nullptr, 0, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "dedup_native_prefill_impl.inc" diff --git a/deploy/nemotron_native/kernels/b32/native/prefill/mk2_abi.cuh b/deploy/nemotron_native/kernels/b32/native/prefill/mk2_abi.cuh new file mode 100644 index 0000000..95a253b --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/prefill/mk2_abi.cuh @@ -0,0 +1,82 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} diff --git a/deploy/nemotron_native/kernels/b32/native/prefill/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/b32/native/prefill/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/prefill/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/b32/native/prefill/prefill.cu b/deploy/nemotron_native/kernels/b32/native/prefill/prefill.cu new file mode 100644 index 0000000..bd17382 --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/prefill/prefill.cu @@ -0,0 +1,18 @@ +#define MK2_BATCH 32 +#define MK2_BATCH_TILE 64 +#define MK2_PREFILL_ONLY 1 +#define STAGES 3 +#define KCHUNK 128 +#define MK2_GEMM_CONSUMERS 3 +#define MK2_THREADS 512 +#define MK2_NO_MAIN 1 +#define MK2_DIRECT_INPUT 1 +#define MK2_B1_LOCAL_BARRIERS 1 +#define MK2_SKIP_HIDDEN_SNAPSHOT 1 +#define MK2_ACTIVE_EXPERTS 1 +#define MK2_SSM_SPLITS 1 +#define PF_SSM_PW_MULTI 12 +#define PF_ACTIVE_TOKEN_EXPERTS 1 +#include +#include +#include "dedup_predictor_mk2.cuh" diff --git a/deploy/nemotron_native/kernels/b32/native/stage0/mk2_abi.cuh b/deploy/nemotron_native/kernels/b32/native/stage0/mk2_abi.cuh new file mode 100644 index 0000000..ee77c78 --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/stage0/mk2_abi.cuh @@ -0,0 +1,91 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_set_native_predictor_active_mask(MK2Predictor *p, const unsigned char *mask); +int mk2_native_startup_skip_abi_version(); +#if MK2_INFLIGHT +int mk2_native_inflight_abi_version(); +int mk2_set_native_progress_mask(MK2Predictor *p, const unsigned char *mask); +int mk2_inflight_noise(const unsigned long long *seeds, const int *generations, + const unsigned char *mask, float *out, void *stream); +#endif +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *, int grid); diff --git a/deploy/nemotron_native/kernels/b32/native/stage0/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/b32/native/stage0/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/stage0/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/b32/native/stage0/native_prefill_impl.inc b/deploy/nemotron_native/kernels/b32/native/stage0/native_prefill_impl.inc new file mode 100644 index 0000000..a254308 --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/stage0/native_prefill_impl.inc @@ -0,0 +1,198 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 16 || BATCH == 32 || BATCH == 64)) + throw std::runtime_error("native prefill requires matching batch 1/16/32/64"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (padded > INT_MAX / ((BB_EXPERTS + 1) * D_MODEL)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + float *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(float))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(float))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} diff --git a/deploy/nemotron_native/kernels/b32/native/stage0/predictor_mk2.cu b/deploy/nemotron_native/kernels/b32/native/stage0/predictor_mk2.cu new file mode 100644 index 0000000..8d5c1c3 --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/stage0/predictor_mk2.cu @@ -0,0 +1,5472 @@ +#include +#include +#ifndef MK2_INFLIGHT +#define MK2_INFLIGHT 0 +#endif +#if MK2_INFLIGHT +#include +#endif +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +__device__ __forceinline__ bool native_progresses( + const unsigned char *mask, int row) { +#if MK2_INFLIGHT + return !mask || mask[row] != 0; +#else + return true; +#endif +} +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +constexpr int BATCH_PAD = ((BATCH + 15) / 16) * 16; +static_assert(BATCH >= 1 && BATCH <= BATCH_PAD); +static_assert(BATCH % 16 != 0 || BATCH_PAD == BATCH); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 4 +#endif +constexpr int THREADS = 256; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act16_tile = st_bf<16, KCHUNK>; +using out16_tile = st_bf; +using out16_float = st; +using act_gl = gl; +using B32QT = kittens::fp8e4m3; +constexpr int B32_STAGES = 2 * STAGES; +constexpr int B32_KCHUNK = KCHUNK; +static_assert(KCHUNK % 2 == 0); +static_assert(D_MODEL % B32_KCHUNK == 0 && D_FFN % B32_KCHUNK == 0); +using B32QW = st; +using B32QA = st; +using B32QA16 = st; +static_assert(B32_STAGES * (sizeof(B32QW) + sizeof(B32QA)) <= + STAGES * (sizeof(weight_tile) + sizeof(act_tile))); +static_assert(B32_STAGES * sizeof(B32QW) % 1024 == 0); +using B32QGL = gl; +struct B32QWeight { B32QT *data=nullptr; float *scale=nullptr; B32QGL *view=nullptr; float activation_scale=1.f; float next_scale=1.f; }; +struct B32QContext { B32QGL views[2]; }; +__device__ unsigned char b32_qencode(float value) { + return __nv_cvt_float_to_fp8(value, __NV_SATFINITE, __NV_E4M3); +} +__device__ __host__ __forceinline__ size_t b32_tiled_offset(int row, int col, int k) { + const int t = row / OUT_ATOM, lr = row - t * OUT_ATOM; + const int c = col / B32_KCHUNK, lc = col - c * B32_KCHUNK; + return ((size_t)(t * (k / B32_KCHUNK) + c) * OUT_ATOM + lr) * B32_KCHUNK + lc; +} +__global__ void b32_pack_fp8(const bf16 *src, B32QT *dst, float *scale, int n, int k) { + __shared__ float values[256]; + for (int row=blockIdx.x; row>=1) { + if (threadIdx.x(dst)[b32_tiled_offset(row,j,k)]=b32_qencode(__bfloat162float(src[(size_t)row*k+j])/qscale); + __syncthreads(); + } +} +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_ACC_F32 }; +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; + const float *r8_initial_ssm = nullptr; + const unsigned int *r8_initial_ready = nullptr; + const int *r8_seq_len = nullptr; + int r8_initial_positions = 0; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; + B32QWeight b32_ff_in[N_LAYERS], b32_ff_out[N_LAYERS]; + B32QContext *b32_quant=nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +#ifndef MK2_SPARSE_MOE +#define MK2_SPARSE_MOE 0 +#endif +constexpr int SPARSE_N = 16; +constexpr bool BB_SPARSE_MOE_OWNED = MK2_SPARSE_MOE && BATCH >= 16; +using sparse_act_tile = st_bf; +using sparse_out_tile = st_bf; +using sparse_act_gl = gl; +struct SparseMoe { + bf16 *gathered; + int *rows; + int *counts; + sparse_act_gl gathered_v; + sparse_act_gl hidden_v; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + const SparseMoe *sparse_moe = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +#define BARRIER_BACKOFF_SPREAD \ + ((unsigned)(BARRIER_BACKOFF_NS - 15) + ((unsigned)(blockIdx.x & 15) << 1)) +#define ROW_BARRIER_LEAD_NS 384u +#define BARRIER_GEN_WORD 32 +__device__ __forceinline__ unsigned int barrier_poll(const unsigned int *p) { + unsigned int v; + asm volatile("ld.relaxed.gpu.global.u32 %0, [%1];" + : "=r"(v) : "l"(p) : "memory"); + return v; +} +template +__device__ __forceinline__ void grid_barrier(unsigned int *state, + unsigned int &next_target) { + __syncthreads(); + if (threadIdx.x == 0) { + if constexpr (ACQUIRE) + asm volatile("fence.acq_rel.gpu;" ::: "memory"); + if constexpr (ACQUIRE) + asm volatile("red.relaxed.gpu.global.add.u32 [%0], %1;" + :: "l"(state), "r"(1u) : "memory"); + else + asm volatile("red.release.gpu.global.add.u32 [%0], %1;" + :: "l"(state), "r"(1u) : "memory"); + const unsigned int target = next_target + gridDim.x; + next_target = target; + { + while (barrier_poll(state) < target) + __nanosleep(BARRIER_BACKOFF_SPREAD); + } + } + __syncthreads(); +} +template +__device__ __forceinline__ void grid_barrier_pf(unsigned int *state, + unsigned int &next_target, + PF prefetch) { + __syncthreads(); + if (threadIdx.x == 0) { + if constexpr (ACQUIRE) + asm volatile("fence.acq_rel.gpu;" ::: "memory"); + if constexpr (ACQUIRE) + asm volatile("red.relaxed.gpu.global.add.u32 [%0], %1;" + :: "l"(state), "r"(1u) : "memory"); + else + asm volatile("red.release.gpu.global.add.u32 [%0], %1;" + :: "l"(state), "r"(1u) : "memory"); + const unsigned int target = next_target + gridDim.x; + next_target = target; + prefetch(); + while (barrier_poll(state) < target) + __nanosleep(BARRIER_BACKOFF_SPREAD); + } + __syncthreads(); +} +template +__device__ __forceinline__ void row_grid_barrier_pf(unsigned int *state, + unsigned int &next_target, + PF prefetch) { + __syncthreads(); + if (threadIdx.x == 0) { + if constexpr (ACQUIRE) + asm volatile("fence.acq_rel.gpu;" ::: "memory"); + unsigned int *counter = state; + if (blockIdx.x < BATCH) { + if constexpr (ACQUIRE) + asm volatile("red.relaxed.gpu.global.add.u32 [%0], %1;" + :: "l"(counter), "r"(1u) : "memory"); + else + asm volatile("red.release.gpu.global.add.u32 [%0], %1;" + :: "l"(counter), "r"(1u) : "memory"); + } + const unsigned int target = next_target + BATCH; + next_target = target; + prefetch(); + if (blockIdx.x >= BATCH) __nanosleep(ROW_BARRIER_LEAD_NS); + while (barrier_poll(counter) < target) + __nanosleep(BARRIER_BACKOFF_SPREAD); + } + __syncthreads(); +} +template +__device__ __forceinline__ void row_grid_barrier(unsigned int *state, + unsigned int &next_target) { + __syncthreads(); + if (threadIdx.x == 0) { + if constexpr (ACQUIRE) + asm volatile("fence.acq_rel.gpu;" ::: "memory"); + unsigned int *counter = state; + if (blockIdx.x < BATCH) { + if constexpr (ACQUIRE) + asm volatile("red.relaxed.gpu.global.add.u32 [%0], %1;" + :: "l"(counter), "r"(1u) : "memory"); + else + asm volatile("red.release.gpu.global.add.u32 [%0], %1;" + :: "l"(counter), "r"(1u) : "memory"); + } + const unsigned int target = next_target + BATCH; + next_target = target; + if (blockIdx.x >= BATCH) __nanosleep(ROW_BARRIER_LEAD_NS); + while (barrier_poll(counter) < target) + __nanosleep(BARRIER_BACKOFF_SPREAD); + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +template +__device__ void layernorm_rows( + bf16 *src, const float *weight, bf16 *dst, float *scratch, + B32QContext *qctx = nullptr, float qscale = 1.f, + const bf16 *position = nullptr) { + constexpr int PER_THREAD = D_MODEL / THREADS; + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const bf16 *in = src + (size_t)row * D_MODEL; + float v[PER_THREAD], wv[PER_THREAD]; + float sum = 0.f, sq = 0.f; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + if constexpr (ADD_POSITION) { + const bf16 added = __float2bfloat16( + __bfloat162float(in[i]) + __bfloat162float(position[i])); + src[(size_t)row * D_MODEL + i] = added; + v[j] = __bfloat162float(added); + } else { + v[j] = __bfloat162float(in[i]); + } + wv[j] = weight[i]; + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) { scratch[warp] = sum; scratch[WARPS + warp] = sq; } + __syncthreads(); + float s = 0.f, q = 0.f; + for (int w = 0; w < WARPS; ++w) { s += scratch[w]; q += scratch[WARPS + w]; } + const float mean = s / D_MODEL; + const float inv = rsqrtf(q / D_MODEL - mean * mean + LN_EPS); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + const bf16 normed = __float2bfloat16((v[j] - mean) * inv * wv[j]); + if (qctx) { + reinterpret_cast(qctx->views[0].raw_ptr) + [(size_t)row * D_MODEL + i] = + b32_qencode(__bfloat162float(normed) / qscale); + } else { + dst[(size_t)row * D_MODEL + i] = normed; + } + } + if (row + gridDim.x < BATCH) __syncthreads(); + } +} +template +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, bf16 *snapshot = nullptr, int snapshot_layer = -1, + float *snapshot_fp32 = nullptr, float *out_fp32 = nullptr, + const bf16 *reset_embedding = nullptr) { + constexpr int PER_THREAD = D_MODEL / THREADS; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const size_t base = (size_t)row * D_MODEL; + float v[PER_THREAD]; + bf16 wv[PER_THREAD]; + float sq = 0.f; + bf16 hv[PER_THREAD], rv[PER_THREAD]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + wv[j] = weight[i]; + hv[j] = EMBED_RESET ? reset_embedding[base + i] : h[base + i]; + if constexpr (!EMBED_RESET) rv[j] = r[base + i]; + } + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + if (snapshot) + snapshot[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = hv[j]; + if (snapshot_fp32) + snapshot_fp32[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = + __bfloat162float(hv[j]); + v[j] = __bfloat162float(hv[j]); + if constexpr (!EMBED_RESET) + v[j] += __bfloat162float(rv[j]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) scratch[warp] = sq; + __syncthreads(); + float s = 0.f; + for (int w = 0; w < WARPS; ++w) s += scratch[w]; + const float inv = rsqrtf(s / D_MODEL + RMS_EPS); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + const bf16 normed = __float2bfloat16( + __bfloat162float(wv[j]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + out[base + i] = normed; + if (out_fp32) out_fp32[base + i] = __bfloat162float(normed); + } + if (row + gridDim.x < n_rows) __syncthreads(); + } +} +__device__ __forceinline__ void mamba_conv_unit( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map, + int row, int group, const unsigned char *progress_mask = nullptr) { + if (!native_progresses(progress_mask, row)) return; + constexpr int GX = MB_INNER / MB_GROUPS; + constexpr int GUNIT = GX + 2 * MB_STATE; + const int base_b = MB_INNER + group * MB_STATE; + const int base_c = MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + for (int j = threadIdx.x; j < GUNIT; j += THREADS) { + const int c = j < GX ? group * GX + j + : j < GX + MB_STATE ? base_b + (j - GX) + : base_c + (j - GX - MB_STATE); + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + const float raw = __bfloat162float( + zxbcdt[(size_t)row * MB_IN_PROJ + MB_INNER + c]); + bf16 held[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) held[k] = state[k]; + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(held[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + xbc_act[(size_t)row * MB_CONV_DIM + c] = __float2bfloat16(acc); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = held[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + } +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map, + const unsigned char *progress_mask = nullptr) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + if (!native_progresses(progress_mask, row)) continue; + const int c = idx - row * MB_CONV_DIM; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + const float raw = __bfloat162float( + zxbcdt[(size_t)row * MB_IN_PROJ + MB_INNER + c]); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + xbc_act[(size_t)row * MB_CONV_DIM + c] = __float2bfloat16(acc); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ __forceinline__ void mamba_ssm_unit( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map, int row, int head, + const unsigned char *progress_mask = nullptr) { + if (!native_progresses(progress_mask, row)) return; + const int lane = threadIdx.x & 31; + constexpr int channels = MB_HEAD_DIM; + const int first_channel = 0; + const int group = head / MB_HEADS_PER_GROUP; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + L.dt_bias[head]); + const float dA = __expf(L.A[head] * dt); + const float Dh = L.D[head]; + const size_t ssm_at = (size_t)slot_map[row] * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE; + bf16 *st = reinterpret_cast(L.ssm_state) + ssm_at; + const bool r8_initial = L.r8_initial_ssm && L.r8_initial_ready && *L.r8_initial_ready == 1 && L.r8_seq_len[row] == L.r8_initial_positions; + const float *st_in = r8_initial ? L.r8_initial_ssm + (size_t)head * MB_HEAD_DIM * MB_STATE : nullptr; + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + constexpr int NSTATE = MB_STATE / 32; + constexpr int SSM_PREFETCH = channels < 8 ? channels : 8; + static_assert(channels % SSM_PREFETCH == 0); + for (int g0 = 0; g0 < channels; g0 += SSM_PREFETCH) { + float sv[SSM_PREFETCH][NSTATE]; + float xvr[SSM_PREFETCH]; + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const size_t prow = (size_t)(first_channel + g0 + gi) * MB_STATE; + xvr[gi] = __bfloat162float( + xa[head * MB_HEAD_DIM + first_channel + g0 + gi]); + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = r8_initial ? st_in[prow + lane + j * 32] : __bfloat162float(st[prow + lane + j * 32]); + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const int p = first_channel + g0 + gi; + const float xv = xvr[gi]; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < NSTATE; ++j) { + const int n = lane + j * 32; + const size_t at = (size_t)p * MB_STATE + n; + const float s = sv[gi][j] * dA + Bv[j] * xv; + st[at] = __float2bfloat16(s); + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } +} +__device__ __forceinline__ void mamba_gated_norm_unit( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int row, int group) { + const int lane = threadIdx.x & 31; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32], wg[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + wg[j] = __bfloat162float(L.gnorm_w[idx]); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16(held[j] * inv * wg[j]); + } +} +__device__ void mamba_mixer_local( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, bf16 *gated, + const int *slot_map, const unsigned char *progress_mask = nullptr) { + constexpr int UNITS = BATCH * MB_GROUPS; + const int warp = threadIdx.x >> 5; + for (int unit = blockIdx.x; unit < UNITS; unit += gridDim.x) + mamba_conv_unit(L, zxbcdt, xbc_act, slot_map, + unit / MB_GROUPS, unit % MB_GROUPS, progress_mask); + __syncthreads(); + for (int unit = blockIdx.x; unit < UNITS; unit += gridDim.x) + mamba_ssm_unit(L, zxbcdt, xbc_act, gated, slot_map, + unit / MB_GROUPS, + (unit % MB_GROUPS) * MB_HEADS_PER_GROUP + warp, progress_mask); + __syncthreads(); + int slot = 0; + for (int unit = blockIdx.x; unit < UNITS; unit += gridDim.x, ++slot) + if (warp == slot) + mamba_gated_norm_unit(L, zxbcdt, gated, + unit / MB_GROUPS, unit % MB_GROUPS); + __syncthreads(); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map, const unsigned char *progress_mask = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + const int pairs = BATCH * MB_HEADS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_HEADS * split); + if (!native_progresses(progress_mask, row)) continue; + const int head = (pair / split) % MB_HEADS; + const int first_channel = (pair % split) * channels; + const int group = head / MB_HEADS_PER_GROUP; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + L.dt_bias[head]); + const float dA = __expf(L.A[head] * dt); + const float Dh = L.D[head]; + const size_t ssm_at = (size_t)slot_map[row] * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE; + bf16 *st = reinterpret_cast(L.ssm_state) + ssm_at; + const bool r8_initial = L.r8_initial_ssm && L.r8_initial_ready && *L.r8_initial_ready == 1 && L.r8_seq_len[row] == L.r8_initial_positions; + const float *st_in = r8_initial ? L.r8_initial_ssm + (size_t)head * MB_HEAD_DIM * MB_STATE : nullptr; + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + constexpr int NSTATE = MB_STATE / 32; + constexpr int SSM_PREFETCH = channels < 8 ? channels : 8; + static_assert(channels % SSM_PREFETCH == 0); + for (int g0 = 0; g0 < channels; g0 += SSM_PREFETCH) { + float sv[SSM_PREFETCH][NSTATE]; + float xvr[SSM_PREFETCH]; + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const size_t prow = (size_t)(first_channel + g0 + gi) * MB_STATE; + xvr[gi] = __bfloat162float( + xa[head * MB_HEAD_DIM + first_channel + g0 + gi]); + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = r8_initial ? st_in[prow + lane + j * 32] : __bfloat162float(st[prow + lane + j * 32]); + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const int p = first_channel + g0 + gi; + const float xv = xvr[gi]; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < NSTATE; ++j) { + const int n = lane + j * 32; + const size_t at = (size_t)p * MB_STATE + n; + const float s = sv[gi][j] * dA + Bv[j] * xv; + st[at] = __float2bfloat16(s); + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, float2 *ssm_params, int n_seq = 1, + const int *seq_slot = nullptr) { + const int total = n_tokens * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_CONV_DIM; + const int c = idx - t * MB_CONV_DIM; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v; +#ifdef PF_CONV_ZERO_STATE + if (src < 0) v = 0.f; +#else + if (src < 0) v = __bfloat162float(state[pos + k]); +#endif + else v = __bfloat162float( + zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]); + acc += __bfloat162float(w[k]) * v; + } +#ifndef PF_CONV_NO_SILU + acc = acc / (1.f + __expf(-acc)); +#endif + xbc_act[(size_t)t * MB_CONV_DIM + c] = __float2bfloat16(acc); + if (c < MB_HEADS) { + const bf16 *dtrow = + zxbcdt + (size_t)t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[c]) + L.dt_bias[c]); + ssm_params[(size_t)t * MB_HEADS + c] = + make_float2(dt, __expf(L.A[c] * dt)); + } + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + int n_tokens, int slot, const float2 *ssm_params, int n_seq = 1, + const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + const int units = n_seq * MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int seq = unit / (MB_HEADS * MB_HEAD_DIM); + const int rem = unit - seq * (MB_HEADS * MB_HEAD_DIM); + const int head = rem / MB_HEAD_DIM; + const int p = rem - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Dh = L.D[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p * MB_STATE; + float sreg[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) sreg[j] = st[lane + j * 32]; + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const bf16 *xa = xbc_act + t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + const float2 params = ssm_params[t * MB_HEADS + head]; + const float dt = params.x; + const float dA = params.y; + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + sreg[j] = sreg[j] * dA + __bfloat162float(Bsec[n]) * dt * xv; + acc += sreg[j] * __bfloat162float(Csec[n]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) st[lane + j * 32] = sreg[j]; + } + return; + } +#if PF_SSM_SPLIT_P + const int units = MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; unit += warps_total) { + const int head = unit / MB_HEAD_DIM; + const int p_only = unit - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; +#else + for (int head = blockIdx.x * WARPS + warp; head < MB_HEADS; + head += warps_total) { + const int group = head / MB_HEADS_PER_GROUP; +#endif + const float Dh = L.D[head]; + float *st = L.ssm_state + (size_t)slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE +#if PF_SSM_SPLIT_P + + (size_t)p_only * MB_STATE +#endif + ; + for (int t = 0; t < n_tokens; ++t) { + const bf16 *xa = xbc_act + (size_t)t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const float2 params = + ssm_params[(size_t)t * MB_HEADS + head]; + const float dt = params.x; + const float dA = params.y; + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } +#if PF_SSM_SPLIT_P + { + const int p = p_only; +#else + for (int p = 0; p < MB_HEAD_DIM; ++p) { +#endif + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { +#if PF_SSM_SPLIT_P + const size_t at = (size_t)lane + j * 32; +#else + const size_t at = (size_t)p * MB_STATE + lane + j * 32; +#endif + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } + } +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32], wg[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + wg[j] = __bfloat162float(L.gnorm_w[idx]); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16(held[j] * inv * wg[j]); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len, const unsigned char *progress_mask = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + constexpr int Q_PER_WARP = 2; + constexpr int GROUPS_PER_KV = 2; + const int pairs = BATCH * BB_KV_HEADS * GROUPS_PER_KV; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x + warp * gridDim.x; pair < pairs; + pair += gridDim.x * WARPS) { + const int row = pair / (BB_KV_HEADS * GROUPS_PER_KV); + if (!native_progresses(progress_mask, row)) continue; + const int within_row = pair % (BB_KV_HEADS * GROUPS_PER_KV); + const int kv = within_row / GROUPS_PER_KV; + const int q_group = within_row % GROUPS_PER_KV; + const int q_first = q_group * Q_PER_WARP; + const int q_count = q_group == 0 ? Q_PER_WARP : 1; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (q_group == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + constexpr int KV_PER_LANE = BB_HEAD_DIM / 32; + bf16 kk[KV_PER_LANE], vv[KV_PER_LANE]; + #pragma unroll + for (int j = 0; j < KV_PER_LANE; ++j) { + kk[j] = k_new[lane + j * 32]; + vv[j] = v_new[lane + j * 32]; + } + #pragma unroll + for (int j = 0; j < KV_PER_LANE; ++j) { + kc[lane + j * 32] = kk[j]; + vc[lane + j * 32] = vv[j]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[h][j] = 0.f; + } + float m[Q_PER_WARP] = {-CUDART_INF_F, -CUDART_INF_F}; + float denom[Q_PER_WARP] = {0.f, 0.f}; + constexpr int PF = 8; + const int tpb = L.tokens_per_block; + const size_t elems_per_block = + (size_t)tpb * BB_KV_HEADS * BB_HEAD_DIM; + const int *bt_k = L.block_table + (size_t)(row * 2) * L.max_blocks_per_seq; + const int *bt_v = bt_k + L.max_blocks_per_seq; + const int head_base = kv * tpb * BB_HEAD_DIM; + float qr[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const bf16 *q = base + (kv * kv_per_q + q_first + h) * BB_HEAD_DIM; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qr[h][j] = __bfloat162float(q[lane + j * 32]); + } + for (int page0 = 0; page0 < T; page0 += tpb) { + const int k_page = bt_k[page0 / tpb] & 0x7FFFFFFF; + const int v_page = bt_v[page0 / tpb] & 0x7FFFFFFF; + const bf16 *kbase = L.pool_base + + (size_t)k_page * elems_per_block + + head_base; + const bf16 *vbase = L.pool_base + + (size_t)v_page * elems_per_block + + head_base; + const int tend = page0 + tpb < T ? page0 + tpb : T; + for (int t0 = page0; t0 < tend; t0 += PF) { + const int navail = tend - t0 < PF ? tend - t0 : PF; + float kr[PF][PER_LANE], vr[PF][PER_LANE]; + #pragma unroll + for (int u = 0; u < PF; ++u) { + const int tpos = t0 + (u < navail ? u : navail - 1); +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = vbase + (tpos - page0) * BB_HEAD_DIM; +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new : kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = cur_token + ? v_new : vbase + (tpos - page0) * BB_HEAD_DIM; +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + kr[u][j] = __bfloat162float(kt[i]); + vr[u][j] = __bfloat162float(vt[i]); + } + } + #pragma unroll + for (int u = 0; u < PF; ++u) { + if (u >= navail) break; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + dot += qr[h][j] * kr[u][j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m[h], s); + const float rescale = __expf(m[h] - m_new); + const float w = __expf(s - m_new); + denom[h] = denom[h] * rescale + w; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[h][j] = acc[h][j] * rescale + w * vr[u][j]; + m[h] = m_new; + } + } + } + } + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const float inv = 1.f / denom[h]; + const int head = kv * kv_per_q + q_first + h; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[h][j] * inv); + } + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * BB_KV_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / BB_KV_HEADS; + const int kv = idx - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv); + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv); + for (int i = 0; i < BB_HEAD_DIM; ++i) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + for (int s = 0; s <= pos; ++s) { + const bf16 *kt = kv_slot(L, seq, 0, s) + kv_offset(L, s, kv); + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + const bf16 *vt = kv_slot(L, seq, 1, s) + kv_offset(L, s, kv); + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + wt * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH) { + __shared__ float score[BB_EXPERTS]; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + float acc0 = 0.f, acc1 = 0.f, acc2 = 0.f; + const bf16 *w0 = gate_w + (size_t)warp * D_MODEL; + const bf16 *w1 = gate_w + (size_t)(warp + WARPS) * D_MODEL; + const bf16 *w2 = gate_w + (size_t)(warp + 2 * WARPS) * D_MODEL; + #pragma unroll 8 + for (int i = lane; i < D_MODEL; i += 32) { + const float x = __bfloat162float(normed[(size_t)row * D_MODEL + i]); + acc0 += x * __bfloat162float(w0[i]); + acc1 += x * __bfloat162float(w1[i]); + acc2 += x * __bfloat162float(w2[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + acc0 += __shfl_xor_sync(0xffffffffu, acc0, off); + acc1 += __shfl_xor_sync(0xffffffffu, acc1, off); + acc2 += __shfl_xor_sync(0xffffffffu, acc2, off); + } + if (lane == 0) { + score[warp] = 1.f / (1.f + __expf(-acc0)); + score[warp + WARPS] = 1.f / (1.f + __expf(-acc1)); + score[warp + 2 * WARPS] = 1.f / (1.f + __expf(-acc2)); + } + __syncthreads(); + if (threadIdx.x < 32) { + const int lane = threadIdx.x; + const bool valid = lane < BB_EXPERTS; + if (valid) moe_w[row * BB_EXPERTS + lane] = 0.f; + float rem = valid ? score[lane] + __bfloat162float(gate_bias[lane]) + : -CUDART_INF_F; + int best_of[BB_TOP_K]; + float total = 0.f; + #pragma unroll + for (int k = 0; k < BB_TOP_K; ++k) { + float v = rem; + int i = lane; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffffu, v, off); + const int oi = __shfl_xor_sync(0xffffffffu, i, off); + if (ov > v || (ov == v && oi < i)) { v = ov; i = oi; } + } + best_of[k] = i; + total += score[i]; + if (lane == i) rem = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + __syncwarp(); + if (lane == 0) { + #pragma unroll + for (int k = 0; k < BB_TOP_K; ++k) + moe_w[row * BB_EXPERTS + best_of[k]] = + score[best_of[k]] * norm; + } + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert <= BB_EXPERTS; ++expert) + value = __fadd_rn(value, parts[base + (size_t)expert * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ bf16 bb_moe_fold_selected( + const float *parts, const float *weights, int idx) { + const int row = idx / D_MODEL; + const size_t base = (size_t)row * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert < BB_EXPERTS; ++expert) { + const float part = weights[row * BB_EXPERTS + expert] != 0.f + ? parts[base + (size_t)expert * D_MODEL] : 0.f; + value = __fadd_rn(value, part); + } + value = __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ bf16 bb_moe_fold_masked( + const float *parts, int row, int col, unsigned selected) { + const size_t base = (size_t)row * (BB_EXPERTS + 1) * D_MODEL + col; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert < BB_EXPERTS; ++expert) { + const float part = (selected & (1u << expert)) + ? parts[base + (size_t)expert * D_MODEL] : 0.f; + value = __fadd_rn(value, part); + } + value = __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ __forceinline__ void sample_argmax_merge( + float &value, int &rank, float other, int other_rank) { + if (other > value || (other == value && other_rank < rank)) { + value = other; + rank = other_rank; + } +} +__device__ __forceinline__ int sample_block_argmax(float value, int rank) { + __shared__ float warp_value[WARPS]; + __shared__ int warp_rank[WARPS]; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + if (lane == 0) { warp_value[warp] = value; warp_rank[warp] = rank; } + __syncthreads(); + if (warp == 0) { + value = lane < WARPS ? warp_value[lane] : -CUDART_INF_F; + rank = lane < WARPS ? warp_rank[lane] : INT_MAX; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + } + return rank; +} +__device__ __forceinline__ int sample_eos_insertion_rank( + const uint64_t *keys, float eos) { + int lo = 0, hi = TOP_K; + while (lo < hi) { + const int mid = (lo + hi) >> 1; + if (eos > ordered_to_float((uint32_t)(keys[mid] >> 32))) hi = mid; + else lo = mid + 1; + } + return lo; +} +__device__ __forceinline__ void sample_topk_serial( + const uint64_t *keys, float eos, const float *noise, int *out) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float value; + int code; + if (!eos_used && eos > av) { + value = eos; + code = AUDIO_EOS; + eos_used = true; + } else { + value = av; + code = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = value + noise[j]; + if (noised > best_noised) { + best_noised = noised; + best_code = code; + } + } + *out = best_code; +} +__device__ __forceinline__ void sample_topk_parallel( + const uint64_t *keys, float eos, const float *noise, int *out, + float noise_j) { + const int j = threadIdx.x; + const bool sorted_nan = j < TOP_K && + isnan(ordered_to_float((uint32_t)(keys[j] >> 32))); + if (__syncthreads_or(sorted_nan)) { + if (threadIdx.x == 0) sample_topk_serial(keys, eos, noise, out); + return; + } + const int lane = threadIdx.x & 31; + int insertion = lane == 0 ? sample_eos_insertion_rank(keys, eos) : 0; + insertion = __shfl_sync(0xffffffffu, insertion, 0); + float value = -CUDART_INF_F; + int rank = INT_MAX; + if (j < TOP_K) { + const int a = j - (j > insertion); + const float v = j == insertion + ? eos : ordered_to_float((uint32_t)(keys[a] >> 32)); + const float noised = v + noise_j; + if (noised > -CUDART_INF_F) { value = noised; rank = j; } + } + const int winner = sample_block_argmax(value, rank); + if (threadIdx.x == 0) { + if (winner == INT_MAX) *out = 0; + else if (winner == insertion) *out = AUDIO_EOS; + else { + const int a = winner - (winner > insertion); + *out = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + } + } +} +constexpr int SORT_ITEMS_PER_THREAD = SORT_N / THREADS; +struct sampler_storage { + uint64_t keys[2][SORT_N]; +}; +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + sampler_storage &storage, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const float noise_j = threadIdx.x < TOP_K + ? gumbel[((size_t)row * N_POS + pos) * TOP_K + threadIdx.x] : 0.f; + uint64_t thread_keys[SORT_ITEMS_PER_THREAD]; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + thread_keys[item] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + static_assert(SORT_ITEMS_PER_THREAD == 4, + "the intra-thread partner permutation below is four wide"); + constexpr int GROUP = SORT_N / WARPS; + static_assert(GROUP == 32 * SORT_ITEMS_PER_THREAD); + static_assert(TOP_K <= GROUP, "top-k must fit one merged group"); + const int sort_warp = threadIdx.x >> 5, sort_lane = threadIdx.x & 31; + auto stage = [&](int jj, int kk) { + uint64_t other[SORT_ITEMS_PER_THREAD]; + if (jj >= 32) { + const uint64_t k0 = thread_keys[0], k1 = thread_keys[1], + k2 = thread_keys[2], k3 = thread_keys[3]; + if (jj == 32) { other[0] = k1; other[1] = k0; + other[2] = k3; other[3] = k2; } + else { other[0] = k2; other[1] = k3; + other[2] = k0; other[3] = k1; } + } else { + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) + other[item] = + __shfl_xor_sync(0xffffffffu, thread_keys[item], jj); + } + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int m = sort_lane + item * 32; + const bool take_hi = kk ? (((m & kk) == 0) == ((m & jj) == 0)) + : ((m & jj) == 0); + const uint64_t mine = thread_keys[item]; + thread_keys[item] = + take_hi ? (mine > other[item] ? mine : other[item]) + : (mine < other[item] ? mine : other[item]); + } + }; + for (int kk = 2; kk <= GROUP; kk <<= 1) { + for (int jj = kk >> 1; jj > 0; jj >>= 1) { stage(jj, kk); } + } + int bank = 0; + for (int step = 1; step < WARPS; step <<= 1, bank ^= 1) { + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) + storage.keys[bank][sort_warp * GROUP + sort_lane + item * 32] = + thread_keys[item]; + __syncthreads(); + if ((sort_warp & ((step << 1) - 1)) == 0) { + const int partner = sort_warp ^ step; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const uint64_t other = storage.keys[bank][ + partner * GROUP + (GROUP - 1 - (sort_lane + item * 32))]; + const uint64_t mine = thread_keys[item]; + thread_keys[item] = mine > other ? mine : other; + } + for (int jj = GROUP >> 1; jj > 0; jj >>= 1) { stage(jj, 0); } + } + } + const int top_bank = bank; + if (sort_warp == 0) { + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int rank = sort_lane + item * 32; + if (rank < TOP_K) storage.keys[top_bank][rank] = thread_keys[item]; + } + } + __syncthreads(); + sample_topk_parallel( + storage.keys[top_bank], eos_v, + gumbel + ((size_t)row * N_POS + pos) * TOP_K, + out_codes + row, noise_j); + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ __forceinline__ int advance_tile_cursor(int &tile_io, int n_tiles) { + int tile_start = (int)blockIdx.x - tile_io; + if (tile_start < 0) tile_start += (int)gridDim.x; + tile_io += n_tiles; + while (tile_io >= (int)gridDim.x) tile_io -= (int)gridDim.x; + return tile_start; +} +template +__device__ __forceinline__ void stage_next(ST (&slots)[NS], const GL &w, + int n_out, int iters, int tile_io, + int iter_io, unsigned round_bytes, + semaphore *filled, + semaphore *drained) { + if (threadIdx.x != 0) return; + const int n_tiles = SPLIT * (n_out / OUT_ATOM); + int t = (int)blockIdx.x - tile_io; + if (t < 0) t += (int)gridDim.x; + if (t >= n_tiles) return; + const int tile_id = SPLIT == 1 ? t : t / SPLIT; + const int n = PF < iters ? PF : iters; + #pragma unroll + for (int k = 0; k < PF; ++k) { + if (k >= n) break; + const int iter = iter_io + k; + const int s = iter % NS; + wait(drained[s], (iter / NS) & 1); + tma::expect_bytes(filled[s], round_bytes); + if constexpr (KITERS) + tma::load_async<2, cache_policy::EVICT_FIRST>( + slots[s], w, {tile_id * KITERS + k, 0}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + slots[s], w, {tile_id, k}, filled[s]); + } +} +__device__ __forceinline__ int staged_rounds(int prestage, int tile_start, + int n_tiles, int iters) { + if (!prestage || tile_start >= n_tiles) return 0; + return prestage < iters ? prestage : iters; +} +__device__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES], act_tile (&a_s)[STAGES], out_tile &o_s, + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + int &tile_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0, + int prestage = 0) { + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + const int iters = k_total / KCHUNK; + const int tile_start = advance_tile_cursor(tile_io, n_tiles); + int tiles_here = 0; + for (int t = tile_start; t < n_tiles; t += gridDim.x) ++tiles_here; + const int iter_base = iter_io; + iter_io += tiles_here * iters * n_tok; + const int my_tiles = (n_tiles - tile_start + gridDim.x - 1) / gridDim.x; + const int staged0 = staged_rounds(prestage, tile_start, n_tiles, iters); + if (warpgroup::groupid() == 1) { + if (warpgroup::laneid() == 0) { + int iter = iter_base; + int staged = staged0; + for (int work = tile_start; work < n_tiles; work += gridDim.x) { + const int tile_id = mk2_expert_tile(work, active_mask, compact_tiles); + for (int tok = 0; tok < n_tok; ++tok) + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + if (staged) { --staged; } else { + wait(drained[s], phase); + tma::expect_bytes(filled[s], + sizeof(weight_tile) + sizeof(act_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s], w, {tile_id, i}, filled[s]); } + const int a_col = out_tiles_per_expert + ? (tile_id / out_tiles_per_expert) * act_k_tiles_per_expert + i + : i; + tma::load_async<2, cache_policy::EVICT_LAST>( + a_s[s], a, {tok, a_col}, filled[s]); + } + } + } + } else { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) + for (int tok = 0; tok < n_tok; ++tok) { + const int tile_id = mk2_expert_tile(work, active_mask, compact_tiles); + rt_fl<16, BATCH_PAD> accum; + warp::zero(accum); + int prior_s = -1; + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s], a_s[s]); + if (prior_s >= 0) { + warpgroup::mma_async_wait<1>(); + if (warp::laneid() == 0) arrive(drained[prior_s]); + } + prior_s = s; + } + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[prior_s]); + const int base = tile_id * OUT_ATOM; + const int grow0 = tok * BATCH_PAD; + if (epi == EPI_BIAS) { + const float bv = __bfloat162float( + bias[base + (warpgroup::laneid() & (OUT_ATOM - 1))]); + warpgroup::store(o_sf, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + dest[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(o_sf[int2{col, row}] + bv); + } + warpgroup::sync(1); + continue; + } + warpgroup::store(o_s, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_s[int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_w[grow * BB_EXPERTS + expert]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * + D_MODEL + local] = __fmul_rn(wt, v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1); + } + (void)my_tiles; + } +} +__device__ __forceinline__ void bf16_wgmma_n8( + rt_fl<16, 16> &d, uint64_t a_desc, uint64_t b_desc) { + const int scale_d = 1; + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %6, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k16.f32.bf16.bf16 " + "{%0, %1, %2, %3}, %4, %5, p, 1, 1, 0, 0;\n" + "}\n" + : "+f"(d.tiles[0][0].data[0].x), "+f"(d.tiles[0][0].data[0].y), + "+f"(d.tiles[0][0].data[1].x), "+f"(d.tiles[0][0].data[1].y) + : "l"(a_desc), "l"(b_desc), "r"(scale_d)); +} +__device__ __forceinline__ void bf16_mma_ABt_n8( + rt_fl<16, 16> &d, const weight_tile &a, + const act16_tile &b, int row_half) { + kittens::st_descriptor a_desc(a); + kittens::st_descriptor b_desc(b); + const uint64_t row_offset = row_half + ? kittens::detail::matrix_descriptor_encode(1024) : 0; + constexpr int K = weight_tile::cols / kittens::TILE_COL_DIM; + warpgroup::mma_fence(d); + #pragma unroll + for (int k = 0; k < K; ++k) + bf16_wgmma_n8(d, a_desc.chunk_descriptor(k), + b_desc.chunk_descriptor(k) + row_offset); + warpgroup::mma_commit_group(); +} +template +__device__ void gemm_narrow( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES], act_tile (&a_s)[STAGES], out_tile &o_s, + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + int &tile_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0, + int prestage = 0) { + out16_tile &o16 = *reinterpret_cast(&o_s); + out16_float &of16 = *reinterpret_cast(&o_sf); + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + static_assert(SPLIT == 2 || SPLIT == 4); + constexpr int ROWS = BATCH / SPLIT; + const int n_tiles = SPLIT * (n_out / OUT_ATOM); + const int iters = k_total / KCHUNK; + const int tile_start = advance_tile_cursor(tile_io, n_tiles); + int tiles_here = 0; + for (int t = tile_start; t < n_tiles; t += gridDim.x) ++tiles_here; + const int iter_base = iter_io; + iter_io += tiles_here * iters * n_tok; + const int my_tiles = (n_tiles - tile_start + gridDim.x - 1) / gridDim.x; + const int staged0 = staged_rounds(prestage, tile_start, n_tiles, iters); + if (epi == EPI_RESIDUAL && n_tiles <= (int)gridDim.x) { + int iter = iter_base; + int staged = staged0; + for (int work = tile_start; work < n_tiles; work += gridDim.x) { + const int tile_id = work / SPLIT; + const int part = work % SPLIT; + const int tok = part / (SPLIT / 2); + const int base = tile_id * OUT_ATOM; + const int grow0 = part * ROWS; + const int col = threadIdx.x & (OUT_ATOM - 1); + const int first_row = threadIdx.x >> 6; + constexpr int NR = ROWS / 4; + float held[NR]; + if (warpgroup::groupid() == 1) { + if (warpgroup::laneid() == 0) { + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + if (staged) { --staged; } else { + wait(drained[s], phase); + tma::expect_bytes(filled[s], + sizeof(weight_tile) + sizeof(act16_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s], w, {tile_id, i}, filled[s]); } + const int a_col = out_tiles_per_expert + ? (tile_id / out_tiles_per_expert) * + act_k_tiles_per_expert + i + : i; + tma::load_async<2, cache_policy::EVICT_LAST>( + *reinterpret_cast(&a_s[s]), + a, {tok, a_col}, filled[s]); + } + } + #pragma unroll + for (int j = 0; j < NR; ++j) + held[j] = __bfloat162float(residual[ + (size_t)(grow0 + first_row + 4 * j) * dest_stride + + base + col]); + } else { + #pragma unroll + for (int j = 0; j < NR; ++j) + held[j] = __bfloat162float(residual[ + (size_t)(grow0 + first_row + 4 * j) * dest_stride + + base + col]); + rt_fl<16, 16> accum; + warp::zero(accum); + int prior_s = -1; + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + if constexpr (SPLIT == 4) + bf16_mma_ABt_n8(accum, w_s[s], + *reinterpret_cast(&a_s[s]), part & 1); + else + warpgroup::mma_ABt(accum, w_s[s], + *reinterpret_cast(&a_s[s])); + if (prior_s >= 0) { + warpgroup::mma_async_wait<1>(); + if (warp::laneid() == 0) arrive(drained[prior_s]); + } + prior_s = s; + } + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[prior_s]); + warpgroup::store(o16, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + } + __syncthreads(); + #pragma unroll + for (int j = 0; j < NR; ++j) { + const int row = first_row + 4 * j; + residual[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16( + held[j] + __bfloat162float(o16[int2{col, row}])); + } + __syncthreads(); + } + return; + } + if (warpgroup::groupid() == 1) { + if (warpgroup::laneid() == 0) { + int iter = iter_base; + int staged = staged0; + for (int work = tile_start; work < n_tiles; work += gridDim.x) { + const int tile_id = work / SPLIT; + const int part = work % SPLIT; + const int tok = part / (SPLIT / 2); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + if (staged) { --staged; } else { + wait(drained[s], phase); + tma::expect_bytes(filled[s], + sizeof(weight_tile) + sizeof(act16_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s], w, {tile_id, i}, filled[s]); } + const int a_col = out_tiles_per_expert + ? (tile_id / out_tiles_per_expert) * act_k_tiles_per_expert + i + : i; + tma::load_async<2, cache_policy::EVICT_LAST>( + *reinterpret_cast(&a_s[s]), a, {tok, a_col}, filled[s]); + } + } + } + } else { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) + { + const int tile_id = work / SPLIT; + const int part = work % SPLIT; + const int tok = part / (SPLIT / 2); + const int base = tile_id * OUT_ATOM; + const int grow0 = part * ROWS; + const float bv0 = (epi == EPI_BIAS && bias) + ? __bfloat162float(bias[base + (warpgroup::laneid() & (OUT_ATOM - 1))]) + : 0.f; + constexpr int NRH = OUT_ATOM * ROWS / 128; + float held[NRH]; + if (epi == EPI_RESIDUAL) { + #pragma unroll + for (int j = 0; j < NRH; ++j) { + const int idx = warpgroup::laneid() + j * 128; + held[j] = __bfloat162float(residual[ + (size_t)(grow0 + idx / OUT_ATOM) * dest_stride + + base + idx % OUT_ATOM]); + } + } + rt_fl<16, 16> accum; + warp::zero(accum); + int prior_s = -1; + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + if constexpr (SPLIT == 4) { + bf16_mma_ABt_n8(accum, w_s[s], + *reinterpret_cast(&a_s[s]), part & 1); + } else { + warpgroup::mma_ABt(accum, w_s[s], + *reinterpret_cast(&a_s[s])); + } + if (prior_s >= 0) { + warpgroup::mma_async_wait<1>(); + if (warp::laneid() == 0) arrive(drained[prior_s]); + } + prior_s = s; + } + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[prior_s]); + if (epi == EPI_BIAS) { + const float bv = bv0; + warpgroup::store(of16, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * ROWS; + idx += 128) { + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + dest[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(of16[int2{col, row}] + bv); + } + warpgroup::sync(1); + continue; + } + if (epi == EPI_RESIDUAL) { + constexpr int NR = NRH; + warpgroup::store(o16, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + #pragma unroll + for (int j = 0; j < NR; ++j) { + const int idx = warpgroup::laneid() + j * 128; + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + residual[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(held[j] + + __bfloat162float(o16[int2{col, row}])); + } + warpgroup::sync(1); + continue; + } + warpgroup::store(o16, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * ROWS; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o16[int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_w[grow * BB_EXPERTS + expert]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * + D_MODEL + local] = __fmul_rn(wt, v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1); + } + (void)my_tiles; + } +} +__device__ void b32_fp8_body(B32QWeight qweight,B32QContext *qctx, + const B32QGL &w, const B32QGL &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + B32QW (&w_s)[B32_STAGES], B32QA (&a_s)[B32_STAGES], out_tile &o_s, + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + int &tile_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0) { + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + const int iters = k_total / B32_KCHUNK; + const int tile_start = advance_tile_cursor(tile_io, n_tiles); + int tiles_here = 0; + for (int t = tile_start; t < n_tiles; t += gridDim.x) ++tiles_here; + const int iter_base = iter_io; + iter_io += tiles_here * iters * n_tok; + const int my_tiles = (n_tiles - tile_start + gridDim.x - 1) / gridDim.x; + if (warpgroup::groupid() == 1) { + if (warpgroup::laneid() == 0) { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) { + const int tile_id = mk2_expert_tile(work, active_mask, compact_tiles); + for (int tok = 0; tok < n_tok; ++tok) + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % B32_STAGES; + const int phase = (iter / B32_STAGES) & 1; + wait(drained[s], phase); + tma::expect_bytes(filled[s], + sizeof(B32QW) + sizeof(B32QA)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s], w, {tile_id * iters + i, 0}, filled[s]); + const int a_col = out_tiles_per_expert + ? (tile_id / out_tiles_per_expert) * act_k_tiles_per_expert + i + : i; + tma::load_async<2, cache_policy::EVICT_LAST>( + a_s[s], a, {tok, a_col}, filled[s]); + } + } + } + } else { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) + for (int tok = 0; tok < n_tok; ++tok) { + const int tile_id = mk2_expert_tile(work, active_mask, compact_tiles); + const int base = tile_id * OUT_ATOM; + const int grow0 = tok * BATCH_PAD; + const int lane = warpgroup::laneid(); + const int col = lane & (OUT_ATOM - 1); + const int first_row = lane >> 6; + const float output_row_scale = qweight.scale[base + col]; + const float bv = (epi == EPI_BIAS && bias) + ? __bfloat162float(bias[base + col]) : 0.f; + unsigned char *const qout_base = + reinterpret_cast(qctx->views[1].raw_ptr); + rt_fl<16, BATCH_PAD> accum; + warp::zero(accum); + int prior_s = -1; + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % B32_STAGES; + const int phase = (iter / B32_STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s], a_s[s]); + if (prior_s >= 0) { + warpgroup::mma_async_wait<1>(); + if (warp::laneid() == 0) arrive(drained[prior_s]); + } + prior_s = s; + } + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[prior_s]); + if (epi == EPI_BIAS) { + warpgroup::store(o_sf, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + dest[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(o_sf[int2{col, row}] + bv); + } + warpgroup::sync(1); + continue; + } + warpgroup::store(o_sf, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + if (rows_valid == BATCH && epi == EPI_GELU) { + unsigned char *const qout = qout_base; + #pragma unroll + for (int row = first_row; row < BATCH; row += 2) { + const size_t at = + (size_t)(grow0 + row) * dest_stride + base + col; + const float v = __bfloat162float(__float2bfloat16( + o_sf[int2{col, row}] * output_row_scale * + qweight.activation_scale)); + const bf16 activated = + __float2bfloat16(gelu_tanh_fp32(v)); + qout[at] = b32_qencode(__bfloat162float(activated) / + qweight.next_scale); + } + } else if (rows_valid == BATCH && epi == EPI_RESIDUAL) { + float held[BATCH / 2]; + #pragma unroll + for (int j = 0; j < BATCH / 2; ++j) + held[j] = __bfloat162float(residual[ + (size_t)(grow0 + first_row + 2 * j) * dest_stride + + base + col]); + #pragma unroll + for (int j = 0; j < BATCH / 2; ++j) { + const int row = first_row + 2 * j; + const float v = __bfloat162float(__float2bfloat16( + o_sf[int2{col, row}] * output_row_scale * + qweight.activation_scale)); + residual[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(held[j] + v); + } + } else for (int idx = lane; idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int generic_col = idx % OUT_ATOM; + const float v = __bfloat162float(__float2bfloat16( + o_sf[int2{generic_col, row}] * + qweight.scale[base + generic_col] * + qweight.activation_scale)); + const int grow = grow0 + row; + const size_t at = + (size_t)grow * dest_stride + base + generic_col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + generic_col; + const float wt = moe_w[grow * BB_EXPERTS + expert]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * + D_MODEL + local] = __fmul_rn(wt, v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + { + const bf16 activated = + __float2bfloat16(gelu_tanh_fp32(v)); + dest[at] = activated; + reinterpret_cast( + qctx->views[1].raw_ptr)[at] = + b32_qencode(__bfloat162float(activated) / + qweight.next_scale); + } + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1); + } + (void)my_tiles; + } +} +__device__ void b32_fp8_gelu_all_warps( + B32QWeight qweight, B32QContext *qctx, + const B32QGL &w, const B32QGL &a, int n_out, int k_total, + B32QW (&w_s)[B32_STAGES], B32QA (&a_s)[B32_STAGES], + out_tile_f &o_sf, semaphore *filled, semaphore *drained, + int &iter_io, int &tile_io, int prestage = 0) { + const int n_tiles = n_out / OUT_ATOM; + const int iters = k_total / B32_KCHUNK; + const int tile_start = advance_tile_cursor(tile_io, n_tiles); + int tiles_here = 0; + for (int t = tile_start; t < n_tiles; t += gridDim.x) ++tiles_here; + const int iter_base = iter_io; + iter_io += tiles_here * iters; + const int staged0 = staged_rounds(prestage, tile_start, n_tiles, iters); + unsigned char *const qout = + reinterpret_cast(qctx->views[1].raw_ptr); + const bool b32_producer = warpgroup::groupid() == 1; + const bool b32_tma_lane = b32_producer && warpgroup::laneid() == 0; + int issue_iter = iter_base; + auto issue_tile = [&](int tile_id, int ordinal) { + for (int i = 0; i < iters; ++i, ++issue_iter) { + const int s = issue_iter % B32_STAGES; + const int phase = (issue_iter / B32_STAGES) & 1; + const int as = (iter_base + i) % B32_STAGES; + const bool reuse_activation = ordinal != 0; + if (ordinal == 0 && i < staged0) {} else { + wait(drained[s], phase); + tma::expect_bytes(filled[s], sizeof(B32QW) + + (reuse_activation ? 0 : sizeof(B32QA))); + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s], w, {tile_id * iters + i, 0}, filled[s]); } + if (!reuse_activation) + tma::load_async<2, cache_policy::EVICT_LAST>( + a_s[as], a, {0, i}, filled[s]); + } + }; + if (b32_tma_lane && tile_start < n_tiles) issue_tile(tile_start, 0); + int iter = iter_base; + int work_ordinal = 0; + for (int work = tile_start; work < n_tiles; + work += gridDim.x, ++work_ordinal) { + const int tile_id = work; + const int next_work = work + (int)gridDim.x; + const int base = tile_id * OUT_ATOM; + const int col = threadIdx.x & (OUT_ATOM - 1); + float output_row_scale; + if (b32_producer) { + if (b32_tma_lane && next_work < n_tiles) + issue_tile(next_work, work_ordinal + 1); + output_row_scale = qweight.scale[base + col]; + } else { + output_row_scale = qweight.scale[base + col]; + rt_fl<16, BATCH_PAD> accum; + warp::zero(accum); + int prior_s = -1; + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % B32_STAGES; + const int phase = (iter / B32_STAGES) & 1; + const int as = (iter_base + i) % B32_STAGES; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s], a_s[as]); + if (prior_s >= 0) { + warpgroup::mma_async_wait<1>(); + if (warp::laneid() == 0) arrive(drained[prior_s]); + } + prior_s = s; + } + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[prior_s]); + warpgroup::store(o_sf, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + } + __syncthreads(); + const int first_row = threadIdx.x >> 6; + #pragma unroll + for (int row = first_row; row < BATCH; row += 4) { + const size_t at = (size_t)row * n_out + base + col; + const float v = __bfloat162float(__float2bfloat16( + o_sf[int2{col, row}] * output_row_scale * + qweight.activation_scale)); + const bf16 activated = __float2bfloat16(gelu_tanh_fp32(v)); + qout[at] = b32_qencode(__bfloat162float(activated) / + qweight.next_scale); + } + if (next_work < n_tiles) __syncthreads(); + } +} +__device__ __forceinline__ void b32_wgmma_n16_e4m3( + rt_fl<16, 16> &d, uint64_t a_desc, uint64_t b_desc) { + const int scale_d = 1; + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %10, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k32.f32.e4m3.e4m3 " + "{%0, %1, %2, %3, %4, %5, %6, %7}, %8, %9, p, 1, 1;\n" + "}\n" + : "+f"(d.tiles[0][0].data[0].x), "+f"(d.tiles[0][0].data[0].y), + "+f"(d.tiles[0][0].data[1].x), "+f"(d.tiles[0][0].data[1].y), + "+f"(d.tiles[0][0].data[2].x), "+f"(d.tiles[0][0].data[2].y), + "+f"(d.tiles[0][0].data[3].x), "+f"(d.tiles[0][0].data[3].y) + : "l"(a_desc), "l"(b_desc), "r"(scale_d)); +} +__device__ __forceinline__ void b32_mma_ABt_n16( + rt_fl<16, 16> &d, const B32QW &a, const B32QA16 &b) { + kittens::st_descriptor a_desc(a); + kittens::st_descriptor b_desc(b); + constexpr int K = B32QW::cols / kittens::TILE_COL_DIM; + warpgroup::mma_fence(d); + #pragma unroll + for (int k = 0; k < K; ++k) + b32_wgmma_n16_e4m3(d, a_desc.chunk_descriptor(k), + b_desc.chunk_descriptor(k)); + warpgroup::mma_commit_group(); +} +__device__ void b32_fp8_body_n16(B32QWeight qweight, B32QContext *qctx, + const B32QGL &w, const B32QGL &a, int n_out, int k_total, + bf16 *residual, int dest_stride, + B32QW (&w_s)[B32_STAGES], B32QA (&a_s)[B32_STAGES], + semaphore *filled, semaphore *drained, out_tile_f &o_sf, + int &iter_io, int &tile_io, int prestage = 0) { + out16_float &of16 = *reinterpret_cast(&o_sf); + const int n_tiles = 2 * (n_out / OUT_ATOM); + const int iters = k_total / B32_KCHUNK; + const int tile_start = advance_tile_cursor(tile_io, n_tiles); + int tiles_here = 0; + for (int t = tile_start; t < n_tiles; t += gridDim.x) ++tiles_here; + const int iter_base = iter_io; + iter_io += tiles_here * iters; + int iter = iter_base; + int staged = staged_rounds(prestage, tile_start, n_tiles, iters); + for (int work = tile_start; work < n_tiles; work += gridDim.x) { + const int tile_id = work >> 1, tok = work & 1; + const int base = tile_id * OUT_ATOM; + const int grow0 = tok * 16; + const int col = threadIdx.x & (OUT_ATOM - 1); + const int first_row = threadIdx.x >> 6; + const float output_row_scale = qweight.scale[base + col]; + constexpr int NR = 16 / 4; + float held[NR]; + if (warpgroup::groupid() == 1) { + if (warpgroup::laneid() == 0) { + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % B32_STAGES; + const int phase = (iter / B32_STAGES) & 1; + if (staged) { --staged; } else { + wait(drained[s], phase); + tma::expect_bytes(filled[s], + sizeof(B32QW) + sizeof(B32QA16)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s], w, {tile_id * iters + i, 0}, filled[s]); } + tma::load_async<2, cache_policy::EVICT_LAST>( + *reinterpret_cast(&a_s[s]), a, {tok, i}, + filled[s]); + } + } + #pragma unroll + for (int j = 0; j < NR; ++j) + held[j] = __bfloat162float(residual[ + (size_t)(grow0 + first_row + 4 * j) * dest_stride + + base + col]); + } else { + #pragma unroll + for (int j = 0; j < NR; ++j) + held[j] = __bfloat162float(residual[ + (size_t)(grow0 + first_row + 4 * j) * dest_stride + + base + col]); + rt_fl<16, 16> accum; + warp::zero(accum); + int prior_s = -1; + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % B32_STAGES; + const int phase = (iter / B32_STAGES) & 1; + wait(filled[s], phase); + b32_mma_ABt_n16(accum, w_s[s], + *reinterpret_cast(&a_s[s])); + if (prior_s >= 0) { + warpgroup::mma_async_wait<1>(); + if (warp::laneid() == 0) arrive(drained[prior_s]); + } + prior_s = s; + } + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[prior_s]); + warpgroup::store(of16, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + } + __syncthreads(); + #pragma unroll + for (int j = 0; j < NR; ++j) { + const int row = first_row + 4 * j; + const float v = __bfloat162float(__float2bfloat16( + of16[int2{col, row}] * output_row_scale * + qweight.activation_scale)); + residual[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(held[j] + v); + } + __syncthreads(); + } +} +__device__ void b32_fp8_gemm_n16(B32QWeight qweight, B32QContext *qctx, + int n_out, int k_total, bf16 *residual, int dest_stride, + weight_tile (&w_s)[STAGES], out_tile_f &o_sf, + semaphore *filled, semaphore *drained, int &iter_io, int &tile_io, + int prestage = 0) { + auto &qw = *reinterpret_cast(&w_s[0]); + auto &qa = *reinterpret_cast( + (char *)&w_s[0] + B32_STAGES * sizeof(B32QW)); + const B32QGL &input = qctx->views[k_total == D_MODEL ? 0 : 1]; + b32_fp8_body_n16(qweight, qctx, *qweight.view, input, n_out, k_total, + residual, dest_stride, qw, qa, filled, drained, o_sf, + iter_io, tile_io, prestage); +} +__device__ void b32_fp8_gemm(B32QWeight qweight,B32QContext *qctx, + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES], act_tile (&a_s)[STAGES], out_tile &o_s, + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + int &tile_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0, + int prestage = 0) { + auto &qw=*reinterpret_cast(&w_s[0]); + auto &qa=*reinterpret_cast((char*)&w_s[0]+B32_STAGES*sizeof(B32QW)); + const B32QGL &input=qctx->views[k_total==D_MODEL?0:1]; + b32_fp8_gelu_all_warps(qweight, qctx, *qweight.view, input, + n_out, k_total, qw, qa, o_sf, filled, drained, iter_io, tile_io, + prestage); +} +#include "sparse_moe_impl.cuh" +constexpr int PF_CHUNKS = 512 / KCHUNK; +constexpr int PF_CHUNKS_Q = + PF_CHUNKS * (int)(sizeof(weight_tile) / sizeof(B32QW)); +static_assert(PF_CHUNKS_Q < B32_STAGES, "pre-staging must leave a free slot"); +static_assert(PF_CHUNKS_Q <= D_MODEL / B32_KCHUNK, + "ff_in sweeps D_MODEL/B32_KCHUNK rounds per tile"); +template +__device__ __forceinline__ void prefetch_next(ST &slot, const GL &w, int n_out, + int tile_io) { + if (threadIdx.x != 0) return; + const int n_tiles = SPLIT * (n_out / OUT_ATOM); + int t = (int)blockIdx.x - tile_io; + if (t < 0) t += (int)gridDim.x; + if (t >= n_tiles) return; + const int tile_id = SPLIT == 1 ? t : t / SPLIT; + #pragma unroll + for (int k = 0; k < PF; ++k) { + if constexpr (KITERS) + tma::prefetch<2, cache_policy::EVICT_FIRST>( + slot, w, {tile_id * KITERS + k, 0}); + else + tma::prefetch<2, cache_policy::EVICT_FIRST>(slot, w, {tile_id, k}); + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr, + const unsigned char *native_predictor_active_mask = nullptr, + const unsigned char *native_progress_mask = nullptr) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&w_s)[STAGES] = alloc.allocate(); + act_tile (&a_s)[STAGES] = alloc.allocate(); + out_tile (&o_s) = alloc.allocate(); + out_tile_f (&o_sf) = alloc.allocate(); + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ semaphore qfilled[B32_STAGES], qdrained[B32_STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * WARPS + 2]; + __shared__ sampler_storage sort_storage; + static_assert(sizeof(SparseMoeSchedule) <= sizeof(sampler_storage)); + SparseMoeSchedule &sparse_schedule = + *reinterpret_cast(&sort_storage); + unsigned int bar_next_target = 0u; + if (threadIdx.x == 0) { + bar_next_target = + *reinterpret_cast(g.barrier + 1); + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS, 0); + for (int c = 0; c < CONSUMER_WARPS; ++c) arrive(drained[s]); + } + for (int s = 0; s < B32_STAGES; ++s) { + init_semaphore(qfilled[s], 0, 1); + init_semaphore(qdrained[s], CONSUMER_WARPS, 0); + for (int c = 0; c < CONSUMER_WARPS; ++c) arrive(qdrained[s]); + } + } + int gemm_iter = 0; + int b32_iter = 0; + int gemm_tile = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + if (layer == 0 && g.native_embedding) { + bb_norm_residual( + bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, g.native_embedding); + } else { + bb_norm_residual( + bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch); + } + const int mi = bb->mamba_of_layer[layer]; + const bool moe_local_publish = + !BB_SKIP_MOE && mi < 0 && bb->attn_of_layer[layer] < 0; + if (moe_local_publish) __syncthreads(); + else row_grid_barrier_pf(g.barrier, bar_next_target, [&] { + if (mi >= 0) { + if constexpr (!BB_SKIP_MAMBA) + stage_next(w_s, bb->mamba[mi].in_proj, + MB_IN_PROJ, D_MODEL / KCHUNK, gemm_tile, gemm_iter, + sizeof(weight_tile) + sizeof(act_tile), filled, drained); + } else { + const int ai = bb->attn_of_layer[layer]; + if constexpr (!BB_SKIP_ATTN) + if (ai >= 0) + stage_next(w_s, bb->attn[ai].qkv, BB_QKV, + D_MODEL / KCHUNK, gemm_tile, gemm_iter, + sizeof(weight_tile) + sizeof(act_tile), filled, + drained); + } + }); + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier, bar_next_target); + grid_barrier(g.barrier, bar_next_target); + grid_barrier(g.barrier, bar_next_target); + grid_barrier(g.barrier, bar_next_target); + grid_barrier(g.barrier, bar_next_target); + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, 1, BATCH, nullptr, 0, + PF_CHUNKS); + grid_barrier(g.barrier, bar_next_target); + if constexpr (MK2_SSM_SPLITS == 1 && MB_HEADS_PER_GROUP == WARPS) { + mamba_mixer_local(L, bb->zxbcdt, bb->xbc_act, bb->gated, + bb->slot_map, native_progress_mask); + } else { + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map, native_progress_mask); + grid_barrier(g.barrier, bar_next_target); + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map, + native_progress_mask); + grid_barrier(g.barrier, bar_next_target); + mamba_gated_norm(L, bb->zxbcdt, bb->gated); + } + if (gridDim.x >= 96) { + grid_barrier_pf(g.barrier, bar_next_target, [&] { + stage_next(w_s, L.out_proj, D_MODEL, + MB_INNER / KCHUNK, gemm_tile, gemm_iter, + sizeof(weight_tile) + sizeof(act16_tile), filled, drained); }); + gemm_narrow<4>(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, 1, BATCH, nullptr, 0, PF_CHUNKS); + } else { + grid_barrier_pf(g.barrier, bar_next_target, [&] { + stage_next(w_s, L.out_proj, D_MODEL, + MB_INNER / KCHUNK, gemm_tile, gemm_iter, + sizeof(weight_tile) + sizeof(act16_tile), filled, drained); }); + gemm_narrow<2>(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, 1, BATCH, nullptr, 0, PF_CHUNKS); + } + grid_barrier(g.barrier, bar_next_target); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier, bar_next_target); + grid_barrier(g.barrier, bar_next_target); + grid_barrier(g.barrier, bar_next_target); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, 1, BATCH, nullptr, 0, + PF_CHUNKS); + grid_barrier(g.barrier, bar_next_target); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len, + native_progress_mask); + const bool oproj_n8 = gridDim.x >= 96; + grid_barrier_pf(g.barrier, bar_next_target, [&] { + if (oproj_n8) + stage_next(w_s, L.o_proj, D_MODEL, + D_MODEL / KCHUNK, gemm_tile, gemm_iter, + sizeof(weight_tile) + sizeof(act16_tile), filled, drained); + else + stage_next(w_s, L.o_proj, D_MODEL, + D_MODEL / KCHUNK, gemm_tile, gemm_iter, + sizeof(weight_tile) + sizeof(act16_tile), filled, drained); }); + if (oproj_n8) + gemm_narrow<4>(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, 1, BATCH, nullptr, 0, PF_CHUNKS); + else + gemm_narrow<2>(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, 1, BATCH, nullptr, 0, PF_CHUNKS); + grid_barrier(g.barrier, bar_next_target); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier, bar_next_target); + grid_barrier(g.barrier, bar_next_target); + grid_barrier(g.barrier, bar_next_target); + grid_barrier(g.barrier, bar_next_target); + grid_barrier(g.barrier, bar_next_target); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; +#ifndef MK2_STAGE0_NO_DIAGNOSTIC_TAPS + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS); + grid_barrier(g.barrier, bar_next_target); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else +#endif + { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w); + } + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + } + grid_barrier(g.barrier, bar_next_target); + if constexpr (MK2_SPARSE_MOE && BATCH >= 16) { + if (bb->sparse_moe) { + sparse_moe_pack(bb, sparse_schedule); + grid_barrier_pf(g.barrier, bar_next_target, [&] { + stage_next_sparse_up(w_s, L.exp_up, sparse_schedule, + gemm_tile, gemm_iter, filled, drained, + PF_CHUNKS); }); + } + } + if (bb->sparse_moe) { + sparse_moe_gemm(L.exp_up, bb->sparse_moe->gathered_v, + *bb->sparse_moe, bb->moe_h, bb->moe_w, bb->moe_out, + w_s, a_s, o_s, filled, drained, gemm_iter, gemm_tile, + sparse_schedule, PF_CHUNKS); + } else if constexpr (!BB_SPARSE_MOE_OWNED) { + gemm(L.exp_up, bb->normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, nullptr, 1, BATCH, + bb->moe_w, BB_MOE_INTER / OUT_ATOM); + } + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + if (bb->sparse_moe) + grid_barrier_pf(g.barrier, bar_next_target, [&] { + stage_next_sparse_down(w_s, L.exp_down, sparse_schedule, + gemm_tile, gemm_iter, filled, drained, PF_CHUNKS); }); + else + grid_barrier(g.barrier, bar_next_target); + if (bb->sparse_moe) { + sparse_moe_gemm(L.exp_down, bb->sparse_moe->hidden_v, + *bb->sparse_moe, bb->moe_h, bb->moe_w, bb->moe_out, + w_s, a_s, o_s, filled, drained, gemm_iter, gemm_tile, + sparse_schedule, PF_CHUNKS); + } else if constexpr (!BB_SPARSE_MOE_OWNED) { + gemm(L.exp_down, bb->moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + 1, BATCH, bb->moe_w, D_MODEL / OUT_ATOM); + } + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier, bar_next_target); + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + constexpr int PER_ROW = D_MODEL / THREADS; + const int lane = threadIdx.x & 31; + const unsigned selected = __ballot_sync( + 0xffffffffu, + lane < BB_EXPERTS && + bb->moe_w[row * BB_EXPERTS + lane] != 0.f); + bf16 folded[PER_ROW]; + #pragma unroll + for (int j = 0; j < PER_ROW; ++j) + folded[j] = bb_moe_fold_masked( + bb->moe_out, row, threadIdx.x + j * THREADS, + selected); + #pragma unroll + for (int j = 0; j < PER_ROW; ++j) + bb->h[(size_t)row * D_MODEL + threadIdx.x + j * THREADS] = + folded[j]; + } + } +#ifndef MK2_STAGE0_NO_DIAGNOSTIC_TAPS + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier, bar_next_target); + } +#endif + } +#ifndef MK2_STAGE0_NO_PREFILL + if (bb && bb->n_prefill_tokens > 0) { + const int n_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = n_total / BATCH; + grid_barrier(g.barrier, bar_next_target); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + bb->pf_h[idx] = bb->prompt[idx]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier, bar_next_target); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = bb->pf_n_seq > 0 ? bb->pf_n_seq : 1; + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T); + grid_barrier(g.barrier, bar_next_target); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier, bar_next_target); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier, bar_next_target); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + reinterpret_cast(bb->pf_moe_out), + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier, bar_next_target); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_gated, + T, slot, + reinterpret_cast(bb->pf_moe_out), + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier, bar_next_target); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier, bar_next_target); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier, bar_next_target); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier, bar_next_target); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier, bar_next_target); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier, bar_next_target); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier, bar_next_target); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier, bar_next_target); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier, bar_next_target); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T); + grid_barrier(g.barrier, bar_next_target); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, T); + } + grid_barrier(g.barrier, bar_next_target); + gemm(L.exp_up, bb->pf_normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier, bar_next_target); + gemm(L.exp_down, bb->pf_moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier, bar_next_target); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier, bar_next_target); + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = + bb_moe_fold_selected(bb->pf_moe_out, bb->pf_moe_w, idx); + grid_barrier(g.barrier, bar_next_target); + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier, bar_next_target); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, bb->pf_final, ln_scratch, + n_total); + grid_barrier(g.barrier, bar_next_target); + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier, bar_next_target); + } + } +#endif + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch); + if (bb->phoneme_out) + row_grid_barrier_pf(g.barrier, bar_next_target, [&] { + stage_next(w_s, bb->phoneme_w, BB_PHON_PADDED, + D_MODEL / KCHUNK, gemm_tile, gemm_iter, + sizeof(weight_tile) + sizeof(act_tile), filled, drained); }); + else + row_grid_barrier(g.barrier, bar_next_target); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + bb->phoneme_b, 0, 0, nullptr, nullptr, 1, BATCH, nullptr, 0, + PF_CHUNKS); + grid_barrier(g.barrier, bar_next_target); + __shared__ float ph_v[WARPS]; + __shared__ int ph_i[WARPS]; + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + if (!native_progresses(native_progress_mask, row)) continue; + float best = -CUDART_INF_F; + int arg = INT_MAX; + for (int n = threadIdx.x; n < BB_PHONEME; n += THREADS) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffffu, best, off); + const int oi = __shfl_xor_sync(0xffffffffu, arg, off); + if (ov > best || (ov == best && oi < arg)) { best = ov; arg = oi; } + } + if (lane == 0) { ph_v[warp] = best; ph_i[warp] = arg; } + __syncthreads(); + if (warp == 0) { + best = lane < WARPS ? ph_v[lane] : -CUDART_INF_F; + arg = lane < WARPS ? ph_i[lane] : INT_MAX; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_down_sync(0xffffffffu, best, off); + const int oi = __shfl_down_sync(0xffffffffu, arg, off); + if (ov > best || (ov == best && oi < arg)) { + best = ov; + arg = oi; + } + } + if (lane == 0) + bb->phoneme_out[row] = arg == INT_MAX ? 0 : arg; + } + if (row + gridDim.x < BATCH) __syncthreads(); + } + __syncthreads(); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier, bar_next_target); + } + } + __shared__ int native_run_predictor; + if (native_predictor_active_mask && bb && g.native_codes) { + const bool mine = threadIdx.x < BATCH && + native_predictor_active_mask[threadIdx.x] != 0 && + native_progresses(native_progress_mask, threadIdx.x); + const int any_active = __syncthreads_or(mine) ? 1 : 0; + if (threadIdx.x == 0) native_run_predictor = any_active; + __syncthreads(); + } + const bool skip_native_predictor = native_predictor_active_mask && bb && + g.native_codes && native_run_predictor == 0; + if (skip_native_predictor) { + for (int index = blockIdx.x * THREADS + threadIdx.x; + index < BATCH * N_POS; index += gridDim.x * THREADS) { + if (!native_progresses(native_progress_mask, index / N_POS)) continue; + g.sampled[index] = 0; + g.native_codes[index] = 0; + } + grid_barrier(g.barrier, bar_next_target); + } else { + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + if (pos == 0 && layer == 0) { + layernorm_rows(g.x, g.norm_self[layer], g.normed, + ln_scratch, nullptr, 1.f, g.pos_emb); + } else { + layernorm_rows( + g.x, g.norm_self[layer], g.normed, ln_scratch); + } +#endif + row_grid_barrier_pf(g.barrier, bar_next_target, [&] { + stage_next(w_s, g.qkv[layer], 3 * D_MODEL, + D_MODEL / KCHUNK, gemm_tile, gemm_iter, + sizeof(weight_tile) + sizeof(act_tile), filled, drained); }); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, 1, BATCH, nullptr, 0, PF_CHUNKS); + grid_barrier(g.barrier, bar_next_target); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x + warp * (int)gridDim.x; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + if (!native_progresses(native_progress_mask, row)) continue; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + constexpr int DPL = D_HEAD / 32; + bf16 kb[DPL], vb[DPL]; + float qv[DPL], kn[DPL], vn[DPL]; + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + kb[j] = k_new[i]; + vb[j] = v_new[i]; + qv[j] = __bfloat162float(q[i]); + kn[j] = __bfloat162float(kb[j]); + vn[j] = __bfloat162float(vb[j]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + g.kcache[cache_base + pos * D_HEAD + i] = kb[j]; + g.vcache[cache_base + pos * D_HEAD + i] = vb[j]; + } + __syncwarp(); + float my_sc = -CUDART_INF_F; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += qv[j] * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == t) my_sc = acc * scale; + } + { + float acc = 0.f; + #pragma unroll + for (int j = 0; j < DPL; ++j) acc += qv[j] * kn[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == pos) my_sc = acc * scale; + } + float m = my_sc; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + m = fmaxf(m, __shfl_xor_sync(0xffffffff, m, off)); + const float my_e = __expf(my_sc - m); + float s = 0.f; + #pragma unroll + for (int t = 0; t < N_POS; ++t) { + const float e = __shfl_sync(0xffffffff, my_e, t); + if (t <= pos) s += e; + } + const float inv = 1.f / s; + const float my_p = my_e * inv; + float ov[DPL] = {}; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *vt = g.vcache + cache_base + t * D_HEAD; + const float s = __shfl_sync(0xffffffff, my_p, t); + #pragma unroll + for (int j = 0; j < DPL; ++j) + ov[j] += s * __bfloat162float(vt[lane + j * 32]); + } + const float p_self = + __shfl_sync(0xffffffff, my_p, pos); + #pragma unroll + for (int j = 0; j < DPL; ++j) { + ov[j] += p_self * vn[j]; + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + + lane + j * 32] = __float2bfloat16(ov[j]); + } + } + } +#endif + if (gridDim.x >= 96) { + grid_barrier_pf(g.barrier, bar_next_target, [&] { + stage_next(w_s, g.o[layer], D_MODEL, + D_MODEL / KCHUNK, gemm_tile, gemm_iter, + sizeof(weight_tile) + sizeof(act16_tile), filled, drained); }); + gemm_narrow<4>(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, 1, BATCH, nullptr, 0, PF_CHUNKS); + } else { + grid_barrier_pf(g.barrier, bar_next_target, [&] { + stage_next(w_s, g.o[layer], D_MODEL, + D_MODEL / KCHUNK, gemm_tile, gemm_iter, + sizeof(weight_tile) + sizeof(act16_tile), filled, drained); }); + gemm_narrow<2>(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, 1, BATCH, nullptr, 0, PF_CHUNKS); + } + grid_barrier(g.barrier, bar_next_target); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch, g.b32_quant, g.b32_ff_in[layer].activation_scale); +#endif + row_grid_barrier_pf(g.barrier, bar_next_target, [&] { + stage_next( + *reinterpret_cast(&w_s[0]), + *g.b32_ff_in[layer].view, D_FFN, D_MODEL / B32_KCHUNK, + gemm_tile, b32_iter, sizeof(B32QW) + sizeof(B32QA), + qfilled, qdrained); }); + b32_fp8_gemm(g.b32_ff_in[layer],g.b32_quant,g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + w_s, a_s, o_s, qfilled, qdrained, o_sf, b32_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, 1, BATCH, nullptr, 0, PF_CHUNKS_Q); + grid_barrier_pf(g.barrier, bar_next_target, [&] { + stage_next( + *reinterpret_cast(&w_s[0]), + *g.b32_ff_out[layer].view, D_MODEL, D_FFN / B32_KCHUNK, + gemm_tile, b32_iter, sizeof(B32QW) + sizeof(B32QA16), + qfilled, qdrained); }); + b32_fp8_gemm_n16(g.b32_ff_out[layer], g.b32_quant, D_MODEL, D_FFN, + g.x, D_MODEL, w_s, o_sf, qfilled, qdrained, b32_iter, gemm_tile, + PF_CHUNKS_Q); + if (layer + 1 == n_layers) + grid_barrier_pf(g.barrier, bar_next_target, [&] { + stage_next(w_s, g.head[pos], VOCAB_PADDED, + D_MODEL / KCHUNK, gemm_tile, gemm_iter, + sizeof(weight_tile) + sizeof(act16_tile), filled, + drained); }); + else + grid_barrier(g.barrier, bar_next_target); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm_narrow<4>(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr); +#else + gemm_narrow<4>(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + g.head_bias[pos], 0, 0, nullptr, nullptr, 1, BATCH, nullptr, 0, + PF_CHUNKS); +#endif + grid_barrier(g.barrier, bar_next_target); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, + g.sampled_pos, sort_storage, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif +#ifdef ABLATE_SAMPLE + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { + __syncthreads(); + } else { + grid_barrier(g.barrier, bar_next_target); + } +#endif + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + constexpr int PER_THREAD = D_MODEL / THREADS; + float seed[PER_THREAD]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + seed[j] = __bfloat162float(emb[i]) + __bfloat162float(pe[i]); + } + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + g.x[(size_t)row * D_MODEL + i] = __float2bfloat16(seed[j]); + } + } + } + if (threadIdx.x == 0) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + g.sampled[row * N_POS + pos] = code; + if (g.native_codes && native_progresses(native_progress_mask, row)) { +#if MK2_INFLIGHT + g.native_codes[row * N_POS + pos] = + native_predictor_active_mask && !native_predictor_active_mask[row] + ? 0 : (long long)code; +#else + g.native_codes[row * N_POS + pos] = (long long)code; +#endif + } + } + } +#ifdef ABLATE_SAMPLE + grid_barrier(g.barrier, bar_next_target); +#else + if (pos + 1 == n_pos) grid_barrier(g.barrier, bar_next_target); +#endif + } + } + if (g.native_seq_len && blockIdx.x == 0 && threadIdx.x < BATCH && + native_progresses(native_progress_mask, threadIdx.x)) + g.native_seq_len[threadIdx.x] += 1; + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } + if (blockIdx.x == 0 && threadIdx.x == 0) + g.barrier[1] = bar_next_target; +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 64 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * (sizeof(weight_tile) + sizeof(act_tile)) + + sizeof(out_tile) + sizeof(out_tile_f) + 4096; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; + const unsigned char *native_predictor_active_mask = nullptr; + const unsigned char *native_progress_mask = nullptr; + int maximum_grid = 0; + bool grid_reset_pending = false; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +static B32QWeight b32_prepare_fp8(MK2Allocations &owned,const void *p,int n,int k,float activation_scale,float next_scale) { + B32QWeight q; + q.activation_scale=activation_scale; + q.next_scale=next_scale; + mk2_check(owned.alloc(&q.data,(size_t)n*k)); + mk2_check(owned.alloc(&q.scale,(size_t)n*sizeof(float))); + mk2_check(owned.alloc(&q.view,sizeof(B32QGL))); + b32_pack_fp8<<>>((const bf16*)p,q.data,q.scale,n,k); + mk2_check(cudaGetLastError()); + mk2_check(cudaDeviceSynchronize()); + if (n % OUT_ATOM || k % B32_KCHUNK) + throw std::runtime_error("tile-major E4M3 packing needs exact 64 x KCHUNK tiles"); + B32QGL view{q.data,nullptr,nullptr,(size_t)n*k/B32_KCHUNK,(size_t)B32_KCHUNK}; + mk2_check(cudaMemcpy(q.view,&view,sizeof(view),cudaMemcpyHostToDevice)); + return q; +} +static B32QContext *b32_prepare_fp8_context(MK2Allocations &owned) { + B32QT *data; + B32QContext *out; + mk2_check(owned.alloc(&data,(size_t)BATCH*(D_MODEL+D_FFN))); + mk2_check(owned.alloc(&out,sizeof(B32QContext))); + B32QContext context{{B32QGL{data,nullptr,nullptr,(size_t)BATCH,(size_t)D_MODEL},B32QGL{data+(size_t)BATCH*D_MODEL,nullptr,nullptr,(size_t)BATCH,(size_t)D_FFN}}}; + mk2_check(cudaMemcpy(out,&context,sizeof(context),cudaMemcpyHostToDevice)); + return out; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * (sizeof(weight_tile) + sizeof(act_tile)) + + sizeof(out_tile) + sizeof(out_tile_f) + 4096), 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = std::min(out->sms * out->blocks_per_sm, 4096); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 64 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + static const float b32_scales[]={4.659598214e-02f,3.152901786e-02f,2.427455357e-02f,1.708984375e-02f,3.166852679e-02f,5.133928571e-02f}; + g.b32_quant=b32_prepare_fp8_context(owned); + for(int i=0;iff_in[i],D_FFN,D_MODEL,b32_scales[2*i],b32_scales[2*i+1]); g.b32_ff_out[i]=b32_prepare_fp8(owned,w->ff_out[i],D_MODEL,D_FFN,b32_scales[2*i+1],1.f); } + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->maximum_grid = info.sms * info.blocks_per_sm; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + if constexpr (MK2_SPARSE_MOE && BATCH >= 16) { + bf16 *gathered = nullptr; + int *rows = nullptr, *counts = nullptr; + SparseMoe *device_sparse = nullptr; + mk2_check(owned.alloc(&gathered, + (size_t)BB_EXPERTS * BATCH_PAD * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&rows, (size_t)BB_EXPERTS * BATCH_PAD * sizeof(int))); + mk2_check(owned.alloc(&counts, BB_EXPERTS * sizeof(int))); + SparseMoe sparse{gathered, rows, counts, + sparse_act_gl{gathered, nullptr, nullptr, (size_t)BB_EXPERTS * BATCH_PAD, + (size_t)D_MODEL}, + sparse_act_gl{moe_h, nullptr, nullptr, (size_t)BB_EXPERTS * BATCH_PAD, + (size_t)BB_MOE_INTER}}; + mk2_check(owned.alloc(&device_sparse, sizeof(SparseMoe))); + mk2_check(cudaMemcpy(device_sparse, &sparse, sizeof(sparse), cudaMemcpyHostToDevice)); + hb.sparse_moe = device_sparse; + } + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.maximum_grid = pred->maximum_grid; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(s, &capture_status)); + if (st->pred.grid_reset_pending || capture_status != cudaStreamCaptureStatusNone) { + mk2_check(cudaMemsetAsync(st->pred.g.barrier, 0, 64 * sizeof(unsigned int), s)); + st->pred.grid_reset_pending = false; + } + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS, nullptr, + native ? st->pred.native_predictor_active_mask : nullptr, + native ? st->pred.native_progress_mask : nullptr); + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" int mk2_set_native_predictor_active_mask(MK2Predictor *p, + const unsigned char *mask) { + try { + if (!p || reinterpret_cast(p)->kind != 1) + throw std::runtime_error("startup mask requires a full Stage0 handle"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && st->pred.native_predictor_active_mask != mask) + throw std::runtime_error("startup mask pointer cannot change after first Stage0 enqueue"); + if (mask) + mk2_stage0_device_pointer(mask, st->pred.device, "native_predictor_active_mask"); + st->pred.native_predictor_active_mask = mask; + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown startup mask binding failure"; return 1; } +} +extern "C" int mk2_native_startup_skip_abi_version() { return 1; } +#include "../../../inflight/runtime.cuh" +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "native_prefill_impl.inc" +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *p, int grid) { + if (!p) { g_last_error = "grid setter requires full Stage0 handle"; return 1; } + auto *st = reinterpret_cast(p); + if (st->kind != 1 || grid < 1 || grid > st->maximum_grid) { + g_last_error = "requested Stage0 grid outside independently resident capacity"; + return 1; + } + if (st->grid != grid) st->grid_reset_pending = true; + st->grid = grid; + g_last_error.clear(); + return 0; +} +extern "C" int mk2_r8_bind_initial_state(MK2Predictor *p, const void *const *initial, const void *ready, const void *seq_len, int positions) { + try { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || st->bindings_initialized || !initial || !ready || !seq_len || positions <= 0) + throw std::runtime_error("invalid immutable initial-state binding"); + for (int i = 0; i < 14; ++i) { + mk2_stage0_device_pointer(initial[i], st->pred.device, "initial state"); + st->mb_host[i].r8_initial_ssm = static_cast(initial[i]); + st->mb_host[i].r8_initial_ready = static_cast(ready); + st->mb_host[i].r8_seq_len = static_cast(seq_len); + st->mb_host[i].r8_initial_positions = positions; + } + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" int mk2_b32_private_state_bytes() { return 2; } diff --git a/deploy/nemotron_native/kernels/b32/native/stage0/sparse_moe_impl.cuh b/deploy/nemotron_native/kernels/b32/native/stage0/sparse_moe_impl.cuh new file mode 100644 index 0000000..a436a6a --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/native/stage0/sparse_moe_impl.cuh @@ -0,0 +1,388 @@ +struct SparseMoeSchedule { + static constexpr int UP_TILES = BB_MOE_INTER / OUT_ATOM; + static constexpr int MAX_TILES = + BB_EXPERTS * (BATCH_PAD / SPARSE_N) * UP_TILES; + int base[BB_EXPERTS + 1]; + unsigned char row_tiles[BB_EXPERTS]; + unsigned char narrow[BB_EXPERTS]; + unsigned int route_mask[BB_EXPERTS]; + unsigned char expert[MAX_TILES]; + int plan_tile_io; + int plan_first; + int plan_step; + int plan_mine; + int plan_first2; + int plan_step2; + int plan_mine2; +}; +template +__device__ __forceinline__ void sparse_moe_tile( + const SparseMoeSchedule &schedule, int f, + int &expert, int &feature_tile, int &token_tile) { + const int lookup = DOWN ? (f >> 1) : f; + expert = schedule.expert[lookup]; + const int local = lookup - schedule.base[expert]; + const bool two_rows = schedule.row_tiles[expert] == 2; + if constexpr (DOWN) { + feature_tile = two_rows ? local : ((local << 1) | (f & 1)); + token_tile = two_rows ? (f & 1) : 0; + } else { + feature_tile = two_rows ? (local >> 1) : local; + token_tile = two_rows ? (local & 1) : 0; + } +} +__device__ __forceinline__ void sparse_moe_build_schedule( + const Backbone *bb, SparseMoeSchedule &schedule) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + for (int e = warp; e < BB_EXPERTS; e += WARPS) { + const bool hit = lane < BATCH && + bb->moe_w[lane * BB_EXPERTS + e] != 0.f; + const unsigned ballot = __ballot_sync(0xffffffffu, hit); + if (lane == 0) schedule.route_mask[e] = ballot; + } + __syncthreads(); + if (threadIdx.x < 32) { + const bool valid = lane < BB_EXPERTS; + const int count = valid ? __popc(schedule.route_mask[lane]) : 0; + const int rt = (count + SPARSE_N - 1) / SPARSE_N; + const int contrib = rt * SparseMoeSchedule::UP_TILES; + int scan = contrib; + #pragma unroll + for (int off = 1; off < 32; off <<= 1) { + const int up = __shfl_up_sync(0xffffffffu, scan, off); + if (lane >= off) scan += up; + } + if (valid) { + const int begin = scan - contrib; + schedule.base[lane] = begin; + schedule.row_tiles[lane] = (unsigned char)rt; + schedule.narrow[lane] = (unsigned char)(count <= 8); + #pragma unroll + for (int i = 0; i < 2 * SparseMoeSchedule::UP_TILES; ++i) + if (i < contrib) + schedule.expert[begin + i] = (unsigned char)lane; + } + if (lane == 31) schedule.base[BB_EXPERTS] = scan; + } + __syncthreads(); +} +__device__ void sparse_moe_pack(const Backbone *bb, + SparseMoeSchedule &schedule) { + const SparseMoe &s = *bb->sparse_moe; + __shared__ int count; + __shared__ int local_rows[BATCH]; + sparse_moe_build_schedule(bb, schedule); + const int experts_per_wave = min(int(gridDim.x), BB_EXPERTS); + const int first_expert = int(blockIdx.x) % experts_per_wave; + const int rank = int(blockIdx.x) / experts_per_wave; + const int team_size = (int(gridDim.x) + experts_per_wave - 1 - first_expert) / + experts_per_wave; + for (int e = first_expert; e < BB_EXPERTS; e += experts_per_wave) { + static_assert(BATCH <= 32, "the compaction below is one warp wide"); + if (threadIdx.x < 32) { + const int row = threadIdx.x; + const unsigned int ballot = schedule.route_mask[e]; + const bool hit = row < BATCH && ((ballot >> row) & 1u); + if (hit) local_rows[__popc(ballot & ((1u << row) - 1u))] = row; + if (row == 0) count = __popc(ballot); + } + __syncthreads(); + if (rank == 0) { + if (threadIdx.x == 0) s.counts[e] = count; + for (int i = threadIdx.x; i < count; i += THREADS) + s.rows[e * BATCH_PAD + i] = local_rows[i]; + } + constexpr int chunks = D_MODEL / 8; + const int padded = count <= 8 + ? 8 : ((count + SPARSE_N - 1) / SPARSE_N) * SPARSE_N; + { + const int first = rank * THREADS + threadIdx.x; + const int step = team_size * THREADS; + const int limit = padded * chunks; + constexpr int MAX_STRIPES = 2; + float4 staged[MAX_STRIPES]; + for (int base_i = first; base_i < limit; base_i += step * MAX_STRIPES) { + #pragma unroll + for (int u = 0; u < MAX_STRIPES; ++u) { + const int i = base_i + u * step; + if (i >= limit) continue; + const int row = i / chunks, col = i % chunks; + staged[u] = row < count + ? *reinterpret_cast( + bb->normed + (size_t)local_rows[row] * D_MODEL + col * 8) + : make_float4(0.f, 0.f, 0.f, 0.f); + } + #pragma unroll + for (int u = 0; u < MAX_STRIPES; ++u) { + const int i = base_i + u * step; + if (i >= limit) continue; + const int row = i / chunks, col = i % chunks; + *reinterpret_cast(s.gathered + + (size_t)(e * BATCH_PAD + row) * D_MODEL + col * 8) = staged[u]; + } + } + } + __syncthreads(); + } +} +template +__device__ __forceinline__ void sparse_moe_plan( + int &tile_io, int n_tiles, int &first, int &step, int &mine, + int &first2, int &step2, int &mine2) { + const int start = advance_tile_cursor(tile_io, n_tiles); + first = start; step = (int)gridDim.x; + mine = start < n_tiles ? 1 + (n_tiles - 1 - start) / gridDim.x : 0; + first2 = 0; step2 = 1; mine2 = 0; + if constexpr (DOWN) { + const int G = (int)gridDim.x; + constexpr int follow_tiles = D_MODEL / OUT_ATOM; + if (G > follow_tiles && n_tiles >= G) { + const int longc = G - follow_tiles; + const int waves = n_tiles / G - 1; + const int head = waves * G; + const int tail = n_tiles - head; + int d = (int)blockIdx.x - tile_io; + if (d < 0) d += G; + mine = waves; + if (d >= follow_tiles) { + const int i = d - follow_tiles; + first2 = head + i; + step2 = longc; + mine2 = i < tail ? 1 + (tail - 1 - i) / longc : 0; + } + } + } +} +__device__ __forceinline__ void stage_next_sparse_down( + weight_tile (&slots)[STAGES], const weight_gl &w, + SparseMoeSchedule &schedule, int tile_io, int iter_io, + semaphore *filled, semaphore *drained, int pf) { + if (threadIdx.x != 0) return; + constexpr int iters = BB_MOE_INTER / KCHUNK; + constexpr int out_tiles = D_MODEL / OUT_ATOM; + const int n_tiles = schedule.base[BB_EXPERTS] * 2; + int first, step, mine, first2, step2, mine2; + int cursor = tile_io; + sparse_moe_plan(cursor, n_tiles, first, step, mine, first2, step2, + mine2); + schedule.plan_tile_io = cursor; + schedule.plan_first = first; + schedule.plan_step = step; + schedule.plan_mine = mine; + schedule.plan_first2 = first2; + schedule.plan_step2 = step2; + schedule.plan_mine2 = mine2; + if (mine + mine2 == 0) return; + const int f = mine ? first : first2; + int e, feature_tile, token_tile; + sparse_moe_tile(schedule, f, e, feature_tile, token_tile); + const int weight_row = e * out_tiles + feature_tile; + const int n = pf < iters ? pf : iters; + for (int k = 0; k < n; ++k) { + const int iter = iter_io + k; + const int stage = iter % STAGES, phase = (iter / STAGES) & 1; + wait(drained[stage], phase); + tma::expect_bytes(filled[stage], + sizeof(weight_tile) + sizeof(sparse_act_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + slots[stage], w, {weight_row, k}, filled[stage]); + } +} +__device__ __forceinline__ void stage_next_sparse_up( + weight_tile (&slots)[STAGES], const weight_gl &w, + SparseMoeSchedule &schedule, int tile_io, int iter_io, + semaphore *filled, semaphore *drained, int pf) { + if (threadIdx.x != 0) return; + constexpr int iters = D_MODEL / KCHUNK; + constexpr int out_tiles = BB_MOE_INTER / OUT_ATOM; + const int n_tiles = schedule.base[BB_EXPERTS]; + int first, step, mine, first2, step2, mine2; + int cursor = tile_io; + sparse_moe_plan(cursor, n_tiles, first, step, mine, first2, step2, + mine2); + schedule.plan_tile_io = cursor; + schedule.plan_first = first; + schedule.plan_step = step; + schedule.plan_mine = mine; + schedule.plan_first2 = first2; + schedule.plan_step2 = step2; + schedule.plan_mine2 = mine2; + if (mine == 0) return; + int e, feature_tile, token_tile; + sparse_moe_tile(schedule, first, e, feature_tile, token_tile); + const int weight_row = e * out_tiles + feature_tile; + const int n = pf < iters ? pf : iters; + for (int k = 0; k < n; ++k) { + const int iter = iter_io + k; + const int stage = iter % STAGES, phase = (iter / STAGES) & 1; + wait(drained[stage], phase); + tma::expect_bytes(filled[stage], + sizeof(weight_tile) + sizeof(sparse_act_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + slots[stage], w, {weight_row, k}, filled[stage]); + } +} +template +__device__ void sparse_moe_gemm( + const weight_gl &w, const sparse_act_gl &a, const SparseMoe &s, + bf16 *hidden, const float *routes, float *contributions, + weight_tile (&w_s)[STAGES], act_tile (&a_s)[STAGES], out_tile &o_s, + semaphore *filled, semaphore *drained, int &iter_io, int &tile_io, + SparseMoeSchedule &schedule, int prestage = 0) { + constexpr int width = DOWN ? D_MODEL : BB_MOE_INTER; + constexpr int k_total = DOWN ? BB_MOE_INTER : D_MODEL; + constexpr int out_tiles = width / OUT_ATOM; + constexpr int iters = k_total / KCHUNK; + static_assert(BATCH_PAD >= SPARSE_N && BATCH_PAD % SPARSE_N == 0); + static_assert(BATCH <= 2 * SPARSE_N, + "sparse tile div/rem assumes one or two token tiles"); + static_assert(sizeof(sparse_act_tile) <= sizeof(act_tile)); + static_assert(sizeof(sparse_out_tile) <= sizeof(out_tile)); + tile_io = schedule.plan_tile_io; + const int first = schedule.plan_first; + const int step = schedule.plan_step; + const int mine = schedule.plan_mine; + const int first2 = schedule.plan_first2; + const int step2 = schedule.plan_step2; + const int mine2 = schedule.plan_mine2; + const int total = mine + mine2; + const int iter_base = iter_io; + iter_io += total * iters; + if (warpgroup::groupid() == 1) { + if (warpgroup::laneid() == 0) { + int iter = iter_base; + int staged = total ? (prestage < iters ? prestage : iters) : 0; + for (int n = 0, f = mine ? first : first2; n < total; ++n, + f = n < mine ? first + n * step + : first2 + (n - mine) * step2) { + int e, feature_tile, token_tile; + sparse_moe_tile(schedule, f, e, feature_tile, token_tile); + const int weight_row = e * out_tiles + feature_tile; + const int act_row = e * (BATCH_PAD / SPARSE_N) + token_tile; + for (int k = 0; k < iters; ++k, ++iter) { + const int stage = iter % STAGES, phase = (iter / STAGES) & 1; + if (staged) { --staged; } else { + wait(drained[stage], phase); + tma::expect_bytes(filled[stage], sizeof(weight_tile) + sizeof(sparse_act_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[stage], w, {weight_row, k}, filled[stage]); } + tma::load_async<2, cache_policy::EVICT_LAST>( + *reinterpret_cast(&a_s[stage]), + a, {act_row, k}, filled[stage]); + } + } + } + } else if (warpgroup::groupid() == 0) { + int iter = iter_base; + for (int n = 0, f = mine ? first : first2; n < total; ++n, + f = n < mine ? first + n * step + : first2 + (n - mine) * step2) { + int e, feature_tile, token_tile; + sparse_moe_tile(schedule, f, e, feature_tile, token_tile); + const bool narrow = schedule.narrow[e] != 0; + const int lane0 = warpgroup::laneid(); + constexpr int NROW0 = SPARSE_N / 2; + int original0[NROW0]; + float weight0[NROW0]; + if constexpr (DOWN) { + const int count0 = s.counts[e]; + const int first_row0 = lane0 >> 6; + if (narrow) { + #pragma unroll + for (int j = 0; j < NROW0 / 2; ++j) { + const int local = first_row0 + 2 * j; + original0[j] = local < count0 + ? s.rows[e * BATCH_PAD + local] : -1; + } + #pragma unroll + for (int j = 0; j < NROW0 / 2; ++j) + weight0[j] = original0[j] >= 0 + ? routes[original0[j] * BB_EXPERTS + e] : 0.f; + } else { + #pragma unroll + for (int j = 0; j < NROW0; ++j) { + const int local = token_tile * SPARSE_N + first_row0 + 2 * j; + original0[j] = local < count0 + ? s.rows[e * BATCH_PAD + local] : -1; + } + #pragma unroll + for (int j = 0; j < NROW0; ++j) + weight0[j] = original0[j] >= 0 + ? routes[original0[j] * BB_EXPERTS + e] : 0.f; + } + } + rt_fl<16, SPARSE_N> accum; + warp::zero(accum); + int prior_stage = -1; + for (int k = 0; k < iters; ++k, ++iter) { + const int stage = iter % STAGES, phase = (iter / STAGES) & 1; + wait(filled[stage], phase); + if (narrow) + bf16_mma_ABt_n8(accum, w_s[stage], + *reinterpret_cast(&a_s[stage]), 0); + else + warpgroup::mma_ABt(accum, w_s[stage], + *reinterpret_cast(&a_s[stage])); + if (prior_stage >= 0) { + warpgroup::mma_async_wait<1>(); + if (warp::laneid() == 0) arrive(drained[prior_stage]); + } + prior_stage = stage; + } + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[prior_stage]); + sparse_out_tile &out = *reinterpret_cast(&o_s); + warpgroup::store(out, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + const int lane = warpgroup::laneid(); + const int col = lane & (OUT_ATOM - 1); + const int first_row = lane >> 6; + if constexpr (DOWN) { + constexpr int NROW = NROW0; + const int *original = original0; + const float *weight = weight0; + const int output = feature_tile * OUT_ATOM + col; + if (narrow) { + #pragma unroll + for (int j = 0; j < NROW / 2; ++j) { + if (original[j] < 0) continue; + const float value = float(out[{col, first_row + 2 * j}]); + contributions[((size_t)original[j] * (BB_EXPERTS + 1) + e) * + D_MODEL + output] = __fmul_rn(weight[j], value); + } + } else { + #pragma unroll + for (int j = 0; j < NROW; ++j) { + if (original[j] < 0) continue; + const float value = float(out[{col, first_row + 2 * j}]); + contributions[((size_t)original[j] * (BB_EXPERTS + 1) + e) * + D_MODEL + output] = __fmul_rn(weight[j], value); + } + } + } else { + if (narrow) { + #pragma unroll + for (int row = first_row; row < 8; row += 2) { + const int packed = e * BATCH_PAD + row; + const int output = feature_tile * OUT_ATOM + col; + const float value = float(out[{col, row}]); + hidden[(size_t)packed * BB_MOE_INTER + output] = + bf16(value / (1.f + __expf(-value))); + } + } else { + #pragma unroll + for (int row = first_row; row < SPARSE_N; row += 2) { + const int local = token_tile * SPARSE_N + row; + const int packed = e * BATCH_PAD + local; + const int output = feature_tile * OUT_ATOM + col; + const float value = float(out[{col, row}]); + hidden[(size_t)packed * BB_MOE_INTER + output] = + bf16(value / (1.f + __expf(-value))); + } + } + } + warpgroup::sync(1); + } + } +} diff --git a/deploy/nemotron_native/kernels/b32/prefill.cu b/deploy/nemotron_native/kernels/b32/prefill.cu new file mode 100644 index 0000000..61fc0b9 --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/prefill.cu @@ -0,0 +1,20 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace b32_prefill { +#include "native/prefill/prefill.cu" +} diff --git a/deploy/nemotron_native/kernels/b32/stage0.cu b/deploy/nemotron_native/kernels/b32/stage0.cu new file mode 100644 index 0000000..f328edb --- /dev/null +++ b/deploy/nemotron_native/kernels/b32/stage0.cu @@ -0,0 +1,35 @@ +#define MK2_BATCH 32 +#define STAGES 4 +#define KCHUNK 256 +#define MK2_NO_MAIN 1 +#define MK2_DIRECT_INPUT 1 +#define MK2_B1_LOCAL_BARRIERS 1 +#define MK2_SKIP_HIDDEN_SNAPSHOT 1 +#define MK2_ACTIVE_EXPERTS 1 +#define MK2_SSM_SPLITS 1 +#define MK2_SPARSE_MOE 1 +#define MK2_STAGE0_NO_PREFILL 1 +#define MK2_STAGE0_NO_DIAGNOSTIC_TAPS 1 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#if MK2_INFLIGHT +#include +#endif +#include +#include +#include +#include +#include +#include +#include +#include +namespace b32_stage0 { +#include "native/stage0/predictor_mk2.cu" +} diff --git a/deploy/nemotron_native/kernels/b4096/codec.cu b/deploy/nemotron_native/kernels/b4096/codec.cu new file mode 100644 index 0000000..9e198ed --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/codec.cu @@ -0,0 +1 @@ +#include "regions/stage1/codec.cu" diff --git a/deploy/nemotron_native/kernels/b4096/prefill.cu b/deploy/nemotron_native/kernels/b4096/prefill.cu new file mode 100644 index 0000000..6d21d7e --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/prefill.cu @@ -0,0 +1 @@ +#include "regions/prefill/prefill.cu" diff --git a/deploy/nemotron_native/kernels/b4096/regions/prefill/native/prefill/dedup_native_prefill_impl.inc b/deploy/nemotron_native/kernels/b4096/regions/prefill/native/prefill/dedup_native_prefill_impl.inc new file mode 100644 index 0000000..ab49c4d --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/prefill/native/prefill/dedup_native_prefill_impl.inc @@ -0,0 +1,267 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 16 || BATCH == 32 || BATCH == 64 || BATCH == 128 || BATCH == 256 || BATCH == 512 || BATCH == 1024 || BATCH == 2048 || BATCH == 4096)) + throw std::runtime_error("native prefill requires a matching supported physical batch"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((std::uint64_t(c.prompt_positions) + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (rows > (INT_MAX - 132 * THREADS) / D_MODEL) + throw std::runtime_error("native prefill full-width prompt exceeds safe device loop bounds"); + if (padded > (INT_MAX - 132 * THREADS) / (BB_EXPERTS * BB_MOE_INTER)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((size_t(config->prompt_positions) + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + float *r9_workspace = nullptr; + mk2_check(owned.alloc(&r9_workspace,size_t(BATCH)*ssm_elements*sizeof(float))); + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + r9_state_t *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(r9_state_t))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(r9_state_t))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = r9_workspace; + st->mb_host[i].r9_private = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + mk2_check(owned.alloc(&st->mb_host[i].r8_canonical_ssm, ssm_elements * sizeof(float))); + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_dtda, 2 * MB_HEADS); + { + const size_t estride = padded; + const size_t grows = size_t(BB_EXPERTS) * estride; + mk2_check(owned.alloc(&hb.pf_moe_gath, grows * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_gath, 0, grows * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_hc, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_hc, 0, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_gidx, grows * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_gidx, 0, grows * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ecnt, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ecnt, 0, BB_EXPERTS * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ntiles, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ntiles, 0, BB_EXPERTS * sizeof(int))); + hb.pf_moe_estride = int(estride); + hb.pf_moe_gath_v = act_gl{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_v = act_gl{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + hb.pf_moe_gath_vn = act_gl_n{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_vn = act_gl_n{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + scratch_bytes += grows * (D_MODEL + BB_MOE_INTER) * sizeof(bf16); + } + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(owned.alloc(&hb.native_pf_cohort_check, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(hb.native_pf_cohort_check, 0, 2 * sizeof(unsigned int))); +#if PF_ACTIVE_TOKEN_EXPERTS + mk2_check(owned.alloc(&hb.native_pf_expert_masks, (padded/BATCH_PAD)*sizeof(unsigned int))); +#endif + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} +extern "C" int mk2_native_prefill_validate(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null compact prefill handle"); + auto *st = reinterpret_cast(p->base); + mk2_check(cudaStreamSynchronize(static_cast(stream))); + unsigned int status[2] = {0, 0}; + mk2_check(cudaMemcpy(status, st->host_bb.native_pf_cohort_check, + sizeof(status), cudaMemcpyDeviceToHost)); + if (status[0] != 0 || status[1] != 1) + throw std::runtime_error("compact prefill requires bit-identical prompts and identity slots"); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown compact prefill validation failure"; return 1; } +} +extern "C" int mk2_native_prefill_workspace_sequences() { return 1; } +extern "C" int mk2_native_prefill_effective_sequences(MK2NativePrefill *p, int *out) { + try { + if (!p || !out) throw std::runtime_error("null effective-sequence query"); + auto *st = reinterpret_cast(p->base); + mk2_check(cudaMemcpy(out, st->host_bb.native_pf_cohort_check + 1, + sizeof(int), cudaMemcpyDeviceToHost)); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" int mk2_r8_initial_state(MK2NativePrefill *p, void **initial, void **ready) { + if (!p || !initial || !ready) return 1; + auto *st = reinterpret_cast(p->base); + for (int i = 0; i < 14; ++i) initial[i] = st->mb_host[i].r8_canonical_ssm; + *ready = st->host_bb.native_pf_cohort_check + 1; + return 0; +} +extern "C" int mk2_r9_prefill_bind(MK2NativePrefill*p,const void*const*scales,const void*const*maps){if(!p)return 1;auto*st=(MK2Stage0State*)p->base;for(int i=0;i<14;++i){st->mb_host[i].r9_scales=(const float*)scales[i];st->mb_host[i].r9_map=(const int*)maps[i];}return (int)cudaMemcpy((void*)st->host_bb.mamba,st->mb_host.data(),14*sizeof(MambaLayer),cudaMemcpyHostToDevice);} diff --git a/deploy/nemotron_native/kernels/b4096/regions/prefill/native/prefill/dedup_predictor_mk2.cuh b/deploy/nemotron_native/kernels/b4096/regions/prefill/native/prefill/dedup_predictor_mk2.cuh new file mode 100644 index 0000000..f98dd53 --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/prefill/native/prefill/dedup_predictor_mk2.cuh @@ -0,0 +1,4600 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef PF_ACTIVE_TOKEN_EXPERTS +#define PF_ACTIVE_TOKEN_EXPERTS 0 +#endif +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +#ifndef MK2_PREFILL_ONLY +#define MK2_PREFILL_ONLY 0 +#endif +#ifndef MK2_BATCH_TILE +#define MK2_BATCH_TILE (((MK2_BATCH + 15) / 16) * 16) +#endif +constexpr int BATCH_PAD = MK2_BATCH_TILE; +static_assert(BATCH >= 1 && (MK2_PREFILL_ONLY || BATCH <= BATCH_PAD)); +static_assert(BATCH_PAD == 16 || BATCH_PAD == 32 || BATCH_PAD == 64 || BATCH_PAD == 128, + "native prefill tile must be 16, 32, 64 or 128"); +static_assert(MK2_PREFILL_ONLY || BATCH_PAD == ((BATCH + 15) / 16) * 16, + "a wider token tile requires a prefill-only library"); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_PW_MULTI +#define PF_SSM_PW_MULTI 4 +#endif +#ifndef PF_SSM_PW_SINGLE +#define PF_SSM_PW_SINGLE 2 +#endif +static_assert(MB_HEAD_DIM % PF_SSM_PW_MULTI == 0 && + MB_HEAD_DIM % PF_SSM_PW_SINGLE == 0, + "head_dim must split evenly into scan channel groups"); +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef MK2_SMEM_SLACK +#define MK2_SMEM_SLACK 1024 +#endif +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +#ifndef MK2_THREADS +#define MK2_THREADS 256 +#endif +constexpr int THREADS = MK2_THREADS; +constexpr int NORM_THREADS = 256; +constexpr int NORM_WARPS = NORM_THREADS / 32; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +#ifndef MK2_GEMM_CONSUMERS +#define MK2_GEMM_CONSUMERS 1 +#endif +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 +#define MK2_HEAD_STAGE_BYTES 0 +#else +#define MK2_HEAD_STAGE_BYTES sizeof(out_tile_f) +#endif +constexpr int NCONS = MK2_GEMM_CONSUMERS; +static_assert(NCONS >= 1 && NCONS <= 3, "consumers occupy warpgroups 0, 2 and 3"); +static_assert(THREADS >= (NCONS + 1) * 128, "one producer plus NCONS consumers"); +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +constexpr int PAD_N = 32; +using act_tile_n = st_bf; +using out_tile_n = st_bf; +using act_gl_n = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_MOE_ACC_C, EPI_ACC_F32 }; +#include +constexpr int R9_STATE_MODE=1; +using r9_state_t=bf16; +__device__ __forceinline__ int r9_scale_index(int head,int channel,int n) {return (((head/8)*4+(head%8)/2)*4+channel/6)*128+n;} +__device__ __forceinline__ float r9_load(r9_state_t v,float scale) { +return __bfloat162float(v); +} +__device__ __forceinline__ r9_state_t r9_store(float x,float scale) { +return __float2bfloat16(x); +} +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; + float *r8_canonical_ssm = nullptr; + r9_state_t *r9_private = nullptr; + const float *r9_scales = nullptr; + const int *r9_map = nullptr; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + unsigned int *native_pf_expert_masks = nullptr; + float *pf_dtda = nullptr; + bf16 *pf_moe_gath = nullptr; + bf16 *pf_moe_hc = nullptr; + int *pf_moe_gidx = nullptr; + int *pf_moe_ecnt = nullptr; + int *pf_moe_ntiles = nullptr; + int pf_moe_estride = 0; + act_gl pf_moe_gath_v; + act_gl pf_moe_hc_v; + act_gl_n pf_moe_gath_vn; + act_gl_n pf_moe_hc_vn; + unsigned int *native_pf_cohort_check = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + volatile unsigned int *generation = state + 1; + const unsigned int seen = *generation; + if (atomicAdd(state, 1u) == gridDim.x - 1) { + atomicExch(state, 0u); + atomicAdd(state + 1, 1u); + } else { + while (*generation == seen) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const bf16 *in = src + (size_t)row * D_MODEL; + float v[PER_THREAD]; + float sum = 0.f, sq = 0.f; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (threadIdx.x < NORM_THREADS) { + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + v[j] = __bfloat162float(in[threadIdx.x + j * NORM_THREADS]); + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + if (lane == 0) { scratch[warp] = sum; scratch[NORM_WARPS + warp] = sq; } + } + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f, q = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) { s += scratch[w]; q += scratch[NORM_WARPS + w]; } + scratch[2 * NORM_WARPS] = s / D_MODEL; + scratch[2 * NORM_WARPS + 1] = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + } + __syncthreads(); + const float mean = scratch[2 * NORM_WARPS], inv = scratch[2 * NORM_WARPS + 1]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + if (threadIdx.x >= NORM_THREADS) break; + const int i = threadIdx.x + j * NORM_THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * weight[i]); + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, + int idx, unsigned mask); +__device__ __forceinline__ void norm_sub_sync(int sub) { + if (THREADS == NORM_THREADS) { __syncthreads(); return; } + asm volatile("bar.sync %0, %1;" :: "r"(6 + sub), "n"(NORM_THREADS) : "memory"); +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, const float *fold_src = nullptr, + const float *fold_routes = nullptr) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + constexpr int NSUB = THREADS / NORM_THREADS; + if (threadIdx.x >= NSUB * NORM_THREADS) return; + const int sub = threadIdx.x / NORM_THREADS; + const int tid = threadIdx.x - sub * NORM_THREADS; + float *const sc = scratch + sub * (NORM_WARPS + 1); + for (int row = blockIdx.x * NSUB + sub; row < n_rows; + row += gridDim.x * NSUB) { + const size_t base = (size_t)row * D_MODEL; + unsigned fold_mask = (1u << BB_EXPERTS) - 1u; + if (fold_src && fold_routes) { + fold_mask = 0u; + const float *rw = fold_routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) + if (rw[e] != 0.f) fold_mask |= 1u << e; + } + float v[PER_THREAD]; + float sq = 0.f; + const int warp = tid >> 5, lane = tid & 31; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + bf16 hv; + if (fold_src) { + hv = bb_moe_fold_deterministic(fold_src, (int)(base + i), fold_mask); + h[base + i] = hv; + } else { + hv = h[base + i]; + } + v[j] = __bfloat162float(hv) + __bfloat162float(r[base + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + if (lane == 0) sc[warp] = sq; + norm_sub_sync(sub); + if (tid == 0) { + float s = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) s += sc[w]; + sc[NORM_WARPS] = rsqrtf(s / D_MODEL + RMS_EPS); + } + norm_sub_sync(sub); + const float inv = sc[NORM_WARPS]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + out[base + i] = __float2bfloat16( + __bfloat162float(weight[i]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + } + norm_sub_sync(sub); + } +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + const int c = idx - row * MB_CONV_DIM; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + const float raw = __bfloat162float( + zxbcdt[(size_t)row * MB_IN_PROJ + MB_INNER + c]); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + xbc_act[(size_t)row * MB_CONV_DIM + c] = __float2bfloat16(acc); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + const int pairs = BATCH * MB_HEADS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_HEADS * split); + const int head = (pair / split) % MB_HEADS; + const int first_channel = (pair % split) * channels; + const int group = head / MB_HEADS_PER_GROUP; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + L.dt_bias[head]); + const float dA = __expf(L.A[head] * dt); + const float Dh = L.D[head]; + float *st = L.ssm_state + (size_t)slot_map[row] * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + for (int p = first_channel; p < first_channel + channels; ++p) { + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + const size_t at = (size_t)p * MB_STATE + n; + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + static_assert(MB_CONV_K == 4, "packed conv taps assume a 4-tap kernel"); + constexpr int CPT = 4; + static_assert(MB_CONV_DIM % CPT == 0, "channel vector width"); + constexpr int CGROUPS = MB_CONV_DIM / CPT; + const int total = n_tokens * CGROUPS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / CGROUPS; + const int c = (idx - t * CGROUPS) * CPT; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + float wv[CPT][MB_CONV_K]; + { + const float4 *wsrc = reinterpret_cast( + L.conv_w + (size_t)c * MB_CONV_K); + #pragma unroll + for (int j = 0; j < CPT / 2; ++j) { + const float4 w4 = wsrc[j]; + const unsigned uu0 = __float_as_uint(w4.x); + const unsigned uu1 = __float_as_uint(w4.y); + const unsigned uu2 = __float_as_uint(w4.z); + const unsigned uu3 = __float_as_uint(w4.w); + wv[2 * j][0] = __uint_as_float(uu0 << 16); + wv[2 * j][1] = __uint_as_float(uu0 & 0xffff0000u); + wv[2 * j][2] = __uint_as_float(uu1 << 16); + wv[2 * j][3] = __uint_as_float(uu1 & 0xffff0000u); + wv[2 * j + 1][0] = __uint_as_float(uu2 << 16); + wv[2 * j + 1][1] = __uint_as_float(uu2 & 0xffff0000u); + wv[2 * j + 1][2] = __uint_as_float(uu3 << 16); + wv[2 * j + 1][3] = __uint_as_float(uu3 & 0xffff0000u); + } + } + float acc[CPT]; + { + const uint2 b2 = *reinterpret_cast(L.conv_b + c); + acc[0] = __uint_as_float(b2.x << 16); + acc[1] = __uint_as_float(b2.x & 0xffff0000u); + acc[2] = __uint_as_float(b2.y << 16); + acc[3] = __uint_as_float(b2.y & 0xffff0000u); + } + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v[CPT]; + if (src < 0) { + #pragma unroll + for (int q = 0; q < CPT; ++q) + v[q] = __bfloat162float( + state[(size_t)q * (MB_CONV_K - 1) + pos + k]); + } else { + const uint2 p = *reinterpret_cast( + zxbcdt + (size_t)src * MB_IN_PROJ + MB_INNER + c); + v[0] = __uint_as_float(p.x << 16); + v[1] = __uint_as_float(p.x & 0xffff0000u); + v[2] = __uint_as_float(p.y << 16); + v[3] = __uint_as_float(p.y & 0xffff0000u); + } + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] += wv[q][k] * v[q]; + } + uint2 out; + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] = acc[q] / (1.f + __expf(-acc[q])); + out.x = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[0])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[1])) << 16); + out.y = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[2])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[3])) << 16); + *reinterpret_cast(xbc_act + (size_t)t * MB_CONV_DIM + c) = out; + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_dtda_prefill( + const MambaLayer &L, const bf16 *zxbcdt, float *dtda, int n_tokens) { + const int total = n_tokens * MB_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_HEADS; + const int head = idx - t * MB_HEADS; + const float dt = mamba_softplus( + __bfloat162float(zxbcdt[(size_t)t * MB_IN_PROJ + MB_INNER + + MB_CONV_DIM + head]) + L.dt_bias[head]); + dtda[(size_t)idx * 2] = dt; + dtda[(size_t)idx * 2 + 1] = __expf(L.A[head] * dt); + } +} +template +__device__ __forceinline__ void mamba_ssm_prefill_impl( + const MambaLayer &L, const bf16 *xbc_act, const float *dtda, bf16 *gated, + int n_pos, int n_seq, int slot, const int *seq_slot) { + constexpr int NS = MB_STATE / 32; + constexpr int PG = MB_HEAD_DIM / PW; + constexpr bool RS = (PW % 4) == 0; + constexpr int H1 = RS ? PW / 2 : 1; + constexpr int H2 = RS ? PW / 4 : 1; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const bool keep_hi = (lane & 16) != 0; + const bool keep_hi2 = (lane & 8) != 0; + const int warps_total = gridDim.x * WARPS; + const int units = n_seq * MB_HEADS * PG; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int pg = unit % PG; + const int rest = unit / PG; + const int head = rest % MB_HEADS; + const int seq = rest / MB_HEADS; + const int p0 = pg * PW; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Dh = L.D[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p0 * MB_STATE; + float sreg[PW][NS]; + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + sreg[q][j] = st[(size_t)q * MB_STATE + lane + j * 32]; + } + auto fetch = [&](int p, float (&oB)[NS], float (&oC)[NS], + float (&oX)[PW], float2 &od) { + const size_t tt = (size_t)p * n_seq + seq; + const bf16 *xa2 = xbc_act + tt * MB_CONV_DIM; + const bf16 *Bs = xa2 + MB_INNER + group * MB_STATE; + const bf16 *Cs = xa2 + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + od = *reinterpret_cast( + dtda + (tt * MB_HEADS + head) * 2); + #pragma unroll + for (int j = 0; j < NS; ++j) { + oB[j] = __bfloat162float(Bs[lane + j * 32]); + oC[j] = __bfloat162float(Cs[lane + j * 32]); + } + if constexpr (RS) { + const uint2 *xp = reinterpret_cast( + xa2 + head * MB_HEAD_DIM + p0); + #pragma unroll + for (int q = 0; q < PW; q += 4) { + const uint2 w = xp[q >> 2]; + oX[q] = __uint_as_float(w.x << 16); + oX[q + 1] = __uint_as_float(w.x & 0xffff0000u); + oX[q + 2] = __uint_as_float(w.y << 16); + oX[q + 3] = __uint_as_float(w.y & 0xffff0000u); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) + oX[q] = __bfloat162float(xa2[head * MB_HEAD_DIM + p0 + q]); + } + }; + float pB[NS], pC[NS], pX[PW]; + float2 pd; + fetch(0, pB, pC, pX, pd); + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const float dt = pd.x; + const float dA = pd.y; + float Bv[NS], Cv[NS]; + #pragma unroll + for (int j = 0; j < NS; ++j) { + Bv[j] = pB[j] * dt; + Cv[j] = pC[j]; + } + float acc[PW], xvs[PW]; + #pragma unroll + for (int q = 0; q < PW; ++q) xvs[q] = pX[q]; + if (pos + 1 < n_pos) fetch(pos + 1, pB, pC, pX, pd); + #pragma unroll + for (int q = 0; q < PW; ++q) { + const float xv = xvs[q]; + float a = 0.f; + #pragma unroll + for (int j = 0; j < NS; ++j) { + sreg[q][j] = sreg[q][j] * dA + Bv[j] * xv; + a += sreg[q][j] * Cv[j]; + } + acc[q] = a; + } + if constexpr (RS) { + float m1[H1], t1[H1], x1[H1]; + #pragma unroll + for (int q = 0; q < H1; ++q) { + m1[q] = keep_hi ? acc[q + H1] : acc[q]; + t1[q] = keep_hi ? acc[q] : acc[q + H1]; + x1[q] = keep_hi ? xvs[q + H1] : xvs[q]; + } + #pragma unroll + for (int q = 0; q < H1; ++q) + m1[q] += __shfl_xor_sync(0xffffffff, t1[q], 16); + float m2[H2], t2[H2], x2[H2]; + #pragma unroll + for (int q = 0; q < H2; ++q) { + m2[q] = keep_hi2 ? m1[q + H2] : m1[q]; + t2[q] = keep_hi2 ? m1[q] : m1[q + H2]; + x2[q] = keep_hi2 ? x1[q + H2] : x1[q]; + } + #pragma unroll + for (int q = 0; q < H2; ++q) + m2[q] += __shfl_xor_sync(0xffffffff, t2[q], 8); + #pragma unroll + for (int q = 0; q < H2; ++q) { + #pragma unroll + for (int off = 4; off > 0; off >>= 1) + m2[q] += __shfl_xor_sync(0xffffffff, m2[q], off); + } + if ((lane & 7) == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + + p0 + (keep_hi ? H1 : 0) + (keep_hi2 ? H2 : 0); + #pragma unroll + for (int q = 0; q < H2; ++q) + dst[q] = __float2bfloat16(m2[q] + x2[q] * Dh); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc[q] += __shfl_xor_sync(0xffffffff, acc[q], off); + } + if (lane == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + p0; + #pragma unroll + for (int q = 0; q < PW; ++q) + dst[q] = __float2bfloat16(acc[q] + xvs[q] * Dh); + } + } + } + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + st[(size_t)q * MB_STATE + lane + j * 32] = sreg[q][j]; + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, + const float *dtda, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + (void)zxbcdt; + if (n_seq > 1) + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens / n_seq, n_seq, slot, seq_slot); + else + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens, 1, slot, nullptr); +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16( + held[j] * inv * __bfloat162float(L.gnorm_w[idx])); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = BATCH * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / BB_ATT_HEADS; + const int head = pair % BB_ATT_HEADS; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const int kv = head / kv_per_q; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *q = base + head * BB_HEAD_DIM; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (head % kv_per_q == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + for (int tpos = 0; tpos < T; ++tpos) { +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new + : kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#endif + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m, s); + const float rescale = __expf(m - m_new); + const float w = __expf(s - m_new); + denom = denom * rescale + w; +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *vt = kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#else + const bf16 *vt = cur_token + ? v_new + : kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + w * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + constexpr int CHUNKS = BB_HEAD_DIM / 8; + const int total = n_tokens * BB_KV_HEADS * CHUNKS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int c = idx % CHUNKS; + const int rest = idx / CHUNKS; + const int t = rest / BB_KV_HEADS; + const int kv = rest - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM + c * 8; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM + c * 8; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv) + c * 8; + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv) + c * 8; + *reinterpret_cast(kc) = *reinterpret_cast(k_new); + *reinterpret_cast(vc) = *reinterpret_cast(v_new); + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float qv[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qv[j] = __bfloat162float(q[lane + j * 32]); + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + const int *const krow = + L.block_table + ((size_t)seq * 2 + 0) * L.max_blocks_per_seq; + const int *const vrow = + L.block_table + ((size_t)seq * 2 + 1) * L.max_blocks_per_seq; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + const int head_bias = kv * L.tokens_per_block * BB_HEAD_DIM; + int page = -1; + const bf16 *kpage = nullptr, *vpage = nullptr; + float kn[PER_LANE], vn[PER_LANE]; + { + page = 0; + kpage = L.pool_base + (size_t)(krow[0] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + (size_t)(vrow[0] & 0x7FFFFFFF) * elems_per_block; + const int at = head_bias + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + for (int s = 0; s <= pos; ++s) { + float kc[PER_LANE], vc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { kc[j] = kn[j]; vc[j] = vn[j]; } + if (s < pos) { + const int np = (s + 1) / L.tokens_per_block; + if (np != page) { + page = np; + kpage = L.pool_base + + (size_t)(krow[np] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + + (size_t)(vrow[np] & 0x7FFFFFFF) * elems_per_block; + } + const int at = head_bias + + ((s + 1) % L.tokens_per_block) * BB_HEAD_DIM + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) dot += qv[j] * kc[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[j] = acc[j] * rescale + wt * vc[j]; + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH, unsigned int *tile_masks = nullptr, + bf16 *gath = nullptr, int *gidx = nullptr, int *ecnt = nullptr, + int estride = 0, bf16 *row_stage = nullptr) { + __shared__ float score[BB_EXPERTS]; + __shared__ int s_pick[BB_TOP_K], s_slot[BB_TOP_K]; + constexpr int CH = D_MODEL / 8; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + { + const float4 *gsrc = + reinterpret_cast(normed + (size_t)row * D_MODEL); + float4 *gdst = reinterpret_cast(row_stage); + for (int i = threadIdx.x; i < CH; i += THREADS) gdst[i] = gsrc[i]; + } + if ((int)threadIdx.x < BB_EXPERTS) + moe_w[(size_t)row * BB_EXPERTS + threadIdx.x] = 0.f; + __syncthreads(); + for (int e = warp; e < BB_EXPERTS; e += WARPS) { + const bf16 *wr = gate_w + (size_t)e * D_MODEL; + float acc = 0.f; + for (int i = lane; i < D_MODEL; i += 32) + acc += __bfloat162float(row_stage[i]) * + __bfloat162float(wr[i]); + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) score[e] = 1.f / (1.f + __expf(-acc)); + } + __syncthreads(); + if (threadIdx.x == 0) { + float remaining[BB_EXPERTS]; + int pick[BB_TOP_K]; + for (int e = 0; e < BB_EXPERTS; ++e) + remaining[e] = score[e] + __bfloat162float(gate_bias[e]); + float total = 0.f; + for (int k = 0; k < BB_TOP_K; ++k) { + int best = 0; + for (int e = 1; e < BB_EXPERTS; ++e) + if (remaining[e] > remaining[best]) best = e; + pick[k] = best; + total += score[best]; + remaining[best] = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + unsigned int bits = 0u; + for (int k = 0; k < BB_TOP_K; ++k) { + moe_w[row * BB_EXPERTS + pick[k]] = score[pick[k]] * norm; + bits |= 1u << pick[k]; + s_pick[k] = pick[k]; + } + if (tile_masks) atomicOr(&tile_masks[row / BATCH_PAD], bits); + } + __syncthreads(); + if (gath) { + if ((int)threadIdx.x < BB_TOP_K) { + const int k = threadIdx.x; + const int slot = atomicAdd(&ecnt[s_pick[k]], 1); + s_slot[k] = slot; + gidx[(size_t)s_pick[k] * estride + slot] = row; + } + __syncthreads(); + const float4 *src = reinterpret_cast(row_stage); + for (int k = 0; k < BB_TOP_K; ++k) { + float4 *dst = reinterpret_cast( + gath + ((size_t)s_pick[k] * estride + s_slot[k]) * D_MODEL); + for (int i = threadIdx.x; i < CH; i += THREADS) dst[i] = src[i]; + } + } + __syncthreads(); + } +} +__device__ __forceinline__ unsigned moe_row_mask(const float *routes, int row) { + unsigned m = 0u; + const float *rw = routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) if (rw[e] != 0.f) m |= 1u << e; + return m; +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx, + unsigned mask) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + for (unsigned m = mask; m; m &= m - 1u) + value = __fadd_rn(value, parts[base + (size_t)(__ffs(m) - 1) * D_MODEL]); + return __float2bfloat16( + __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL])); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + uint64_t *keys, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int n = threadIdx.x; n < SORT_N; n += THREADS) { + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + keys[n] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + __syncthreads(); + for (int k = 2; k <= SORT_N; k <<= 1) { + if (k >= 64) __syncthreads(); + for (int j = k >> 1; j > 0; j >>= 1) { + for (int i = threadIdx.x; i < SORT_N; i += THREADS) { + const int partner = i ^ j; + if (partner > i) { + const bool descending = ((i & k) == 0); + if ((keys[i] < keys[partner]) == descending) { + const uint64_t t = keys[i]; + keys[i] = keys[partner]; + keys[partner] = t; + } + } + } + if (j >= 32) __syncthreads(); + else __syncwarp(); + } + } + __syncthreads(); + if (threadIdx.x == 0) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float v; + int n; + if (!eos_used && eos_v > av) { + v = eos_v; + n = AUDIO_EOS; + eos_used = true; + } else { + v = av; + n = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = + v + gumbel[((size_t)row * N_POS + pos) * TOP_K + j]; + if (noised > best_noised) { best_noised = noised; best_code = n; } + } + out_codes[row] = best_code; + } + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void prefill_expert_masks(const float *routes, unsigned int *masks, int rows) { + const int lane=threadIdx.x&31, warp=threadIdx.x>>5; + const int nt=(rows+BATCH_PAD-1)/BATCH_PAD; + for (int t=blockIdx.x*WARPS+warp;t(gath + g * D_MODEL + c * 8); + if (local < ecnt[e]) + *dst = *reinterpret_cast( + normed + (size_t)gidx[g] * D_MODEL + c * 8); + else + *dst = make_float4(0.f, 0.f, 0.f, 0.f); + } +} +__device__ __forceinline__ int gemm_eff(int tiles, int reps) { + if (NCONS <= 1) return 1; + const int g = (int)gridDim.x; + const int rounds1 = (tiles * reps + g - 1) / g; + if constexpr (NCONS == 3) { + const int max_eff = reps <= BB_EXPERTS ? 2 : 3; + const int rounds2 = (((tiles + 1) / 2) * reps + g - 1) / g; + if (max_eff == 2) return rounds2 < rounds1 ? 2 : 1; + const int rounds3 = (((tiles + 2) / 3) * reps + g - 1) / g; + if (rounds3 < rounds2 && rounds3 < rounds1) return 3; + return rounds2 < rounds1 ? 2 : 1; + } else { + const int roundsN = (((tiles + NCONS - 1) / NCONS) * reps + g - 1) / g; + return roundsN < rounds1 ? NCONS : 1; + } +} +__device__ __forceinline__ int gemm_groups(int tiles, int reps) { + const int eff = gemm_eff(tiles, reps); + return (tiles + eff - 1) / eff; +} +__device__ __forceinline__ int gemm_consumer_id() { + const int g = warpgroup::groupid(); + if (g == 1) return -1; + const int c = g == 0 ? 0 : g - 1; + return c < NCONS ? c : -1; +} +__device__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][BATCH_PAD]; + __shared__ int moe_gidx_stage[NCONS][BATCH_PAD]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + a_s[s], a, {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + a_s[s], a, {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / BATCH_PAD) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, BATCH_PAD> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], a_s[s]); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * BATCH_PAD; + if (epi == EPI_BIAS) { + warpgroup::store(o_sf, accum); + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + dest[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(o_sf[int2{col, row}] + + __bfloat162float(bias[base + col])); + } + warpgroup::sync(1 + cons); + continue; + } + warpgroup::store(o_s[cons], accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_s[cons][int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__device__ void gemm_narrow( + const weight_gl &w, const act_gl_n &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = PAD_N, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][PAD_N]; + __shared__ int moe_gidx_stage[NCONS][PAD_N]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + PAD_N - 1) / PAD_N; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + PAD_N - 1) / PAD_N; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / PAD_N) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, PAD_N> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], + *reinterpret_cast(&a_s[s])); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * PAD_N; + out_tile_n &o_sn = *reinterpret_cast(&o_s[cons]); + warpgroup::store(o_sn, accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_sn[int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&w_s)[STAGES * NCONS] = + alloc.allocate(); + act_tile (&a_s)[STAGES] = alloc.allocate(); + out_tile (&o_s)[NCONS] = alloc.allocate(); +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 + static_assert(NCONS * sizeof(out_tile) >= sizeof(out_tile_f), + "aliased head staging must fit the output tiles"); + out_tile_f &o_sf = *reinterpret_cast(&o_s[0]); +#else + out_tile_f (&o_sf) = alloc.allocate(); +#endif + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_WARPS + 2]; +#if MK2_PREFILL_ONLY + bf16 *const route_stage = reinterpret_cast(&o_s[0]); + static_assert(sizeof(o_s) >= D_MODEL * sizeof(bf16), + "router row stage must fit the output tiles"); +#else + __shared__ __align__(16) uint64_t sort_keys[SORT_N]; + bf16 *const route_stage = reinterpret_cast(&sort_keys[0]); + static_assert(sizeof(sort_keys) >= D_MODEL * sizeof(bf16), + "router row stage must fit the sampler workspace"); +#endif + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS * NCONS, 0); + for (int c = 0; c < CONSUMER_WARPS * NCONS; ++c) arrive(drained[s]); + } + } + int gemm_iter = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map); + grid_barrier(g.barrier); + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map); + grid_barrier(g.barrier); + mamba_gated_norm(L, bb->zxbcdt, bb->gated); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, nullptr, 1, BATCH, + bb->moe_w, BB_MOE_INTER / OUT_ATOM); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + 1, BATCH, bb->moe_w, D_MODEL / OUT_ATOM); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_deterministic(bb->moe_out, idx, + (1u << BB_EXPERTS) - 1u); + grid_barrier(g.barrier); + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int original_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + bool identical_cohort = false; + if constexpr (BATCH > 1) { + if (bb->native_pf_cohort_check && bb->native_prefill_reset && + bb->pf_n_seq == BATCH && !bb->pf_fuse_decode && + !bb->pf_layer_out && !bb->moe_w_inject && !bb->moe_w_audit) { + if (blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[0] = 0; + grid_barrier(g.barrier); + const auto *bits = reinterpret_cast(bb->prompt); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < original_total * D_MODEL; idx += gridDim.x * THREADS) { + const int canonical = (idx / (BATCH * D_MODEL)) * BATCH * D_MODEL + idx % D_MODEL; + if (bits[idx] != bits[canonical]) atomicOr(bb->native_pf_cohort_check, 1u); + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + if (bb->pf_seq_slot[row] != row || bb->slot_map[row] != row) + atomicOr(bb->native_pf_cohort_check, 1u); + grid_barrier(g.barrier); + identical_cohort = bb->native_pf_cohort_check[0] == 0; + } + } + if (!identical_cohort) { + if (bb->native_pf_cohort_check && blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[1] = 0; + grid_barrier(g.barrier); + return; + } + const int n_total = original_total / BATCH; + if (bb->native_pf_cohort_check && blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[1] = identical_cohort ? 1 : BATCH; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int reset_slots = identical_cohort ? 1 : BATCH; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * reset_slots * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % reset_slots; + const MambaLayer &L = bb->mamba[row_layer / reset_slots]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = original_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + const int src = identical_cohort ? (idx / D_MODEL) * BATCH * D_MODEL + idx % D_MODEL : idx; + bb->pf_h[idx] = bb->prompt[src]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = identical_cohort ? 1 : (bb->pf_n_seq > 0 ? bb->pf_n_seq : 1); + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T, + (!bb->pf_layer_out && layer > 0 && + bb->moe_of_layer[layer - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + if (bb->native_pf_expert_masks) + for (int i = blockIdx.x * THREADS + threadIdx.x; i < n_tok; + i += gridDim.x * THREADS) + bb->native_pf_expert_masks[i] = 0u; + if (bb->pf_moe_ecnt) + for (int e = blockIdx.x * THREADS + threadIdx.x; e < BB_EXPERTS; + e += gridDim.x * THREADS) + bb->pf_moe_ecnt[e] = 0; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + for(size_t j=blockIdx.x*THREADS+threadIdx.x;jpf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + mamba_dtda_prefill(L, bb->pf_zxbcdt, bb->pf_dtda, T); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_dtda, + bb->pf_gated, T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + for(size_t j=blockIdx.x*THREADS+threadIdx.x;jpf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, + T, bb->native_pf_expert_masks, bb->pf_moe_gath, + bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, route_stage); + } + grid_barrier(g.barrier); + int max_ecnt = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + max_ecnt = max(max_ecnt, bb->pf_moe_ecnt[e]); + const bool narrow = BATCH_PAD > PAD_N && max_ecnt <= PAD_N; + const int etile = narrow ? PAD_N : BATCH_PAD; + int nt_sum = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_sum += (bb->pf_moe_ecnt[e] + etile - 1) / etile; + const int up_steps = + gemm_groups(BB_MOE_INTER / OUT_ATOM, nt_sum) * + nt_sum * (D_MODEL / KCHUNK); + const int shup_steps = + gemm_groups(BB_SHARED_INTER / OUT_ATOM, n_tok) * + n_tok * (D_MODEL / KCHUNK); + const int down_steps = + gemm_groups(D_MODEL / OUT_ATOM, nt_sum) * + nt_sum * (BB_MOE_INTER / KCHUNK); + const int shdown_steps = + gemm_groups(D_MODEL / OUT_ATOM, n_tok) * + n_tok * (BB_SHARED_INTER / KCHUNK); + if (narrow) + gemm_narrow(L.exp_up, bb->pf_moe_gath_vn, + BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, PAD_N, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + else + gemm(L.exp_up, bb->pf_moe_gath_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, BATCH_PAD, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + up_steps, up_steps + shup_steps); + grid_barrier(g.barrier); + if (narrow) + gemm_narrow(L.exp_down, bb->pf_moe_hc_vn, + BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, PAD_N, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + else + gemm(L.exp_down, bb->pf_moe_hc_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + down_steps, down_steps + shdown_steps); + grid_barrier(g.barrier); + if (bb->pf_layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * D_MODEL; idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = + bb_moe_fold_deterministic(bb->pf_moe_out, idx, + moe_row_mask(bb->pf_moe_w, idx / D_MODEL)); + grid_barrier(g.barrier); + } + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, + identical_cohort ? bb->pf_normed : bb->pf_final, ln_scratch, + n_total, + (!bb->pf_layer_out && n_backbone_layers > 0 && + bb->moe_of_layer[n_backbone_layers - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + grid_barrier(g.barrier); + if (identical_cohort) { + const int worker = blockIdx.x * THREADS + threadIdx.x; + const int workers = gridDim.x * THREADS; + for (int i = worker; i < original_total * D_MODEL; i += workers) + bb->pf_final[i] = bb->pf_normed[(i / (BATCH * D_MODEL)) * D_MODEL + i % D_MODEL]; + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = worker; i < 14 * (BATCH - 1) * CV; i += workers) { + const int cell = i % CV, pair = i / CV; + const int row = pair % (BATCH - 1) + 1; + const MambaLayer &L = bb->mamba[pair / (BATCH - 1)]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = L.conv_state[cell]; + } + const int kv_columns = BB_KV_HEADS * BB_HEAD_DIM; + for (size_t i = worker; i < size_t(3) * 2 * (BATCH - 1) * n_total * kv_columns; i += workers) { + size_t q = i / kv_columns; + const int col = i % kv_columns; + const int pos = q % n_total; q /= n_total; + const int row = q % (BATCH - 1) + 1; q /= BATCH - 1; + const int kind = q % 2, layer = q / 2; + const AttnLayer &L = bb->attn[layer]; + const int offset = kv_offset(L, pos, col / BB_HEAD_DIM) + col % BB_HEAD_DIM; + kv_slot(L, row, kind, pos)[offset] = kv_slot(L, 0, kind, pos)[offset]; + } + grid_barrier(g.barrier); + } + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } +#if MK2_PREFILL_ONLY + return; +#else + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, bb->phoneme_b); + grid_barrier(g.barrier); + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + if (threadIdx.x == 0) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = 0; n < BB_PHONEME; ++n) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + bb->phoneme_out[row] = arg; + } + __syncthreads(); + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + for (int i = lane; i < D_HEAD; i += 32) { + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + for (int t = 0; t <= pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + for (int i = lane; i < D_HEAD; i += 32) { + float acc = 0.f; + for (int t = 0; t <= pos; ++t) + acc += sc[t] * __bfloat162float( + g.vcache[cache_base + t * D_HEAD + i]); + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + i] = + __float2bfloat16(acc); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, g.sampled_pos, + sort_keys, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { +#ifdef ABLATE_SAMPLE + __syncthreads(); +#endif + } else { + grid_barrier(g.barrier); + } + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x < BATCH && blockIdx.x == 0) { + const int code = g.sampled_pos[threadIdx.x]; + g.sampled[threadIdx.x * N_POS + pos] = code; + if (g.native_codes) + g.native_codes[threadIdx.x * N_POS + pos] = (long long)code; + } + grid_barrier(g.barrier); + } + if (g.native_seq_len && blockIdx.x == 0 && threadIdx.x < BATCH) + g.native_seq_len[threadIdx.x] += 1; + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +#endif +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 2 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK), + 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = MK2_PREFILL_ONLY + ? out->sms * out->blocks_per_sm + : std::min(out->sms * out->blocks_per_sm, D_FFN / OUT_ATOM); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + 0, + nullptr, + nullptr, + nullptr, nullptr, nullptr, nullptr, nullptr, 0, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "dedup_native_prefill_impl.inc" diff --git a/deploy/nemotron_native/kernels/b4096/regions/prefill/native/prefill/mk2_abi.cuh b/deploy/nemotron_native/kernels/b4096/regions/prefill/native/prefill/mk2_abi.cuh new file mode 100644 index 0000000..95a253b --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/prefill/native/prefill/mk2_abi.cuh @@ -0,0 +1,82 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} diff --git a/deploy/nemotron_native/kernels/b4096/regions/prefill/native/prefill/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/b4096/regions/prefill/native/prefill/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/prefill/native/prefill/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/b4096/regions/prefill/native/prefill/prefill.cu b/deploy/nemotron_native/kernels/b4096/regions/prefill/native/prefill/prefill.cu new file mode 100644 index 0000000..271ac63 --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/prefill/native/prefill/prefill.cu @@ -0,0 +1,18 @@ +#define MK2_BATCH 4096 +#define MK2_BATCH_TILE 64 +#define MK2_PREFILL_ONLY 1 +#define STAGES 3 +#define KCHUNK 128 +#define MK2_GEMM_CONSUMERS 3 +#define MK2_THREADS 512 +#define MK2_NO_MAIN 1 +#define MK2_DIRECT_INPUT 1 +#define MK2_B1_LOCAL_BARRIERS 1 +#define MK2_SKIP_HIDDEN_SNAPSHOT 1 +#define MK2_ACTIVE_EXPERTS 1 +#define MK2_SSM_SPLITS 1 +#define PF_SSM_PW_MULTI 12 +#define PF_ACTIVE_TOKEN_EXPERTS 1 +#include +#include +#include "dedup_predictor_mk2.cuh" diff --git a/deploy/nemotron_native/kernels/b4096/regions/prefill/prefill.cu b/deploy/nemotron_native/kernels/b4096/regions/prefill/prefill.cu new file mode 100644 index 0000000..dce5f15 --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/prefill/prefill.cu @@ -0,0 +1,20 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace b4096_prefill { +#include "native/prefill/prefill.cu" +} diff --git a/deploy/nemotron_native/kernels/b4096/regions/stage0/native/stage0/mk2_abi.cuh b/deploy/nemotron_native/kernels/b4096/regions/stage0/native/stage0/mk2_abi.cuh new file mode 100644 index 0000000..bf1e415 --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/stage0/native/stage0/mk2_abi.cuh @@ -0,0 +1,85 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_set_native_predictor_active_mask(MK2Predictor *p, const unsigned char *mask); +int mk2_native_startup_skip_abi_version(); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *, int grid); diff --git a/deploy/nemotron_native/kernels/b4096/regions/stage0/native/stage0/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/b4096/regions/stage0/native/stage0/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/stage0/native/stage0/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/b4096/regions/stage0/native/stage0/native_prefill_impl.inc b/deploy/nemotron_native/kernels/b4096/regions/stage0/native/stage0/native_prefill_impl.inc new file mode 100644 index 0000000..a254308 --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/stage0/native/stage0/native_prefill_impl.inc @@ -0,0 +1,198 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 16 || BATCH == 32 || BATCH == 64)) + throw std::runtime_error("native prefill requires matching batch 1/16/32/64"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (padded > INT_MAX / ((BB_EXPERTS + 1) * D_MODEL)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + float *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(float))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(float))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} diff --git a/deploy/nemotron_native/kernels/b4096/regions/stage0/native/stage0/predictor_mk2.cu b/deploy/nemotron_native/kernels/b4096/regions/stage0/native/stage0/predictor_mk2.cu new file mode 100644 index 0000000..c2b54f3 --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/stage0/native/stage0/predictor_mk2.cu @@ -0,0 +1,6107 @@ +#include "../../../../../inflight/progress.cuh" +#if MK2_INFLIGHT +#include +#endif +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#define MK2_FUSE_CONV (MK2_SSM_SPLITS == 1) +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +constexpr int BATCH_PAD = 64; +constexpr int DECODE_TILES = BATCH / BATCH_PAD; +constexpr int DECODE_ROWS = BATCH_PAD; +static_assert(BATCH == 512 || BATCH == 1024 || BATCH == 2048 || BATCH == 4096); +static_assert(BATCH % BATCH_PAD == 0); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +#ifndef NMULT +#define NMULT 1 +#endif +#ifndef SSM_PREFETCH_MAX +#define SSM_PREFETCH_MAX 8 +#endif +constexpr int NM_MAX = (NMULT > 1 && DECODE_TILES % NMULT == 0) ? NMULT : 1; +#ifndef MK2_OUT_MULT +#define MK2_OUT_MULT 1 +#endif +constexpr int MT = MK2_OUT_MULT; +static_assert(MT >= 1 && MT <= NM_MAX, "an item holds at most NM_MAX accumulators"); +constexpr int SLOTS = 1 + NM_MAX; +static_assert(MT + 1 <= SLOTS); +constexpr int THREADS = 256; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +constexpr int NORM_SCRATCH_STRIDE = 2 * WARPS + 2; +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_ACC_F32 }; +#define EPI_ROW(idx) ((OUT_ATOM == 64 && BATCH_PAD == 64) \ + ? (((idx) & 15) | ((((idx) >> 6) & 3) << 4)) : ((idx) / (OUT_ATOM / 8))) +#define EPI_OCT(idx) ((OUT_ATOM == 64 && BATCH_PAD == 64) \ + ? ((((idx) >> 4) & 3) | ((((idx) >> 8) & 1) << 2)) : ((idx) % (OUT_ATOM / 8))) +struct alignas(16) bf16_oct { bf16 x0, x1, x2, x3, x4, x5, x6, x7; }; +static_assert(sizeof(bf16_oct) == 16); +#include +constexpr int R9_STATE_MODE=1; +using r9_state_t=bf16; +__device__ __forceinline__ int r9_scale_index(int head,int channel,int n) {return (((head/8)*4+(head%8)/2)*4+channel/6)*128+n;} +__device__ __forceinline__ float r9_load(r9_state_t v,float scale) { +return __bfloat162float(v); +} +__device__ __forceinline__ r9_state_t r9_store(float x,float scale) { +return __float2bfloat16(x); +} +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; + const float *r8_initial_ssm = nullptr; + const unsigned int *r8_initial_ready = nullptr; + const int *r8_seq_len = nullptr; + int r8_initial_positions = 0; + const float *r9_scales = nullptr; + const int *r9_map = nullptr; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + bf16 *r6_packed = nullptr; + int *r6_map = nullptr,*r6_counts = nullptr; + act_gl *r6_views = nullptr; + int *r6_lut = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +#define BARRIER_GEN_WORD 32 +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + constexpr unsigned long long GRID_BLOCKS = 132ull; + auto *wide_state = reinterpret_cast(state); + const unsigned long long arrival = atomicAdd(wide_state, 1ull); + const unsigned int rank = (unsigned int)(arrival % GRID_BLOCKS); + if (rank != GRID_BLOCKS - 1u) { + const unsigned long long target = arrival - rank + GRID_BLOCKS; + volatile unsigned long long *counter = wide_state; + while (*counter < target) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / THREADS; + float wv[PER_THREAD]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) wv[j] = weight[threadIdx.x + j * THREADS]; + float v[PER_THREAD], v_next[PER_THREAD]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + v[j] = __bfloat162float( + src[(size_t)blockIdx.x * D_MODEL + threadIdx.x + j * THREADS]); + for (int row = blockIdx.x, row_iter = 0; row < BATCH; + row += gridDim.x, ++row_iter) { + float *row_scratch = scratch + (row_iter & 1) * NORM_SCRATCH_STRIDE; + const int next_row = row + (int)gridDim.x; + if (next_row < BATCH) { + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + v_next[j] = __bfloat162float( + src[(size_t)next_row * D_MODEL + threadIdx.x + j * THREADS]); + } + float sum = 0.f, sq = 0.f; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) { + row_scratch[warp] = sum; + row_scratch[WARPS + warp] = sq; + } + __syncthreads(); + float s = 0.f, q = 0.f; + #pragma unroll + for (int w = 0; w < WARPS; ++w) { + s += row_scratch[w]; + q += row_scratch[WARPS + w]; + } + const float mean = s / D_MODEL; + const float inv = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * wv[j]); + } + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) v[j] = v_next[j]; + } +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, bf16 *snapshot = nullptr, int snapshot_layer = -1, + float *snapshot_fp32 = nullptr, float *out_fp32 = nullptr, + const float *fold_src = nullptr, const float *fold_routes = nullptr) { + constexpr int PER_THREAD = D_MODEL / THREADS; + bf16 wv[PER_THREAD]; + extern __shared__ int __shm[]; + bf16 *const row_stage = reinterpret_cast(__shm); + constexpr int ROW_SLOT = 2 * D_MODEL; + constexpr int ROW_VECTORS = D_MODEL / 8; + if ((int)blockIdx.x < n_rows && threadIdx.x < ROW_VECTORS) { + const int i = 8 * threadIdx.x; + const size_t at = (size_t)blockIdx.x * D_MODEL + i; + const uint32_t rs = static_cast( + __cvta_generic_to_shared(row_stage + D_MODEL + i)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;" + :: "r"(rs), "l"(r + at) : "memory"); + if (!fold_src) { + const uint32_t hs = static_cast( + __cvta_generic_to_shared(row_stage + i)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;" + :: "r"(hs), "l"(h + at) : "memory"); + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + wv[j] = weight[threadIdx.x + j * THREADS]; + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); + __syncthreads(); + for (int row = blockIdx.x, row_iter = 0; row < n_rows; + row += gridDim.x, ++row_iter) { + bf16 *const current = row_stage + (row_iter & 1) * ROW_SLOT; + const int next_row = row + (int)gridDim.x; + if (next_row < n_rows) { + bf16 *const next = row_stage + ((row_iter + 1) & 1) * ROW_SLOT; + if (threadIdx.x < ROW_VECTORS) { + const int i = 8 * threadIdx.x; + const size_t at = (size_t)next_row * D_MODEL + i; + const uint32_t rs = static_cast( + __cvta_generic_to_shared(next + D_MODEL + i)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;" + :: "r"(rs), "l"(r + at) : "memory"); + if (!fold_src) { + const uint32_t hs = static_cast( + __cvta_generic_to_shared(next + i)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;" + :: "r"(hs), "l"(h + at) : "memory"); + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } + float *row_scratch = scratch + (row_iter & 1) * NORM_SCRATCH_STRIDE; + const size_t base = (size_t)row * D_MODEL; + float fv[PER_THREAD]; + if (fold_src) { + unsigned mask = 0u; + const float *rw = fold_routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) + if (rw[e] != 0.f) mask |= 1u << e; + const float *pe = fold_src + + (size_t)row * (BB_EXPERTS + 1) * D_MODEL + threadIdx.x; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) fv[j] = 0.f; + #pragma unroll + for (int e = 0; e <= BB_EXPERTS; ++e) { + const bool take = e == BB_EXPERTS || ((mask >> e) & 1u); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + fv[j] = __fadd_rn(fv[j], + take ? pe[(size_t)e * D_MODEL + j * THREADS] : 0.f); + } + } + float v[PER_THREAD]; + float sq = 0.f; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + const bf16 hv = fold_src ? __float2bfloat16(fv[j]) : current[i]; + if (fold_src) h[base + i] = hv; + if (snapshot) + snapshot[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = hv; + if (snapshot_fp32) + snapshot_fp32[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = + __bfloat162float(hv); + v[j] = __bfloat162float(hv) + + __bfloat162float(current[D_MODEL + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) row_scratch[warp] = sq; + __syncthreads(); + float s = 0.f; + #pragma unroll + for (int w = 0; w < WARPS; ++w) s += row_scratch[w]; + const float inv = rsqrtf(s / D_MODEL + RMS_EPS); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + const bf16 normed = __float2bfloat16( + __bfloat162float(wv[j]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + out[base + i] = normed; + if (out_fp32) out_fp32[base + i] = __bfloat162float(normed); + } + if (next_row < n_rows) { + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); + __syncthreads(); + } + } +} +__device__ __forceinline__ bf16 mamba_conv_channel( + const MambaLayer &L, bf16 *state_slot, const bf16 *zrow, int c) { + const unsigned ch = (unsigned)c; + const bf16 *w = L.conv_w + ch * (unsigned)MB_CONV_K; + bf16 *state = state_slot + ch * (unsigned)(MB_CONV_K - 1); + const bf16 rawb = zrow[c]; + const float raw = __bfloat162float(rawb); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = rawb; + return __float2bfloat16(acc); +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map MK2_PROGRESS_PARAMETER(progress_mask)) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + MK2_PROGRESS_GUARD(progress_mask, row); + const int c = idx - row * MB_CONV_DIM; + xbc_act[(size_t)row * MB_CONV_DIM + c] = mamba_conv_channel( + L, L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride, + zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER, c); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map MK2_PROGRESS_PARAMETER(progress_mask)) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int sb4 = (lane >> 4) & 1, sb3 = (lane >> 3) & 1; + const int sb2 = (lane >> 2) & 1, sb1 = (lane >> 1) & 1; + const int ssm_own_ch = 6 * sb4 + 3 * sb3 + 2 * sb2 + sb1; + const bool ssm_owner = (2 * sb2 + sb1) < 3 && (lane & 1) == 0; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + constexpr bool FUSE_GNORM = (split == 1); + extern __shared__ int __shm[]; + bf16 (*const gn_stage)[MB_GROUP_SIZE] = + reinterpret_cast(&__shm[0]); + constexpr bool FUSE_CONV = (MK2_FUSE_CONV != 0); + bf16 (*const conv_x)[MB_GROUP_SIZE] = gn_stage + WARPS; + constexpr int SSM_NSTATE = MB_STATE / 32; + constexpr int SSM_PF = channels < SSM_PREFETCH_MAX ? channels : SSM_PREFETCH_MAX; + constexpr int SSM_GROUP_BYTES = SSM_PF * MB_STATE * (int)sizeof(r9_state_t); + constexpr bool SSM_SMEM_OK = R9_STATE_MODE == 1 && channels == 2 * SSM_PF && + SSM_GROUP_BYTES % 512 == 0; + constexpr int SSM_RING = 4; + constexpr int SSM_GROUPS_PER_HEAD = channels / SSM_PF; + constexpr int SSM_PAIRS = MB_HEADS_PER_GROUP * SSM_GROUPS_PER_HEAD; + static_assert(!SSM_SMEM_OK || (SSM_RING & (SSM_RING - 1)) == 0, + "the flat (head, channel group) ring index is masked"); + static_assert(!SSM_SMEM_OK || SSM_PAIRS >= SSM_RING, + "the prologue must be able to fill the ring"); + static_assert(!SSM_SMEM_OK || + 2 * WARPS * MB_GROUP_SIZE * sizeof(bf16) + + WARPS * SSM_RING * SSM_GROUP_BYTES <= STAGES * SLOTS * sizeof(weight_tile)); + constexpr int SSM_F32_RING = 2; + constexpr int SSM_F32_GROUP_BYTES = SSM_PF * MB_STATE * (int)sizeof(float); + constexpr bool SSM_F32_OK = SSM_SMEM_OK && SSM_F32_GROUP_BYTES % 512 == 0 && + (SSM_F32_RING & (SSM_F32_RING - 1)) == 0 && SSM_PAIRS >= SSM_F32_RING && + SSM_F32_RING * SSM_F32_GROUP_BYTES <= SSM_RING * SSM_GROUP_BYTES; + char *const ssm_stage = reinterpret_cast(conv_x + WARPS); + const int pairs = BATCH * MB_GROUPS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_GROUPS * split); + MK2_PROGRESS_GUARD(progress_mask, row); + const int group = (pair / split) % MB_GROUPS; + const int first_channel = (pair % split) * channels; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const bool r8_initial = L.r8_initial_ssm && L.r8_initial_ready && + *L.r8_initial_ready == 1 && L.r8_seq_len[row] == L.r8_initial_positions; + r9_state_t *st = reinterpret_cast(L.ssm_state) + + (size_t)slot_map[row] * L.ssm_slot_stride; + char *const warp_ssm = + ssm_stage + (size_t)warp * SSM_RING * SSM_GROUP_BYTES; + auto issue_ssm = [&](int head_in_group, int channel0, int buffer) { + const int head = group * MB_HEADS_PER_GROUP + head_in_group; + const r9_state_t *src = st + + (size_t)(head * MB_HEAD_DIM + channel0) * MB_STATE; + const uint32_t dst = (uint32_t)__cvta_generic_to_shared( + warp_ssm + (size_t)buffer * SSM_GROUP_BYTES + lane * 16); + #pragma unroll + for (int chunk = 0; chunk < SSM_GROUP_BYTES / 512; ++chunk) + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: + "r"(dst + (uint32_t)chunk * 512), + "l"(reinterpret_cast(src) + + (size_t)chunk * 512 + lane * 16) : "memory"); + asm volatile("cp.async.commit_group;\n" ::: "memory"); + }; + auto issue_ssm_f32 = [&](int head_in_group, int channel0, int buffer) { + const int head = group * MB_HEADS_PER_GROUP + head_in_group; + const float *src = L.r8_initial_ssm + + (size_t)(head * MB_HEAD_DIM + channel0) * MB_STATE; + const uint32_t dst = (uint32_t)__cvta_generic_to_shared( + warp_ssm + (size_t)buffer * SSM_F32_GROUP_BYTES + lane * 16); + #pragma unroll + for (int chunk = 0; chunk < SSM_F32_GROUP_BYTES / 512; ++chunk) + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: + "r"(dst + (uint32_t)chunk * 512), + "l"(reinterpret_cast(src) + + (size_t)chunk * 512 + lane * 16) : "memory"); + asm volatile("cp.async.commit_group;\n" ::: "memory"); + }; + if constexpr (SSM_SMEM_OK) { + if (!r8_initial) { + #pragma unroll + for (int k = 0; k < SSM_RING; ++k) + issue_ssm(k / SSM_GROUPS_PER_HEAD, + (k % SSM_GROUPS_PER_HEAD) * SSM_PF, k); + } else if constexpr (SSM_F32_OK) { + #pragma unroll + for (int k = 0; k < SSM_F32_RING; ++k) + issue_ssm_f32(k / SSM_GROUPS_PER_HEAD, + (k % SSM_GROUPS_PER_HEAD) * SSM_PF, k); + } + } + float Braw[MB_STATE / 32]; + float Cv[MB_STATE / 32]; + const bf16 *xrow; + if constexpr (FUSE_CONV) { + bf16 *const cslot = + L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride; + const bf16 *const zrow = + zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER; + static_assert(MB_CONV_K == 4, "the state record is three 32-bit words"); + static_assert(MB_GROUP_SIZE % 64 == 0 && MB_STATE % 64 == 0, + "one paired round covers 64 channels of a warp"); + auto wide_lo = [](unsigned int v) { return __uint_as_float(v << 16); }; + auto wide_hi = [](unsigned int v) { return __uint_as_float(v & 0xffff0000u); }; + struct ConvCursor { + const uint2 *w; unsigned int *s; + const unsigned int *z; const unsigned int *b; + }; + auto conv_cursor = [&](int c) { + const unsigned int cu = (unsigned int)c; + return ConvCursor{ + reinterpret_cast( + L.conv_w + cu * (unsigned int)MB_CONV_K), + reinterpret_cast( + cslot + cu * (unsigned int)(MB_CONV_K - 1)), + reinterpret_cast(zrow + cu), + reinterpret_cast(L.conv_b + cu)}; + }; + auto conv_two = [&](uint2 w0, uint2 w1, unsigned int s0, + unsigned int s1, unsigned int s2, + unsigned int zp, unsigned int bp, float &a0, + float &a1, unsigned int &n0, unsigned int &n1, + unsigned int &n2) { + a0 = wide_lo(bp); + a0 += wide_lo(w0.x) * wide_lo(s0); + a0 += wide_hi(w0.x) * wide_hi(s0); + a0 += wide_lo(w0.y) * wide_lo(s1); + a0 += wide_hi(w0.y) * wide_lo(zp); + a0 = a0 / (1.f + __expf(-a0)); + a1 = wide_hi(bp); + a1 += wide_lo(w1.x) * wide_hi(s1); + a1 += wide_hi(w1.x) * wide_lo(s2); + a1 += wide_lo(w1.y) * wide_hi(s2); + a1 += wide_hi(w1.y) * wide_hi(zp); + a1 = a1 / (1.f + __expf(-a1)); + n0 = (s0 >> 16) | (s1 << 16); + n1 = (zp & 0xffffu) | (s2 << 16); + n2 = (s2 >> 16) | (zp & 0xffff0000u); + }; + struct ConvPairIn { + uint2 w0, w1; + unsigned int s0, s1, s2, zp, bp; + }; + auto conv_pair_fetch = [](const ConvCursor &q, int r) { + const uint2 *const wp = q.w + 64 * r; + const unsigned int *const sp = q.s + 96 * r; + return ConvPairIn{wp[0], wp[1], sp[0], sp[1], sp[2], + q.z[32 * r], q.b[32 * r]}; + }; + auto conv_pair_apply = [&](unsigned int *const sp, + const ConvPairIn &v, float &a0, + float &a1) { + unsigned int n0, n1, n2; + conv_two(v.w0, v.w1, v.s0, v.s1, v.s2, v.zp, v.bp, a0, a1, + n0, n1, n2); + sp[0] = n0; + sp[1] = n1; + sp[2] = n2; + }; + constexpr int QUAD_SPAN = 4 * 32; + static_assert(MB_INNER % 4 == 0 && MB_STATE % 4 == 0 && + MB_GROUP_SIZE % 4 == 0, + "every quad base is a four-channel boundary"); + static_assert(MB_STATE == QUAD_SPAN && + MB_GROUP_SIZE == QUAD_SPAN + 2 * 32, + "one quad round covers B and C; x adds a pair round"); + struct ConvQuad { + const uint4 *w; uint2 *s; const uint2 *z; const uint2 *b; + }; + auto conv_quad_cursor = [&](int c) { + const unsigned int cu = (unsigned int)c; + return ConvQuad{ + reinterpret_cast( + L.conv_w + cu * (unsigned int)MB_CONV_K), + reinterpret_cast( + cslot + cu * (unsigned int)(MB_CONV_K - 1)), + reinterpret_cast(zrow + cu), + reinterpret_cast(L.conv_b + cu)}; + }; + struct ConvQuadIn { + uint4 w01, w23; + uint2 sa, sb, sc, zp, bp; + }; + auto conv_quad_fetch = [](const ConvQuad &q) { + const uint2 *const sp = q.s; + return ConvQuadIn{q.w[0], q.w[1], sp[0], sp[1], sp[2], + q.z[0], q.b[0]}; + }; + auto conv_quad_apply = [&](uint2 *const sp, const ConvQuadIn &v, + float &a0, float &a1, float &a2, + float &a3) { + unsigned int n0, n1, n2, n3, n4, n5; + conv_two(make_uint2(v.w01.x, v.w01.y), + make_uint2(v.w01.z, v.w01.w), + v.sa.x, v.sa.y, v.sb.x, v.zp.x, v.bp.x, + a0, a1, n0, n1, n2); + conv_two(make_uint2(v.w23.x, v.w23.y), + make_uint2(v.w23.z, v.w23.w), + v.sb.y, v.sc.x, v.sc.y, v.zp.y, v.bp.y, + a2, a3, n3, n4, n5); + sp[0] = make_uint2(n0, n1); + sp[1] = make_uint2(n2, n3); + sp[2] = make_uint2(n4, n5); + }; + auto conv_pack = [](float a, float b) { + return (unsigned int)__bfloat16_as_ushort(__float2bfloat16(a)) | + ((unsigned int)__bfloat16_as_ushort(__float2bfloat16(b)) << 16); + }; + const ConvQuad qb = conv_quad_cursor( + MB_INNER + group * MB_STATE + 4 * lane); + const ConvQuad qc = conv_quad_cursor( + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE + 4 * lane); + static_assert(MB_GROUP_SIZE * (int)sizeof(bf16) % 8 == 0, + "each warp's staging row starts 8-byte aligned"); + const ConvQuad qx = conv_quad_cursor( + group * MB_GROUP_SIZE + 4 * lane); + const ConvCursor qt = conv_cursor( + group * MB_GROUP_SIZE + QUAD_SPAN + 2 * lane); + const ConvQuadIn vb = conv_quad_fetch(qb); + const ConvQuadIn vc = conv_quad_fetch(qc); + const ConvQuadIn vx = conv_quad_fetch(qx); + const ConvPairIn vt = conv_pair_fetch(qt, 0); + uint2 bpk, cpk; + { + float a0, a1, a2, a3; + conv_quad_apply(qb.s, vb, a0, a1, a2, a3); + bpk = make_uint2(conv_pack(a0, a1), conv_pack(a2, a3)); + conv_quad_apply(qc.s, vc, a0, a1, a2, a3); + cpk = make_uint2(conv_pack(a0, a1), conv_pack(a2, a3)); + } + { + float a0, a1, a2, a3; + conv_quad_apply(qx.s, vx, a0, a1, a2, a3); + *reinterpret_cast(&conv_x[warp][4 * lane]) = + make_uint2(conv_pack(a0, a1), conv_pack(a2, a3)); + } + { + float a0, a1; + conv_pair_apply(qt.s, vt, a0, a1); + conv_x[warp][QUAD_SPAN + 2 * lane] = __float2bfloat16(a0); + conv_x[warp][QUAD_SPAN + 2 * lane + 1] = __float2bfloat16(a1); + } + const bool qword = ((lane >> 1) & 1) != 0; + const bool qhalf = (lane & 1) != 0; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int src = (lane >> 2) + 8 * j; + const unsigned int bx = __shfl_sync(0xffffffffu, bpk.x, src); + const unsigned int by = __shfl_sync(0xffffffffu, bpk.y, src); + const unsigned int cx = __shfl_sync(0xffffffffu, cpk.x, src); + const unsigned int cy = __shfl_sync(0xffffffffu, cpk.y, src); + const unsigned int bg = qword ? by : bx; + const unsigned int cg = qword ? cy : cx; + Braw[j] = qhalf ? wide_hi(bg) : wide_lo(bg); + Cv[j] = qhalf ? wide_hi(cg) : wide_lo(cg); + } + __syncwarp(); + xrow = &conv_x[warp][0]; + } else { + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Braw[j] = __bfloat162float(Bsec[lane + j * 32]); + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + xrow = xa + (size_t)group * MB_GROUP_SIZE; + } + const int head_lane = group * MB_HEADS_PER_GROUP + + (lane < MB_HEADS_PER_GROUP ? lane : 0); + const float dt_lane = mamba_softplus( + __bfloat162float(dtrow[head_lane]) + L.dt_bias[head_lane]); + const float dA_lane = __expf(L.A[head_lane] * dt_lane); + const float D_lane = L.D[head_lane]; + #pragma unroll 1 + for (int head_in_group = 0; head_in_group < MB_HEADS_PER_GROUP; + ++head_in_group) { + const int head = group * MB_HEADS_PER_GROUP + head_in_group; + const float dt = __shfl_sync(0xffffffffu, dt_lane, head_in_group); + const float dA = __shfl_sync(0xffffffffu, dA_lane, head_in_group); + const float Dh = __shfl_sync(0xffffffffu, D_lane, head_in_group); + const float *st_in = L.r8_initial_ssm + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float Bv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) + Bv[j] = Braw[j] * dt; + constexpr int NSTATE = SSM_NSTATE; + constexpr int SSM_PREFETCH = SSM_PF; + static_assert(channels % SSM_PREFETCH == 0); + for (int g0 = 0; g0 < channels; g0 += SSM_PREFETCH) { + float sv[SSM_PREFETCH][NSTATE]; + float xvr[SSM_PREFETCH]; + if (r8_initial && SSM_F32_OK) { + asm volatile("cp.async.wait_group 1;\n" ::: "memory"); + __syncwarp(); + const int pair = head_in_group * SSM_GROUPS_PER_HEAD + + g0 / SSM_PREFETCH; + const int buffer = pair & (SSM_F32_RING - 1); + const float *const smf = reinterpret_cast( + warp_ssm + (size_t)buffer * SSM_F32_GROUP_BYTES); + static_assert(SSM_PREFETCH % 2 == 0 && MB_HEAD_DIM % 2 == 0, + "x staging vectorises in element pairs"); + const __nv_bfloat162 *const xq = + reinterpret_cast( + xrow + head_in_group * MB_HEAD_DIM + + first_channel + g0); + #pragma unroll + for (int gq = 0; gq < SSM_PREFETCH / 2; ++gq) { + const float2 xp = __bfloat1622float2(xq[gq]); + xvr[2 * gq] = xp.x; + xvr[2 * gq + 1] = xp.y; + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = smf[(size_t)gi * MB_STATE + lane + j * 32]; + } + __syncwarp(); + const int ahead = pair + SSM_F32_RING; + if (ahead < SSM_PAIRS) + issue_ssm_f32(ahead / SSM_GROUPS_PER_HEAD, + (ahead % SSM_GROUPS_PER_HEAD) * SSM_PF, + buffer); + else + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } else if (r8_initial) { + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const size_t prow = (size_t)(first_channel + g0 + gi) * MB_STATE; + xvr[gi] = __bfloat162float( + xrow[head_in_group * MB_HEAD_DIM + + first_channel + g0 + gi]); + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = st_in[prow + lane + j * 32]; + } + } else if constexpr (SSM_SMEM_OK) { + asm volatile("cp.async.wait_group 3;\n" ::: "memory"); + __syncwarp(); + const int pair = head_in_group * SSM_GROUPS_PER_HEAD + + g0 / SSM_PREFETCH; + const int buffer = pair & (SSM_RING - 1); + const r9_state_t *sm = reinterpret_cast( + warp_ssm + (size_t)buffer * SSM_GROUP_BYTES); + static_assert(SSM_PREFETCH % 2 == 0 && MB_HEAD_DIM % 2 == 0, + "x staging vectorises in element pairs"); + const __nv_bfloat162 *const xq = + reinterpret_cast( + xrow + head_in_group * MB_HEAD_DIM + + first_channel + g0); + #pragma unroll + for (int gq = 0; gq < SSM_PREFETCH / 2; ++gq) { + const float2 xp = __bfloat1622float2(xq[gq]); + xvr[2 * gq] = xp.x; + xvr[2 * gq + 1] = xp.y; + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = r9_load( + sm[(size_t)gi * MB_STATE + lane + j * 32], 1.f); + } + __syncwarp(); + const int ahead = pair + SSM_RING; + if (ahead < SSM_PAIRS) + issue_ssm(ahead / SSM_GROUPS_PER_HEAD, + (ahead % SSM_GROUPS_PER_HEAD) * SSM_PF, buffer); + else + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } else { + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + xvr[gi] = __bfloat162float( + xrow[head_in_group * MB_HEAD_DIM + + first_channel + g0 + gi]); + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + { + int p = first_channel+g0+gi; + int mapped = R9_STATE_MODE==4 ? L.r9_map[head*24+p] : head*24+p; + int n = lane+j*32; + float scale = R9_STATE_MODE==1 ? 1.f : L.r9_scales[r9_scale_index(mapped/24,mapped%24,n)]; + sv[gi][j] = r9_load(st[(size_t)mapped*128+n],scale); + } + } + } + if constexpr (SSM_PREFETCH == 12) { + float v6[6]; + #pragma unroll + for (int k = 0; k < 6; ++k) { + float pair_acc[2]; + #pragma unroll + for (int half = 0; half < 2; ++half) { + const int gi = k + 6 * half; + const int p = first_channel + g0 + gi; + const float xv = xvr[gi]; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < NSTATE; ++j) { + const int n = lane + j * 32; + const float s = sv[gi][j] * dA + Bv[j] * xv; + const int mapped = R9_STATE_MODE == 4 + ? L.r9_map[head * 24 + p] : head * 24 + p; + const float scale = R9_STATE_MODE == 1 ? 1.f + : L.r9_scales[r9_scale_index( + mapped / 24, mapped % 24, n)]; + st[(size_t)mapped * 128 + n] = r9_store(s, scale); + acc += s * Cv[j]; + } + pair_acc[half] = acc; + } + const float mine = sb4 ? pair_acc[1] : pair_acc[0]; + const float away = sb4 ? pair_acc[0] : pair_acc[1]; + v6[k] = mine + __shfl_xor_sync(0xffffffff, away, 16); + } + float v4[4]; + #pragma unroll + for (int k = 0; k < 3; ++k) { + const float mine = sb3 ? v6[3 + k] : v6[k]; + const float away = sb3 ? v6[k] : v6[3 + k]; + v4[k] = mine + __shfl_xor_sync(0xffffffff, away, 8); + } + v4[3] = 0.f; + float v2[2]; + #pragma unroll + for (int k = 0; k < 2; ++k) { + const float mine = sb2 ? v4[2 + k] : v4[k]; + const float away = sb2 ? v4[k] : v4[2 + k]; + v2[k] = mine + __shfl_xor_sync(0xffffffff, away, 4); + } + float acc = (sb1 ? v2[1] : v2[0]) + + __shfl_xor_sync(0xffffffff, sb1 ? v2[0] : v2[1], 2); + acc += __shfl_xor_sync(0xffffffff, acc, 1); + if (ssm_owner) { + const int p = first_channel + g0 + ssm_own_ch; + const float xv = __bfloat162float( + xrow[head_in_group * MB_HEAD_DIM + p]); + const bf16 outv = __float2bfloat16(acc + xv * Dh); + if constexpr (FUSE_GNORM) + gn_stage[warp][head_in_group * MB_HEAD_DIM + p] = outv; + else + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = outv; + } + } else { + float cacc[SSM_PREFETCH]; + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const int p = first_channel + g0 + gi; + const float xv = xvr[gi]; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < NSTATE; ++j) { + const int n = lane + j * 32; + const float s = sv[gi][j] * dA + Bv[j] * xv; + const int mapped = R9_STATE_MODE == 4 + ? L.r9_map[head * 24 + p] : head * 24 + p; + const float scale = R9_STATE_MODE == 1 ? 1.f + : L.r9_scales[r9_scale_index( + mapped / 24, mapped % 24, n)]; + st[(size_t)mapped * 128 + n] = r9_store(s, scale); + acc += s * Cv[j]; + } + cacc[gi] = acc; + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const int p = first_channel + g0 + gi; + float acc = cacc[gi]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) { + const bf16 outv = __float2bfloat16(acc + xvr[gi] * Dh); + if constexpr (FUSE_GNORM) + gn_stage[warp][head_in_group * MB_HEAD_DIM + p] = outv; + else + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = outv; + } + } + } + } + } + if constexpr (SSM_SMEM_OK) { + if (!r8_initial || SSM_F32_OK) { + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); + __syncwarp(); + } + } + if constexpr (FUSE_GNORM) { + __syncwarp(); + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32], wg[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float gv = __bfloat162float(gn_stage[warp][lane + j * 32]) * + (z / (1.f + __expf(-z))); + wg[j] = __bfloat162float(L.gnorm_w[idx]); + held[j] = gv; + sq += gv * gv; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16(held[j] * inv * wg[j]); + } + __syncwarp(); + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_CONV_DIM; + const int c = idx - t * MB_CONV_DIM; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v; +#ifdef PF_CONV_ZERO_STATE + if (src < 0) v = 0.f; +#else + if (src < 0) v = __bfloat162float(state[pos + k]); +#endif + else v = __bfloat162float( + zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]); + acc += __bfloat162float(w[k]) * v; + } +#ifndef PF_CONV_NO_SILU + acc = acc / (1.f + __expf(-acc)); +#endif + xbc_act[(size_t)t * MB_CONV_DIM + c] = __float2bfloat16(acc); + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + const int units = n_seq * MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int seq = unit / (MB_HEADS * MB_HEAD_DIM); + const int rem = unit - seq * (MB_HEADS * MB_HEAD_DIM); + const int head = rem / MB_HEAD_DIM; + const int p = rem - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p * MB_STATE; + float sreg[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) sreg[j] = st[lane + j * 32]; + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const bf16 *xa = xbc_act + t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + const bf16 *dtrow = zxbcdt + t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + sreg[j] = sreg[j] * dA + __bfloat162float(Bsec[n]) * dt * xv; + acc += sreg[j] * __bfloat162float(Csec[n]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) st[lane + j * 32] = sreg[j]; + } + return; + } +#if PF_SSM_SPLIT_P + const int units = MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; unit += warps_total) { + const int head = unit / MB_HEAD_DIM; + const int p_only = unit - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; +#else + for (int head = blockIdx.x * WARPS + warp; head < MB_HEADS; + head += warps_total) { + const int group = head / MB_HEADS_PER_GROUP; +#endif + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *st = L.ssm_state + (size_t)slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE +#if PF_SSM_SPLIT_P + + (size_t)p_only * MB_STATE +#endif + ; + for (int t = 0; t < n_tokens; ++t) { + const bf16 *xa = xbc_act + (size_t)t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = + zxbcdt + (size_t)t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } +#if PF_SSM_SPLIT_P + { + const int p = p_only; +#else + for (int p = 0; p < MB_HEAD_DIM; ++p) { +#endif + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { +#if PF_SSM_SPLIT_P + const size_t at = (size_t)lane + j * 32; +#else + const size_t at = (size_t)p * MB_STATE + lane + j * 32; +#endif + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } + } +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32], wg[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + wg[j] = __bfloat162float(L.gnorm_w[idx]); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16(held[j] * inv * wg[j]); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len MK2_PROGRESS_PARAMETER(progress_mask)) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + constexpr int Q_PER_WARP = 3; + constexpr int GROUPS_PER_KV = 1; + const int pairs = BATCH * BB_KV_HEADS * GROUPS_PER_KV; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x + warp * gridDim.x; pair < pairs; + pair += gridDim.x * WARPS) { + const int row = pair / (BB_KV_HEADS * GROUPS_PER_KV); + MK2_PROGRESS_GUARD(progress_mask, row); + const int within_row = pair % (BB_KV_HEADS * GROUPS_PER_KV); + const int kv = within_row / GROUPS_PER_KV; + const int q_group = within_row % GROUPS_PER_KV; + const int q_first = q_group * Q_PER_WARP; + const int q_count = q_group == 0 ? Q_PER_WARP : 1; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (q_group == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[h][j] = 0.f; + } + float m[Q_PER_WARP], denom[Q_PER_WARP]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + m[h] = -CUDART_INF_F; + denom[h] = 0.f; + } + constexpr int PF = 8; + const int tpb = L.tokens_per_block; + const size_t elems_per_block = + (size_t)tpb * BB_KV_HEADS * BB_HEAD_DIM; + const int *bt_k = L.block_table + (size_t)(row * 2) * L.max_blocks_per_seq; + const int *bt_v = bt_k + L.max_blocks_per_seq; + const int head_base = kv * tpb * BB_HEAD_DIM; + float qr[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const bf16 *q = base + (kv * kv_per_q + q_first + h) * BB_HEAD_DIM; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qr[h][j] = __bfloat162float(q[lane + j * 32]); + } + for (int page0 = 0; page0 < T; page0 += tpb) { + const int k_page = bt_k[page0 / tpb] & 0x7FFFFFFF; + const int v_page = bt_v[page0 / tpb] & 0x7FFFFFFF; + const bf16 *kbase = L.pool_base + + (size_t)k_page * elems_per_block + + head_base; + const bf16 *vbase = L.pool_base + + (size_t)v_page * elems_per_block + + head_base; + const int tend = page0 + tpb < T ? page0 + tpb : T; + for (int t0 = page0; t0 < tend; t0 += PF) { + const int navail = tend - t0 < PF ? tend - t0 : PF; + float kr[PF][PER_LANE], vr[PF][PER_LANE]; + #pragma unroll + for (int u = 0; u < PF; ++u) { + const int tpos = t0 + (u < navail ? u : navail - 1); +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = vbase + (tpos - page0) * BB_HEAD_DIM; +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new : kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = cur_token + ? v_new : vbase + (tpos - page0) * BB_HEAD_DIM; +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + kr[u][j] = __bfloat162float(kt[i]); + vr[u][j] = __bfloat162float(vt[i]); + } + } + #pragma unroll + for (int u = 0; u < PF; ++u) { + if (u >= navail) break; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + dot += qr[h][j] * kr[u][j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m[h], s); + const float rescale = __expf(m[h] - m_new); + const float w = __expf(s - m_new); + denom[h] = denom[h] * rescale + w; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[h][j] = acc[h][j] * rescale + w * vr[u][j]; + m[h] = m_new; + } + } + } + } + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const float inv = 1.f / denom[h]; + const int head = kv * kv_per_q + q_first + h; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[h][j] * inv); + } + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * BB_KV_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / BB_KV_HEADS; + const int kv = idx - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv); + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv); + for (int i = 0; i < BB_HEAD_DIM; ++i) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + for (int s = 0; s <= pos; ++s) { + const bf16 *kt = kv_slot(L, seq, 0, s) + kv_offset(L, s, kv); + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + const bf16 *vt = kv_slot(L, seq, 1, s) + kv_offset(L, s, kv); + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + wt * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH, bf16 *row_stage = nullptr) { + __shared__ float score[BB_EXPERTS]; + constexpr int CH = D_MODEL / 8; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const char *gsrc = reinterpret_cast( + normed + (size_t)row * D_MODEL); + char *gdst = reinterpret_cast(row_stage); + const uint32_t act_s = (uint32_t)__cvta_generic_to_shared(gdst); + for (int i = threadIdx.x; i < CH; i += THREADS) { + const uint32_t dst = act_s + (uint32_t)i * 16; + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: + "r"(dst), "l"(gsrc + (size_t)i * 16) : "memory"); + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + if ((int)threadIdx.x < BB_EXPERTS) + moe_w[(size_t)row * BB_EXPERTS + threadIdx.x] = 0.f; + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); + __syncthreads(); + float acc0 = 0.f, acc1 = 0.f, acc2 = 0.f; + const bf16 *w0 = gate_w + (size_t)warp * D_MODEL; + const bf16 *w1 = gate_w + (size_t)(warp + WARPS) * D_MODEL; + const bf16 *w2 = gate_w + (size_t)(warp + 2 * WARPS) * D_MODEL; + for (int i = lane; i < D_MODEL; i += 128) { + bf16 sx0, sx1; + move::lds(sx0, act_s + (uint32_t)i * sizeof(bf16)); + move::lds(sx1, act_s + (uint32_t)(i + 32) * sizeof(bf16)); + float x0 = __bfloat162float(sx0); + float x1 = __bfloat162float(sx1); + acc0 += x0 * __bfloat162float(w0[i]); + acc1 += x0 * __bfloat162float(w1[i]); + acc2 += x0 * __bfloat162float(w2[i]); + acc0 += x1 * __bfloat162float(w0[i + 32]); + acc1 += x1 * __bfloat162float(w1[i + 32]); + acc2 += x1 * __bfloat162float(w2[i + 32]); + move::lds(sx0, act_s + (uint32_t)(i + 64) * sizeof(bf16)); + move::lds(sx1, act_s + (uint32_t)(i + 96) * sizeof(bf16)); + x0 = __bfloat162float(sx0); + x1 = __bfloat162float(sx1); + acc0 += x0 * __bfloat162float(w0[i + 64]); + acc1 += x0 * __bfloat162float(w1[i + 64]); + acc2 += x0 * __bfloat162float(w2[i + 64]); + acc0 += x1 * __bfloat162float(w0[i + 96]); + acc1 += x1 * __bfloat162float(w1[i + 96]); + acc2 += x1 * __bfloat162float(w2[i + 96]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + acc0 += __shfl_xor_sync(0xffffffffu, acc0, off); + acc1 += __shfl_xor_sync(0xffffffffu, acc1, off); + acc2 += __shfl_xor_sync(0xffffffffu, acc2, off); + } + if (lane == 0) { + score[warp] = 1.f / (1.f + __expf(-acc0)); + score[warp + WARPS] = 1.f / (1.f + __expf(-acc1)); + score[warp + 2 * WARPS] = 1.f / (1.f + __expf(-acc2)); + } + __syncthreads(); + if (threadIdx.x < 32) { + float remaining = (lane < BB_EXPERTS) + ? score[lane] + __bfloat162float(gate_bias[lane]) + : -CUDART_INF_F; + int pick[BB_TOP_K]; + float total = 0.f; + #pragma unroll + for (int k = 0; k < BB_TOP_K; ++k) { + float bv = remaining; + int bi = lane; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffff, bv, off); + const int oi = __shfl_xor_sync(0xffffffff, bi, off); + if (ov > bv || (ov == bv && oi < bi)) { bv = ov; bi = oi; } + } + pick[k] = bi; + total += score[bi]; + if (lane == bi) remaining = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + #pragma unroll + for (int k = 0; k < BB_TOP_K; ++k) + if (lane == pick[k]) moe_w[row * BB_EXPERTS + lane] = score[lane] * norm; + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert <= BB_EXPERTS; ++expert) + value = __fadd_rn(value, parts[base + (size_t)expert * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ bf16 bb_moe_fold_selected( + const float *parts, const float *weights, int idx) { + const int row = idx / D_MODEL; + const size_t base = (size_t)row * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert < BB_EXPERTS; ++expert) { + const float part = weights[row * BB_EXPERTS + expert] != 0.f + ? parts[base + (size_t)expert * D_MODEL] : 0.f; + value = __fadd_rn(value, part); + } + value = __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ __forceinline__ void sample_argmax_merge( + float &value, int &rank, float other, int other_rank) { + if (other > value || (other == value && other_rank < rank)) { + value = other; + rank = other_rank; + } +} +__device__ __forceinline__ int phoneme_warp_argmax(const bf16 *lg) { + const int lane = threadIdx.x & 31; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = lane; n < BB_PHONEME; n += 32) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, best, off); + const int other_arg = __shfl_down_sync(0xffffffffu, arg, off); + sample_argmax_merge(best, arg, other, other_arg); + } + return arg; +} +__device__ __forceinline__ int sample_block_argmax(float value, int rank) { + __shared__ float warp_value[WARPS]; + __shared__ int warp_rank[WARPS]; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + if (lane == 0) { warp_value[warp] = value; warp_rank[warp] = rank; } + __syncthreads(); + if (warp == 0) { + value = lane < WARPS ? warp_value[lane] : -CUDART_INF_F; + rank = lane < WARPS ? warp_rank[lane] : INT_MAX; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + } + return rank; +} +__device__ __forceinline__ int sample_eos_insertion_rank( + const uint64_t *keys, float eos) { + int lo = 0, hi = TOP_K; + while (lo < hi) { + const int mid = (lo + hi) >> 1; + if (eos > ordered_to_float((uint32_t)(keys[mid] >> 32))) hi = mid; + else lo = mid + 1; + } + return lo; +} +__device__ __forceinline__ void sample_topk_serial( + const uint64_t *keys, float eos, const float *noise, int *out) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float value; + int code; + if (!eos_used && eos > av) { + value = eos; + code = AUDIO_EOS; + eos_used = true; + } else { + value = av; + code = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = value + noise[j]; + if (noised > best_noised) { + best_noised = noised; + best_code = code; + } + } + *out = best_code; +} +__device__ __forceinline__ void sample_topk_parallel( + const uint64_t *keys, float eos, const float *noise, int *out) { + const int j = threadIdx.x; + const bool sorted_nan = j < TOP_K && + isnan(ordered_to_float((uint32_t)(keys[j] >> 32))); + if (__syncthreads_or(sorted_nan)) { + if (threadIdx.x == 0) sample_topk_serial(keys, eos, noise, out); + return; + } + const int lane = threadIdx.x & 31; + int insertion = lane == 0 ? sample_eos_insertion_rank(keys, eos) : 0; + insertion = __shfl_sync(0xffffffffu, insertion, 0); + float value = -CUDART_INF_F; + int rank = INT_MAX; + if (j < TOP_K) { + const int a = j - (j > insertion); + const float v = j == insertion + ? eos : ordered_to_float((uint32_t)(keys[a] >> 32)); + const float noised = v + noise[j]; + if (noised > -CUDART_INF_F) { value = noised; rank = j; } + } + const int winner = sample_block_argmax(value, rank); + if (threadIdx.x == 0) { + if (winner == INT_MAX) *out = 0; + else if (winner == insertion) *out = AUDIO_EOS; + else { + const int a = winner - (winner > insertion); + *out = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + } + } +} +constexpr int SORT_ITEMS_PER_THREAD = SORT_N / THREADS; +struct sampler_storage { + uint64_t keys[2][SORT_N]; +}; +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + sampler_storage &storage, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + uint64_t thread_keys[SORT_ITEMS_PER_THREAD]; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + thread_keys[item] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + int cross_stage = 0; + for (int k = 2; k <= SORT_N; k <<= 1) { + for (int j = k >> 1; j > 0; j >>= 1) { + if (j >= 32) { + const int bank = cross_stage++ & 1; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + storage.keys[bank][n] = thread_keys[item]; + } + __syncthreads(); + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + const uint64_t other = storage.keys[bank][n ^ j]; + const bool take_hi = ((n & k) == 0) == ((n & j) == 0); + const uint64_t mine = thread_keys[item]; + thread_keys[item] = take_hi + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } else { + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + const uint64_t other = __shfl_xor_sync( + 0xffffffffu, thread_keys[item], j); + const bool take_hi = ((n & k) == 0) == ((n & j) == 0); + const uint64_t mine = thread_keys[item]; + thread_keys[item] = take_hi + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } + } + } + const int top_bank = cross_stage & 1; + if (threadIdx.x < TOP_K) + storage.keys[top_bank][threadIdx.x] = thread_keys[0]; + __syncthreads(); + sample_topk_parallel( + storage.keys[top_bank], eos_v, + gumbel + ((size_t)row * N_POS + pos) * TOP_K, + out_codes + row); + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void r6_sparse_map(const float* routes,int* map,int* counts) { + __shared__ int r7_warp_counts[THREADS/32]; + const int lane=threadIdx.x%32, warp=threadIdx.x/32; + const unsigned lower=(1u<>1; + for(int e=0;e(packed)[(size_t)e*BATCH*row_vectors+i]= + row<0?make_uint4(0,0,0,0):reinterpret_cast(source)[(size_t)row*row_vectors+i%row_vectors]; + } + } +} +__device__ void r6_tile(int work,int output_tiles,const int* lut,int &tile,int &tok, + int mtg) { + const int groups=output_tiles/mtg; + tok=lut[1+work/groups]; + const int e=tok/DECODE_TILES; + tile=e*output_tiles+(work%groups)*mtg; +} +__device__ __forceinline__ void mk2_group_tile( + int work, int grouped, int &tile, int &tok, int &here) { + here = work < grouped ? NM_MAX : 1; + const int unit = work < grouped ? work * NM_MAX : grouped * (NM_MAX - 1) + work; + tile = unit / DECODE_TILES; + tok = unit - tile * DECODE_TILES; +} +__device__ __forceinline__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&p_s)[STAGES][SLOTS], out_tile &o_s, + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + int &tile_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = DECODE_TILES, + int rows_valid = DECODE_ROWS, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const int *r6_lut = nullptr,const int *r6_map = nullptr) { + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + constexpr int QSLOT = MT + NM_MAX; + constexpr bool WIDE = MT > 1 && NM_MAX > 1 && STAGES * SLOTS >= 2 * QSLOT; + constexpr int PSTAGES = WIDE ? 2 : STAGES; + constexpr int PSLOT = WIDE ? QSLOT : SLOTS; + constexpr int ACCN = WIDE ? MT * NM_MAX : NM_MAX; + constexpr int HEX_MT = 3, HEX_NM = 2; + constexpr int HEX_SLOT = HEX_MT + HEX_NM; + auto epi_publish = [](out_tile &st, const rt_fl<16, BATCH_PAD> &src) { + using acc_t = rt_fl<16, BATCH_PAD>; + static_assert(acc_t::height == 1, "one row block per warp"); + static_assert(out_tile::rows == BATCH_PAD && out_tile::cols == OUT_ATOM, + "staging tile must be square for an in-place transpose"); + const int t = ::kittens::laneid(); + const uint32_t base = + static_cast(__cvta_generic_to_shared(&st.data[0])); + const int row0 = 8 * (t >> 4) + (t & 7); + const int col0 = 16 * warpgroup::warpid() + 8 * ((t >> 3) & 1); + #pragma unroll + for (int j = 0; j < acc_t::width; ++j) { + bf16_2 q0 = base_types::convertor::convert( + src.tiles[0][j].data[0]); + bf16_2 q1 = base_types::convertor::convert( + src.tiles[0][j].data[1]); + bf16_2 q2 = base_types::convertor::convert( + src.tiles[0][j].data[2]); + bf16_2 q3 = base_types::convertor::convert( + src.tiles[0][j].data[3]); + move::stsm4t( + out_tile::idx(base, {16 * j + row0, col0}), q0, q1, q2, q3); + } + }; + weight_tile *const pf = &p_s[0][0]; + const bool arena10 = + reinterpret_cast(&o_sf) == + reinterpret_cast(pf + STAGES * SLOTS); + const bool two_d = !r6_lut && out_tiles_per_expert == 0 + && n_tok == DECODE_TILES && DECODE_TILES % NM_MAX == 0; + int n_tiles = compact_tiles ? __popc(active_mask)*compact_tiles : n_out/OUT_ATOM; + constexpr int QTOK = DECODE_TILES / NM_MAX; + const int out_tiles = n_tiles; + const bool tall3 = WIDE && two_d && out_tiles == 24; + const bool split72 = WIDE && two_d && DECODE_TILES == 16 && out_tiles == 72; + const bool split81 = WIDE && two_d && DECODE_TILES == 16 && out_tiles == 81; + const bool split17 = WIDE && two_d && DECODE_TILES == 16 && out_tiles == 17; + const bool mixed40 = WIDE && two_d && DECODE_TILES == 16 && out_tiles == 40; + const int n_tiles_hex_all = (out_tiles / HEX_MT) * (DECODE_TILES / HEX_NM); + const bool hex = WIDE && two_d && arena10 && DECODE_TILES == 16 && DECODE_TILES % HEX_NM == 0 && + (out_tiles == 96 || out_tiles == 81 || out_tiles == 72 || + out_tiles == 40) && + (epi == EPI_STORE || epi == EPI_GELU) && !r6_lut && + out_tiles_per_expert == 0; + const int hex_tok = DECODE_TILES / HEX_NM; + const int hex_wide = (out_tiles == 72 || out_tiles == 81) + ? 128 : n_tiles_hex_all; + const int pslot = hex ? HEX_SLOT : PSLOT; + int q_units = 0, q_rem_base = 0; + int grouped = 0; + if constexpr (WIDE) { + if (two_d) { + q_units = (out_tiles / MT) * QTOK; + q_rem_base = out_tiles - out_tiles % MT; + } + } else if (two_d) { + const int g = (int)gridDim.x; + const int units = n_tiles * DECODE_TILES; + const int rounds = (units + g - 1) / g; + const int cap = units / NM_MAX; + grouped = g * (rounds / NM_MAX); + if (grouped > cap) grouped = cap; + } + const int r6_output_tiles = r6_lut ? n_out/BB_EXPERTS/OUT_ATOM : 0; + const int r6_mtg = (MT > 1 && r6_lut && r6_output_tiles % MT == 0) ? MT : 1; + const int r6_groups = r6_lut ? r6_output_tiles/r6_mtg : 0; + const int r6_rows = r6_lut ? (r6_lut[0] & 0xFFF) : 0; + const int r6_pairs = (r6_lut && WIDE && r6_mtg == MT) ? (r6_lut[0] >> 12) : 0; + const int r6_pair_items = r6_pairs * r6_groups; + const int r6_single_mtg = (r6_lut && WIDE && r6_output_tiles % 3 == 0) + ? 3 : r6_mtg; + const int r6_single_groups = r6_lut ? r6_output_tiles/r6_single_mtg : 0; + const int r6_single_items = r6_lut + ? (r6_rows - 2 * r6_pairs) * r6_single_groups : 0; + if(r6_lut) { + n_tiles=r6_pair_items+r6_single_items; + n_tok=1; + } else if(two_d) { + if constexpr (WIDE) { + if (hex) n_tiles = (out_tiles == 72) ? 256 + : (out_tiles == 81) ? 264 + : (out_tiles == 40) + ? n_tiles_hex_all + hex_tok + : n_tiles_hex_all; + else if (mixed40) n_tiles = 244; + else if (split72) n_tiles = 360; + else if (split81) n_tiles = 384; + else if (split17) n_tiles = 132; + else if (tall3) n_tiles = (out_tiles / 3) * DECODE_TILES; + else n_tiles = q_units + (out_tiles % MT) * QTOK; + } + else n_tiles = grouped + (n_tiles*DECODE_TILES - grouped*NM_MAX); + n_tok=1; } + const int r6_extra = r6_lut ? n_tiles % (int)gridDim.x : 0; + const int r6_skew = r6_lut + ? (r6_single_items < r6_extra ? r6_single_items : r6_extra) : 0; + const int iters = k_total / KCHUNK; + int tile_start = (int)blockIdx.x - tile_io % (int)gridDim.x; + if (tile_start < 0) tile_start += (int)gridDim.x; + tile_io += n_tiles; + int tiles_here = 0; + for (int t = tile_start; t < n_tiles; t += gridDim.x) ++tiles_here; + const int iter_base = iter_io; + iter_io += tiles_here * iters * n_tok; + const int my_tiles = (n_tiles - tile_start + gridDim.x - 1) / gridDim.x; + if (warpgroup::groupid() == 1) { + if (hex) { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) { + if (out_tiles == 81 && work >= hex_wide) { + const int j = work - hex_wide; + const bool pair = j >= 128; + const int nt = pair ? 80 : 48 + 2 * (j / hex_tok); + const int ntok = HEX_NM * (j % hex_tok); + if (pair) { + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + wait(drained[s], (iter / PSTAGES) & 1); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], sizeof(weight_tile) + + HEX_NM * sizeof(act_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot], w, {nt, i}, filled[s]); + #pragma unroll + for (int m = 0; m < HEX_NM; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + 1 + m], a, + {ntok + m, i}, filled[s]); + } + } + } else { + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + wait(drained[s], (iter / PSTAGES) & 1); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], 2 * sizeof(weight_tile) + + HEX_NM * sizeof(act_tile)); + #pragma unroll + for (int c = 0; c < 2; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {nt + c, i}, filled[s]); + #pragma unroll + for (int m = 0; m < HEX_NM; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + 2 + m], a, + {ntok + m, i}, filled[s]); + } + } + } + continue; + } + if (out_tiles == 40 && work >= hex_wide) { + const int ntok = HEX_NM * (work - hex_wide); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + wait(drained[s], (iter / PSTAGES) & 1); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], sizeof(weight_tile) + + HEX_NM * sizeof(act_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot], w, {out_tiles - 1, i}, filled[s]); + #pragma unroll + for (int m = 0; m < HEX_NM; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + 1 + m], a, + {ntok + m, i}, filled[s]); + } + } + continue; + } + if (work >= hex_wide) { + const int j = work - hex_wide; + const int nt = HEX_MT * ((hex_wide + (j >> 1)) / hex_tok); + const int ntok = HEX_NM * ((hex_wide + (j >> 1)) % hex_tok) + (j & 1); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + wait(drained[s], (iter / PSTAGES) & 1); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], HEX_MT * sizeof(weight_tile) + + sizeof(act_tile)); + #pragma unroll + for (int c = 0; c < HEX_MT; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {nt + c, i}, filled[s]); + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + HEX_MT], a, {ntok, i}, filled[s]); + } + } + continue; + } + const int tile_id = HEX_MT * (work / hex_tok); + const int tok = HEX_NM * (work % hex_tok); + rt_fl<16, BATCH_PAD> pacc[HEX_NM]; + #pragma unroll + for (int u = 0; u < HEX_NM; ++u) warp::zero(pacc[u]); + int pend = -1, pend_ph = 0; + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(drained[s], phase); + if (warpgroup::laneid() == 0) { + tma::expect_bytes(filled[s], HEX_MT * sizeof(weight_tile) + + HEX_NM * sizeof(act_tile)); + #pragma unroll + for (int c = 0; c < HEX_MT; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {tile_id + c, i}, filled[s]); + #pragma unroll + for (int m = 0; m < HEX_NM; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + HEX_MT + m], a, {tok + m, i}, filled[s]); + } + if (pend >= 0) { { + wait(filled[pend], pend_ph); + warpgroup::mma_ABt(pacc[0], pf[pend * pslot + HEX_MT - 1], + pf[pend * pslot + HEX_MT]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + pacc[1], pf[pend * pslot + HEX_MT - 1], + pf[pend * pslot + HEX_MT + 1]); + warpgroup::mma_async_wait(); + } } + pend = s; pend_ph = phase; + } + { + wait(filled[pend], pend_ph); + warpgroup::mma_ABt(pacc[0], pf[pend * pslot + HEX_MT - 1], + pf[pend * pslot + HEX_MT]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + pacc[1], pf[pend * pslot + HEX_MT - 1], + pf[pend * pslot + HEX_MT + 1]); + warpgroup::mma_async_wait(); + } + const int es = iter % PSTAGES; + wait(drained[es], (iter / PSTAGES) & 1); + out_tile_f &q_sf = + *reinterpret_cast(&pf[es * pslot]); + out_tile &q_s = *reinterpret_cast(&q_sf); + const int qbase = (tile_id + HEX_MT - 1) * OUT_ATOM; + constexpr int OUT_OCTS = OUT_ATOM / 8; + #pragma unroll + for (int u = 0; u < HEX_NM; ++u) { + const int grow0 = (tok + u) * BATCH_PAD; + epi_publish(q_s, pacc[u]); + warpgroup::sync(3); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); + idx < OUT_OCTS * rows_valid; idx += 128) { + const int row = idx >> 3, col = (idx & 7) << 3; + const size_t at = + (size_t)(grow0 + row) * dest_stride + qbase + col; + const bf16_oct q = *reinterpret_cast( + &q_s[int2{row, col}]); + if (epi == EPI_STORE) { + *reinterpret_cast(dest + at) = q; + } else { + *reinterpret_cast(dest + at) = bf16_oct{ + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x0))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x1))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x2))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x3))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x4))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x5))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x6))), + __float2bfloat16(gelu_tanh_fp32(__bfloat162float(q.x7)))}; + } + } + warpgroup::sync(3); + } + } + } else if (warpgroup::laneid() == 0) { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) { + int tile_id = mk2_expert_tile(work, active_mask, compact_tiles), r6_tok=0; + int nm = 1, mt = 1; + if(r6_lut) { + int rwork = work - r6_skew; + if (rwork < 0) rwork += n_tiles; + const bool pair = rwork < r6_pair_items; + const int q = pair ? rwork / r6_groups + : (rwork-r6_pair_items) / r6_single_groups; + const int group = pair ? rwork-q*r6_groups + : rwork-r6_pair_items-q*r6_single_groups; + const int base_item = pair ? 2*q : 2*r6_pairs+q; + const int ltok = r6_lut[1 + base_item]; + r6_tok = ltok; + tile_id = (ltok / DECODE_TILES) * r6_output_tiles + + group * (pair ? r6_mtg : r6_single_mtg); + mt = pair ? r6_mtg : r6_single_mtg; + nm = pair ? NM_MAX : 1; + } + else if(two_d) { + if constexpr (WIDE) { + if (mixed40) { + const bool quad = work < 132; + const int item = quad ? work : work - 132; + if (quad && item < 84) { + tile_id = 2 * (item / 7); + r6_tok = 2 * (item % 7); + } else if (quad) { + const int tail = item - 84; + tile_id = 2 * (12 + tail / 6); + r6_tok = 2 * (tail % 6); + } else if (item < 48) { + tile_id = 2 * (item / 4) + (item % 4) / 2; + r6_tok = 14 + item % 2; + } else { + const int tail = item - 48; + tile_id = 2 * (12 + tail / 8) + (tail % 8) / 4; + r6_tok = 12 + tail % 4; + } + mt = quad ? MT : 1; + nm = quad ? NM_MAX : 1; + } else if (split72) { + if (work < 264) { + tile_id = 2 * (work / 8); + r6_tok = 2 * (work % 8); + mt = MT; nm = NM_MAX; + } else { + const int i = work - 264, c = i >> 2, q = i & 3; + tile_id = 2 * (33 + (c >> 3)) + (q >> 1); + r6_tok = 2 * (c & 7) + (q & 1); + mt = 1; nm = 1; + } + } else if (split81) { + if (work < 264) { + tile_id = 2 * (work / 8); + r6_tok = 2 * (work % 8); + mt = MT; nm = NM_MAX; + } else if (work < 376) { + const int i = work - 264; + tile_id = 2 * (33 + i / 16); + r6_tok = i % 16; + mt = MT; nm = 1; + } else { + tile_id = 80; + r6_tok = 2 * (work - 376); + mt = 1; nm = NM_MAX; + } + } else if (split17) { + if (work < 8) { + tile_id = 14; + r6_tok = work; + mt = 3; nm = 1; + } else if (work < 120) { + const int i = work - 8; + tile_id = i >> 3; + r6_tok = 2 * (i & 7); + mt = 1; nm = NM_MAX; + } else { + const int i = work - 120; + tile_id = 14 + (i >> 2); + r6_tok = 8 + 2 * (i & 3); + mt = 1; nm = NM_MAX; + } + } else if (tall3) { + tile_id = 3 * (work / DECODE_TILES); + r6_tok = work % DECODE_TILES; + mt = 3; + nm = 1; + } else { + const int j = work < q_units ? work : work - q_units; + tile_id = work < q_units ? MT * (j / QTOK) + : q_rem_base + j / QTOK; + r6_tok = NM_MAX * (j % QTOK); + mt = work < q_units ? MT : 1; + nm = NM_MAX; + } + } else mk2_group_tile(work,grouped,tile_id,r6_tok,nm); + } + for (int token_iteration = 0; token_iteration < n_tok; ++token_iteration) + for (int i = 0; i < iters; ++i, ++iter) { + const int tok = (r6_lut || two_d) ? r6_tok : token_iteration; + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(drained[s], phase); + tma::expect_bytes(filled[s], + mt * sizeof(weight_tile) + nm * sizeof(act_tile)); + if (mt == 1) { + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot], w, {tile_id, i}, filled[s]); + } else if (mt == 3) { + #pragma unroll + for (int c = 0; c < 3; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {tile_id + c, i}, filled[s]); + } else { + #pragma unroll + for (int c = 0; c < MT; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + pf[s * pslot + c], w, {tile_id + c, i}, filled[s]); + } + const int a_col = out_tiles_per_expert && !r6_lut + ? (tile_id / out_tiles_per_expert) * act_k_tiles_per_expert + i + : i; + if (NM_MAX == 1 || nm == NM_MAX) { + #pragma unroll + for (int m = 0; m < NM_MAX; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + mt + m], a, {tok + m, a_col}, filled[s]); + } else { + tma::load_async<2, cache_policy::EVICT_LAST>( + pf[s * pslot + mt], a, {tok, a_col}, filled[s]); + } + } + } + } + } else { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) + for (int token_iteration = 0; token_iteration < n_tok; ++token_iteration) { + int tile_id=mk2_expert_tile(work, active_mask, compact_tiles),tok=token_iteration; + int nm = 1, mt = 1; + if(r6_lut) { + int rwork = work - r6_skew; + if (rwork < 0) rwork += n_tiles; + const bool pair = rwork < r6_pair_items; + const int q = pair ? rwork / r6_groups + : (rwork-r6_pair_items) / r6_single_groups; + const int group = pair ? rwork-q*r6_groups + : rwork-r6_pair_items-q*r6_single_groups; + const int base_item = pair ? 2*q : 2*r6_pairs+q; + const int ltok = r6_lut[1 + base_item]; + tok = ltok; + tile_id = (ltok / DECODE_TILES) * r6_output_tiles + + group * (pair ? r6_mtg : r6_single_mtg); + mt = pair ? r6_mtg : r6_single_mtg; + nm = pair ? NM_MAX : 1; + } + else if(two_d) { + if constexpr (WIDE) { + if (hex) { + if (out_tiles == 81 && work >= hex_wide) { + const int j = work - hex_wide; + const bool pair = j >= 128; + tile_id = pair ? 80 : 48 + 2 * (j / hex_tok); + tok = HEX_NM * (j % hex_tok); + mt = pair ? 1 : 2; + nm = HEX_NM; + } else if (out_tiles == 40 && work >= hex_wide) { + tile_id = out_tiles - 1; + tok = HEX_NM * (work - hex_wide); + mt = 1; + nm = HEX_NM; + } else { + const int hw = work < hex_wide + ? work : hex_wide + ((work - hex_wide) >> 1); + tile_id = HEX_MT * (hw / hex_tok); + tok = HEX_NM * (hw % hex_tok) + + (work < hex_wide ? 0 : ((work - hex_wide) & 1)); + mt = HEX_MT; nm = work < hex_wide ? HEX_NM : 1; + } + } else if (mixed40) { + const bool quad = work < 132; + const int item = quad ? work : work - 132; + if (quad && item < 84) { + tile_id = 2 * (item / 7); + tok = 2 * (item % 7); + } else if (quad) { + const int tail = item - 84; + tile_id = 2 * (12 + tail / 6); + tok = 2 * (tail % 6); + } else if (item < 48) { + tile_id = 2 * (item / 4) + (item % 4) / 2; + tok = 14 + item % 2; + } else { + const int tail = item - 48; + tile_id = 2 * (12 + tail / 8) + (tail % 8) / 4; + tok = 12 + tail % 4; + } + mt = quad ? MT : 1; + nm = quad ? NM_MAX : 1; + } else if (split72) { + if (work < 264) { + tile_id = 2 * (work / 8); + tok = 2 * (work % 8); + mt = MT; nm = NM_MAX; + } else { + const int i = work - 264, c = i >> 2, q = i & 3; + tile_id = 2 * (33 + (c >> 3)) + (q >> 1); + tok = 2 * (c & 7) + (q & 1); + mt = 1; nm = 1; + } + } else if (split81) { + if (work < 264) { + tile_id = 2 * (work / 8); + tok = 2 * (work % 8); + mt = MT; nm = NM_MAX; + } else if (work < 376) { + const int i = work - 264; + tile_id = 2 * (33 + i / 16); + tok = i % 16; + mt = MT; nm = 1; + } else { + tile_id = 80; + tok = 2 * (work - 376); + mt = 1; nm = NM_MAX; + } + } else if (split17) { + if (work < 8) { + tile_id = 14; + tok = work; + mt = 3; nm = 1; + } else if (work < 120) { + const int i = work - 8; + tile_id = i >> 3; + tok = 2 * (i & 7); + mt = 1; nm = NM_MAX; + } else { + const int i = work - 120; + tile_id = 14 + (i >> 2); + tok = 8 + 2 * (i & 3); + mt = 1; nm = NM_MAX; + } + } else if (tall3) { + tile_id = 3 * (work / DECODE_TILES); + tok = work % DECODE_TILES; + mt = 3; + nm = 1; + } else { + const int j = work < q_units ? work : work - q_units; + tile_id = work < q_units ? MT * (j / QTOK) + : q_rem_base + j / QTOK; + tok = NM_MAX * (j % QTOK); + mt = work < q_units ? MT : 1; + nm = NM_MAX; + } + } else mk2_group_tile(work,grouped,tile_id,tok,nm); + } + int epi_hold = -1; + rt_fl<16, BATCH_PAD> accum[ACCN]; + if (hex && nm == 1) { + #pragma unroll + for (int u = 0; u < 3; ++u) warp::zero(accum[u]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + HEX_MT]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[1], pf[s * pslot + 1], pf[s * pslot + HEX_MT]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[2], pf[s * pslot + 2], pf[s * pslot + HEX_MT]); + warpgroup::mma_async_wait(); + if (i + 1 < iters) { + if (warpgroup::laneid() == 0) + arrive(drained[s], CONSUMER_WARPS); + } else epi_hold = s; + } + } else if (hex && mt == HEX_MT) { + #pragma unroll + for (int u = 0; u < ACCN; ++u) warp::zero(accum[u]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], + pf[s * pslot + HEX_MT]); + #pragma unroll + for (int u = 1; u < ACCN; ++u) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[u], pf[s * pslot + u / HEX_NM], + pf[s * pslot + HEX_MT + u % HEX_NM]); + warpgroup::mma_async_wait(); + if (i + 1 < iters) { + if (warpgroup::laneid() == 0) + arrive(drained[s], CONSUMER_WARPS); + } else epi_hold = s; + } + } else if (WIDE && mt == 3) { + #pragma unroll + for (int u = 0; u < 3; ++u) warp::zero(accum[u]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + 3]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[1], pf[s * pslot + 1], pf[s * pslot + 3]); + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[2], pf[s * pslot + 2], pf[s * pslot + 3]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else if (WIDE && mt == MT && nm == NM_MAX) { + #pragma unroll + for (int u = 0; u < ACCN; ++u) warp::zero(accum[u]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + MT]); + #pragma unroll + for (int u = 1; u < ACCN; ++u) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[u], pf[s * pslot + u / NM_MAX], + pf[s * pslot + MT + u % NM_MAX]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else if (MT > 1 && mt == MT) { + #pragma unroll + for (int c = 0; c < MT; ++c) warp::zero(accum[c]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + MT]); + #pragma unroll + for (int c = 1; c < MT; ++c) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[c], pf[s * pslot + c], pf[s * pslot + MT]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else if (NM_MAX == 1 || nm == NM_MAX) { + #pragma unroll + for (int m = 0; m < NM_MAX; ++m) warp::zero(accum[m]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + 1]); + #pragma unroll + for (int m = 1; m < NM_MAX; ++m) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[m], pf[s * pslot], pf[s * pslot + 1 + m]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else { + warp::zero(accum[0]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % PSTAGES; + const int phase = (iter / PSTAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], pf[s * pslot], pf[s * pslot + 1]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } + out_tile_f &e_sf = epi_hold >= 0 + ? *reinterpret_cast(&pf[epi_hold * pslot]) + : o_sf; + out_tile *const e_sp = reinterpret_cast(&e_sf); + const int nacc = mt * nm; + if (epi != EPI_BIAS) { + epi_publish(e_sp[0], accum[0]); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + } + #pragma unroll + for (int u = 0; u < ACCN; ++u) { + if (u >= mt * nm) break; + const int c = (mt > 1) ? ((nm > 1) ? (u / NM_MAX) : u) : 0; + const int m = (nm > 1) ? (u % NM_MAX) : 0; + const int base = (r6_lut ? (tile_id + c)%r6_output_tiles + : (tile_id + c)) * OUT_ATOM; + const int grow0 = (tok + m) * BATCH_PAD; + if (epi == EPI_BIAS) { + const int bw = 16 * warpgroup::warpid() + + (::kittens::laneid() >> 2); + const float b0 = __bfloat162float(bias[base + bw]); + const float b1 = __bfloat162float(bias[base + bw + 8]); + #pragma unroll + for (int j = 0; j < rt_fl<16, BATCH_PAD>::width; ++j) { + accum[u].tiles[0][j].data[0].x += b0; + accum[u].tiles[0][j].data[0].y += b0; + accum[u].tiles[0][j].data[1].x += b1; + accum[u].tiles[0][j].data[1].y += b1; + accum[u].tiles[0][j].data[2].x += b0; + accum[u].tiles[0][j].data[2].y += b0; + accum[u].tiles[0][j].data[3].x += b1; + accum[u].tiles[0][j].data[3].y += b1; + } + epi_publish(e_sp[0], accum[u]); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + constexpr int OUT_OCTS = OUT_ATOM / 8; + for (int idx = warpgroup::laneid(); idx < OUT_OCTS * rows_valid; + idx += 128) { + const int row = idx >> 3, col = (idx & 7) << 3; + const size_t at = (size_t)(grow0 + row) * dest_stride + base + col; + *reinterpret_cast(dest + at) = + *reinterpret_cast( + &e_sp[0][int2{row, col}]); + } + warpgroup::sync(1); + continue; + } + if (u + 1 < ACCN && u + 1 < nacc) + epi_publish(e_sp[(u + 1) & 1], + accum[u + 1 < ACCN ? u + 1 : 0]); + const out_tile &e_sr = e_sp[u & 1]; + constexpr int OUT_OCTS = OUT_ATOM / 8; + for (int idx = warpgroup::laneid(); idx < OUT_OCTS * rows_valid; + idx += 128) { + const int row = idx >> 3; + const int col = (idx & 7) << 3; + const int grow = grow0 + row; + const size_t at = + (size_t)grow * dest_stride + base + col; + const bf16_oct src = *reinterpret_cast( + &e_sr[int2{row, col}]); + if (epi == EPI_STORE) { + *reinterpret_cast(dest + at) = src; + continue; + } + const float v0 = __bfloat162float(src.x0); + const float v1 = __bfloat162float(src.x1); + const float v2 = __bfloat162float(src.x2); + const float v3 = __bfloat162float(src.x3); + const float v4 = __bfloat162float(src.x4); + const float v5 = __bfloat162float(src.x5); + const float v6 = __bfloat162float(src.x6); + const float v7 = __bfloat162float(src.x7); + switch (epi) { + case EPI_STORE: break; + case EPI_SILU: + *reinterpret_cast(dest + at) = bf16_oct{ + __float2bfloat16(v0 / (1.f + __expf(-v0))), + __float2bfloat16(v1 / (1.f + __expf(-v1))), + __float2bfloat16(v2 / (1.f + __expf(-v2))), + __float2bfloat16(v3 / (1.f + __expf(-v3))), + __float2bfloat16(v4 / (1.f + __expf(-v4))), + __float2bfloat16(v5 / (1.f + __expf(-v5))), + __float2bfloat16(v6 / (1.f + __expf(-v6))), + __float2bfloat16(v7 / (1.f + __expf(-v7)))}; + break; + case EPI_MOE_ACC: { + const int expert = (tile_id + c) / out_tiles_per_expert; + const int local = ((tile_id + c) % out_tiles_per_expert) * OUT_ATOM + + col; + const int original_row=r6_map?r6_map[grow]:grow; + if(original_row<0) break; + const float wt = moe_w[original_row * BB_EXPERTS + expert]; + if (wt != 0.f) { + const size_t dst = + ((size_t)original_row * (BB_EXPERTS + 1) + expert) * + D_MODEL + local; + *reinterpret_cast(moe_out + dst) = make_float4( + __fmul_rn(wt, v0), __fmul_rn(wt, v1), + __fmul_rn(wt, v2), __fmul_rn(wt, v3)); + *reinterpret_cast(moe_out + dst + 4) = make_float4( + __fmul_rn(wt, v4), __fmul_rn(wt, v5), + __fmul_rn(wt, v6), __fmul_rn(wt, v7)); + } + break; + } + case EPI_ACC_F32: + *reinterpret_cast(moe_out + + ((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * + D_MODEL + base + col) = make_float4(v0, v1, v2, v3); + *reinterpret_cast(moe_out + + ((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * + D_MODEL + base + col + 4) = make_float4(v4, v5, v6, v7); + break; + case EPI_BIAS: break; + case EPI_GELU: + *reinterpret_cast(dest + at) = bf16_oct{ + __float2bfloat16(gelu_tanh_fp32(v0)), + __float2bfloat16(gelu_tanh_fp32(v1)), + __float2bfloat16(gelu_tanh_fp32(v2)), + __float2bfloat16(gelu_tanh_fp32(v3)), + __float2bfloat16(gelu_tanh_fp32(v4)), + __float2bfloat16(gelu_tanh_fp32(v5)), + __float2bfloat16(gelu_tanh_fp32(v6)), + __float2bfloat16(gelu_tanh_fp32(v7))}; + break; + case EPI_RESIDUAL: { + const bf16_oct prior = + *reinterpret_cast(residual + at); + *reinterpret_cast(residual + at) = bf16_oct{ + __float2bfloat16(__bfloat162float(prior.x0) + v0), + __float2bfloat16(__bfloat162float(prior.x1) + v1), + __float2bfloat16(__bfloat162float(prior.x2) + v2), + __float2bfloat16(__bfloat162float(prior.x3) + v3), + __float2bfloat16(__bfloat162float(prior.x4) + v4), + __float2bfloat16(__bfloat162float(prior.x5) + v5), + __float2bfloat16(__bfloat162float(prior.x6) + v6), + __float2bfloat16(__bfloat162float(prior.x7) + v7)}; + break; + } + } + } + if (u + 1 < nacc) { + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + } + } + if (epi != EPI_BIAS) warpgroup::sync(1); + if (epi_hold >= 0 && warpgroup::laneid() == 0) + arrive(drained[epi_hold], CONSUMER_WARPS); + } + (void)my_tiles; + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr, + const unsigned char *native_predictor_active_mask = nullptr + MK2_PROGRESS_PARAMETER(native_progress_mask)) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&p_s)[STAGES][SLOTS] = alloc.allocate(); + out_tile_f (&o_sf) = alloc.allocate(); + out_tile &o_s = *reinterpret_cast(&o_sf); + bf16 *const route_stage = reinterpret_cast(&o_sf); + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_SCRATCH_STRIDE]; + static_assert(STAGES * SLOTS * sizeof(weight_tile) >= sizeof(sampler_storage)); + sampler_storage &sort_storage = *reinterpret_cast(&__shm[0]); + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS, 0); + arrive(drained[s], CONSUMER_WARPS); + } + } + int gemm_iter = 0; + int gemm_tile = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + const float *fold_src = nullptr, *fold_routes = nullptr; + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + fold_src = nullptr; + fold_routes = nullptr; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + grid_barrier(g.barrier); +#endif + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); +#endif + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, BATCH, + route_stage); + } + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + } + grid_barrier(g.barrier); + r6_sparse_map(bb->moe_w,bb->r6_map,bb->r6_counts); + grid_barrier(g.barrier); + r6_sparse_gather(bb->normed,bb->r6_packed,bb->r6_map,bb->r6_counts, + bb->r6_lut); + grid_barrier(g.barrier); + gemm(L.exp_up, bb->r6_views[0], BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0,0,nullptr,nullptr,DECODE_TILES,DECODE_ROWS, + nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->r6_views[1], BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + DECODE_TILES, DECODE_ROWS, nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + if (bb->layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_selected(bb->moe_out, bb->moe_w, idx); + grid_barrier(g.barrier); + } else { + fold_src = bb->moe_out; + fold_routes = bb->moe_w; + } + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int n_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (size_t i = size_t(blockIdx.x) * THREADS + threadIdx.x; i < size_t(14) * BATCH * CV; + i += size_t(gridDim.x) * THREADS) { + const int cell = int(i % CV); + const size_t row_layer = i / CV; + const int row = int(row_layer % BATCH); + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (size_t i = size_t(blockIdx.x) * THREADS + threadIdx.x; i < size_t(14) * BATCH * SS; + i += size_t(gridDim.x) * THREADS) { + const int cell = int(i % SS); + const size_t row_layer = i / SS; + const int row = int(row_layer % BATCH); + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = n_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + bb->pf_h[idx] = bb->prompt[idx]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = bb->pf_n_seq > 0 ? bb->pf_n_seq : 1; + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_gated, + T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, T, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * (BB_EXPERTS + 1) * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->pf_normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->pf_moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = bb_moe_fold_deterministic(bb->pf_moe_out, idx); + grid_barrier(g.barrier); + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, bb->pf_final, ln_scratch, + n_total); + grid_barrier(g.barrier); + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, bb->phoneme_b); + grid_barrier(g.barrier); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + for (int row = blockIdx.x + warp * gridDim.x; row < BATCH; + row += gridDim.x * WARPS) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + MK2_PROGRESS_GUARD(native_progress_mask, row); + const int arg = phoneme_warp_argmax(lg); + if (lane == 0) bb->phoneme_out[row] = arg; + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + __shared__ int native_run_predictor; + if (native_predictor_active_mask && bb && g.native_codes) { + if (threadIdx.x == 0) { + int any_active = 0; + for (int row = 0; row < BATCH; ++row) + any_active |= int(native_predictor_active_mask[row] != 0 && + MK2_PROGRESS_ENABLED(native_progress_mask, row)); + native_run_predictor = any_active; + } + __syncthreads(); + } + const bool skip_native_predictor = native_predictor_active_mask && bb && + g.native_codes && native_run_predictor == 0; + if (skip_native_predictor) { + for (int index = blockIdx.x * THREADS + threadIdx.x; + index < BATCH * N_POS; index += gridDim.x * THREADS) { + MK2_PROGRESS_GUARD(native_progress_mask, index / N_POS); + g.sampled[index] = 0; + g.native_codes[index] = 0; + } + grid_barrier(g.barrier); + } else { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + constexpr int DPL = D_HEAD / 32; + float qv[DPL], kn[DPL], vn[DPL]; + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + qv[j] = __bfloat162float(q[i]); + kn[j] = __bfloat162float(k_new[i]); + vn[j] = __bfloat162float(v_new[i]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += qv[j] * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + { + float acc = 0.f; + #pragma unroll + for (int j = 0; j < DPL; ++j) acc += qv[j] * kn[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[pos] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + float ov[DPL] = {}; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *vt = g.vcache + cache_base + t * D_HEAD; + const float s = sc[t]; + #pragma unroll + for (int j = 0; j < DPL; ++j) + ov[j] += s * __bfloat162float(vt[lane + j * 32]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + ov[j] += sc[pos] * vn[j]; + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + + lane + j * 32] = __float2bfloat16(ov[j]); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, + g.sampled_pos, sort_storage, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif +#ifdef ABLATE_SAMPLE + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { + __syncthreads(); + } else { + grid_barrier(g.barrier); + } +#endif + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x == 0) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + g.sampled[row * N_POS + pos] = code; + if (g.native_codes && MK2_PROGRESS_ENABLED(native_progress_mask, row)) { +#if MK2_INFLIGHT + g.native_codes[row * N_POS + pos] = + native_predictor_active_mask && !native_predictor_active_mask[row] + ? 0 : (long long)code; +#else + g.native_codes[row * N_POS + pos] = (long long)code; +#endif + } + } + } +#ifdef ABLATE_SAMPLE + grid_barrier(g.barrier); +#else + if (pos + 1 == n_pos) grid_barrier(g.barrier); +#endif + } + } + if (g.native_seq_len && blockIdx.x == 0) + for (int row = threadIdx.x; row < BATCH; row += blockDim.x) { + MK2_PROGRESS_GUARD(native_progress_mask, row); + g.native_seq_len[row] += 1; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2_startup( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr, + const unsigned char *native_predictor_active_mask = nullptr + MK2_PROGRESS_PARAMETER(native_progress_mask)) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&p_s)[STAGES][SLOTS] = alloc.allocate(); + out_tile_f (&o_sf) = alloc.allocate(); + out_tile &o_s = *reinterpret_cast(&o_sf); + bf16 *const route_stage = reinterpret_cast(&o_sf); + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_SCRATCH_STRIDE]; + static_assert(STAGES * SLOTS * sizeof(weight_tile) >= sizeof(sampler_storage)); + sampler_storage &sort_storage = *reinterpret_cast(&__shm[0]); + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS, 0); + arrive(drained[s], CONSUMER_WARPS); + } + } + int gemm_iter = 0; + int gemm_tile = 0; + __syncthreads(); + constexpr bool prefilling = false; + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + const float *fold_src = nullptr, *fold_routes = nullptr; + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + fold_src = nullptr; + fold_routes = nullptr; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + grid_barrier(g.barrier); +#endif + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); +#endif + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, BATCH, + route_stage); + } + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + } + grid_barrier(g.barrier); + r6_sparse_map(bb->moe_w,bb->r6_map,bb->r6_counts); + grid_barrier(g.barrier); + r6_sparse_gather(bb->normed,bb->r6_packed,bb->r6_map,bb->r6_counts, + bb->r6_lut); + grid_barrier(g.barrier); + gemm(L.exp_up, bb->r6_views[0], BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0,0,nullptr,nullptr,DECODE_TILES,DECODE_ROWS, + nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->r6_views[1], BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + DECODE_TILES, DECODE_ROWS, nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + if (bb->layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_selected(bb->moe_out, bb->moe_w, idx); + grid_barrier(g.barrier); + } else { + fold_src = bb->moe_out; + fold_routes = bb->moe_w; + } + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if constexpr (false) { + const int n_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (size_t i = size_t(blockIdx.x) * THREADS + threadIdx.x; i < size_t(14) * BATCH * CV; + i += size_t(gridDim.x) * THREADS) { + const int cell = int(i % CV); + const size_t row_layer = i / CV; + const int row = int(row_layer % BATCH); + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (size_t i = size_t(blockIdx.x) * THREADS + threadIdx.x; i < size_t(14) * BATCH * SS; + i += size_t(gridDim.x) * THREADS) { + const int cell = int(i % SS); + const size_t row_layer = i / SS; + const int row = int(row_layer % BATCH); + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = n_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + bb->pf_h[idx] = bb->prompt[idx]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = bb->pf_n_seq > 0 ? bb->pf_n_seq : 1; + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_gated, + T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, T, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * (BB_EXPERTS + 1) * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->pf_normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->pf_moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = bb_moe_fold_deterministic(bb->pf_moe_out, idx); + grid_barrier(g.barrier); + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, bb->pf_final, ln_scratch, + n_total); + grid_barrier(g.barrier); + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, bb->phoneme_b); + grid_barrier(g.barrier); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + for (int row = blockIdx.x + warp * gridDim.x; row < BATCH; + row += gridDim.x * WARPS) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + MK2_PROGRESS_GUARD(native_progress_mask, row); + const int arg = phoneme_warp_argmax(lg); + if (lane == 0) bb->phoneme_out[row] = arg; + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + __shared__ int native_run_predictor; + if (native_predictor_active_mask && bb && g.native_codes) { + if (threadIdx.x == 0) { + int any_active = 0; + for (int row = 0; row < BATCH; ++row) + any_active |= int(native_predictor_active_mask[row] != 0 && + MK2_PROGRESS_ENABLED(native_progress_mask, row)); + native_run_predictor = any_active; + } + __syncthreads(); + } + const bool skip_native_predictor = native_predictor_active_mask && bb && + g.native_codes && native_run_predictor == 0; + if constexpr (true) { + for (int index = blockIdx.x * THREADS + threadIdx.x; + index < BATCH * N_POS; index += gridDim.x * THREADS) { + MK2_PROGRESS_GUARD(native_progress_mask, index / N_POS); + g.sampled[index] = 0; + g.native_codes[index] = 0; + } + grid_barrier(g.barrier); + } else { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + constexpr int DPL = D_HEAD / 32; + float qv[DPL], kn[DPL], vn[DPL]; + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + qv[j] = __bfloat162float(q[i]); + kn[j] = __bfloat162float(k_new[i]); + vn[j] = __bfloat162float(v_new[i]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += qv[j] * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + { + float acc = 0.f; + #pragma unroll + for (int j = 0; j < DPL; ++j) acc += qv[j] * kn[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[pos] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + float ov[DPL] = {}; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *vt = g.vcache + cache_base + t * D_HEAD; + const float s = sc[t]; + #pragma unroll + for (int j = 0; j < DPL; ++j) + ov[j] += s * __bfloat162float(vt[lane + j * 32]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + ov[j] += sc[pos] * vn[j]; + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + + lane + j * 32] = __float2bfloat16(ov[j]); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, + g.sampled_pos, sort_storage, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif +#ifdef ABLATE_SAMPLE + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { + __syncthreads(); + } else { + grid_barrier(g.barrier); + } +#endif + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x == 0) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + g.sampled[row * N_POS + pos] = code; + if (g.native_codes && MK2_PROGRESS_ENABLED(native_progress_mask, row)) { +#if MK2_INFLIGHT + g.native_codes[row * N_POS + pos] = + native_predictor_active_mask && !native_predictor_active_mask[row] + ? 0 : (long long)code; +#else + g.native_codes[row * N_POS + pos] = (long long)code; +#endif + } + } + } +#ifdef ABLATE_SAMPLE + grid_barrier(g.barrier); +#else + if (pos + 1 == n_pos) grid_barrier(g.barrier); +#endif + } + } + if (g.native_seq_len && blockIdx.x == 0) + for (int row = threadIdx.x; row < BATCH; row += blockDim.x) { + MK2_PROGRESS_GUARD(native_progress_mask, row); + g.native_seq_len[row] += 1; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 64 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * SLOTS * sizeof(weight_tile) + + sizeof(out_tile_f) + 2048; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; + const unsigned char *native_predictor_active_mask = nullptr; + MK2_PROGRESS_MEMBER + int maximum_grid = 0; + bool grid_reset_pending = false; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * SLOTS * sizeof(weight_tile) + + sizeof(out_tile_f) + 2048), 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaFuncSetAttribute(predictor_mk2_startup, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = std::min(out->sms * out->blocks_per_sm, 4096); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 64 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->maximum_grid = info.sms * info.blocks_per_sm; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + mk2_check(owned.alloc(&hb.r6_packed,(size_t)BB_EXPERTS*BATCH*D_MODEL*sizeof(bf16))); + mk2_check(owned.alloc(&hb.r6_map,(size_t)BB_EXPERTS*BATCH*sizeof(int))); + mk2_check(owned.alloc(&hb.r6_counts,BB_EXPERTS*sizeof(int))); + mk2_check(owned.alloc(&hb.r6_lut,(1+BB_EXPERTS*DECODE_TILES)*sizeof(int))); + act_gl sparse_views[]={ + act_gl{hb.r6_packed,nullptr,nullptr,(size_t)BB_EXPERTS*BATCH,(size_t)D_MODEL}, + act_gl{moe_h,nullptr,nullptr,(size_t)BB_EXPERTS*BATCH,(size_t)BB_MOE_INTER}}; + mk2_check(owned.alloc(&hb.r6_views,sizeof(sparse_views))); + mk2_check(cudaMemcpy(hb.r6_views,sparse_views,sizeof(sparse_views),cudaMemcpyHostToDevice)); + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.maximum_grid = pred->maximum_grid; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native, bool startup = false) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(s, &capture_status)); + if (st->pred.grid_reset_pending || capture_status != cudaStreamCaptureStatusNone) { + mk2_check(cudaMemsetAsync(st->pred.g.barrier, 0, 64 * sizeof(unsigned int), s)); + st->pred.grid_reset_pending = false; + } + if (startup) { + predictor_mk2_startup<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS, nullptr, + st->pred.native_predictor_active_mask + MK2_PROGRESS_ARGUMENT(st->pred.native_progress_mask)); + } else { + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS, nullptr, + native ? st->pred.native_predictor_active_mask : nullptr + MK2_PROGRESS_ARGUMENT(native ? st->pred.native_progress_mask : nullptr)); + } + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" int mk2_run_native_startup_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream) { + return mk2_enqueue_stage0(p, b, stream, true, true); +} +extern "C" int mk2_set_native_predictor_active_mask(MK2Predictor *p, + const unsigned char *mask) { + try { + if (!p || reinterpret_cast(p)->kind != 1) + throw std::runtime_error("startup mask requires a full Stage0 handle"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && st->pred.native_predictor_active_mask != mask) + throw std::runtime_error("startup mask pointer cannot change after first Stage0 enqueue"); + if (mask) + mk2_stage0_device_pointer(mask, st->pred.device, "native_predictor_active_mask"); + st->pred.native_predictor_active_mask = mask; + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown startup mask binding failure"; return 1; } +} +extern "C" int mk2_native_startup_skip_abi_version() { return 1; } +#include "../../../../../inflight/runtime.cuh" +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "native_prefill_impl.inc" +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *p, int grid) { + if (!p) { g_last_error = "grid setter requires full Stage0 handle"; return 1; } + auto *st = reinterpret_cast(p); + if (st->kind != 1 || grid < 1 || grid > st->maximum_grid) { + g_last_error = "requested Stage0 grid outside independently resident capacity"; + return 1; + } + if (st->grid != grid) st->grid_reset_pending = true; + st->grid = grid; + g_last_error.clear(); + return 0; +} +extern "C" int mk2_r8_bind_initial_state(MK2Predictor *p, const void *const *initial, const void *ready, const void *seq_len, int positions) { + try { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || st->bindings_initialized || !initial || !ready || !seq_len || positions <= 0) + throw std::runtime_error("invalid immutable initial-state binding"); + for (int i = 0; i < 14; ++i) { + mk2_stage0_device_pointer(initial[i], st->pred.device, "initial state"); + st->mb_host[i].r8_initial_ssm = static_cast(initial[i]); + st->mb_host[i].r8_initial_ready = static_cast(ready); + st->mb_host[i].r8_seq_len = static_cast(seq_len); + st->mb_host[i].r8_initial_positions = positions; + } + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" int mk2_r9_state_bind(void* p,const void*const* scales,const void*const* maps) {auto*st=(MK2Stage0State*)p;if(st->bindings_initialized)return 1;for(int i=0;i<14;++i){st->mb_host[i].r9_scales=(const float*)scales[i];st->mb_host[i].r9_map=(const int*)maps[i];}return 0;} +extern "C" int mk2_r9_state_bytes(){return sizeof(r9_state_t);} diff --git a/deploy/nemotron_native/kernels/b4096/regions/stage0/stage0.cu b/deploy/nemotron_native/kernels/b4096/regions/stage0/stage0.cu new file mode 100644 index 0000000..115086c --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/stage0/stage0.cu @@ -0,0 +1,35 @@ +#define MK2_BATCH 4096 +#define STAGES 3 +#define KCHUNK 128 +#define MK2_NO_MAIN 1 +#define MK2_DIRECT_INPUT 1 +#define MK2_B1_LOCAL_BARRIERS 1 +#define MK2_SKIP_HIDDEN_SNAPSHOT 1 +#define MK2_ACTIVE_EXPERTS 1 +#define MK2_SSM_SPLITS 1 +#define NMULT 2 +#define MK2_OUT_MULT 2 +#define SSM_PREFETCH_MAX 12 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#if MK2_INFLIGHT +#include +#endif +#include +#include +#include +#include +#include +#include +#include +#include +namespace b4096_stage0 { +#include "native/stage0/predictor_mk2.cu" +} diff --git a/deploy/nemotron_native/kernels/b4096/regions/stage1/codec.cu b/deploy/nemotron_native/kernels/b4096/regions/stage1/codec.cu new file mode 100644 index 0000000..7cf7265 --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/stage1/codec.cu @@ -0,0 +1,4 @@ +#define NATIVE_CODEC_PREFETCH 1 +#define NATIVE_CODEC_PACK_SUBCTAS 1 +#define D7W 96 +#include "native/codec/native_stage1.cu" diff --git a/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/kernel_kf_codec.h b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/kernel_kf_codec.h new file mode 100644 index 0000000..9a8b4d5 --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/kernel_kf_codec.h @@ -0,0 +1,13 @@ +#pragma once +#include +struct CodecArgs { + const void *inp[41]; + const void *st[26]; + void *ns[26]; + void *waveform; + void *scratch; + int B; + int F; +}; +void codec_launch(const CodecArgs &a, void *stream); +size_t codec_scratch_bytes(int B, int F); diff --git a/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/legacy_codec.cuh b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/legacy_codec.cuh new file mode 100644 index 0000000..b360a3c --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/legacy_codec.cuh @@ -0,0 +1,1387 @@ +#include "kernel_kf_codec.h" +#include +#include +#include +#include +using namespace nvcuda; +template __device__ void prepare_attribute(F,int) {} +#ifndef TC2_QPT +#define TC2_QPT 4 +#endif +#define DIVUP(a,b) (((a) + (b) - 1) / (b)) +__device__ __forceinline__ float half_snake(float v, int c, int half, + const float *__restrict__ alpha) { + if (c < half) { + const float a = alpha[c]; + const float s = __sinf(a * v); + return __fmaf_rn(s, s / (a + 1e-9f), v); + } + return v > 0.f ? v : 0.01f * v; +} +#ifndef WG_ST3 +#define WG_ST3 1 +#define WG_MB3 3 +#define WG_ST7 1 +#define WG_MB7 3 +#define WG_BN3 128 +#define WG_BN7 128 +#endif +#define PDL_ENTER() do { (void)0; if (threadIdx.x == 0) (void)0; } while (0) +#define PDL_TRIGGER() do { if (threadIdx.x == 0) (void)0; } while (0) +template +__device__ __forceinline__ void pdl_launch(KFn kern, dim3 grid, dim3 blk, size_t shmem, + cudaStream_t s, KArgs... args) { + native_invoke(kern, grid, blk, shmem, args...); +} +template +__device__ __forceinline__ void pdl_launch_pf(KFn kern, dim3 grid, dim3 blk, size_t shmem, + cudaStream_t s, const void *pf, size_t pfb, KArgs... args) { + native_invoke_pf(kern, grid, blk, shmem, pf, pfb, args...); +} +#include "wg_kernel.inc" +#include "wg2_kernel.inc" +#include "wg3_kernel.inc" +#include "wg4_step_widths.inc" +#include "wg5_step_widths.inc" +#include "repack3.inc" +#include "tail_kernel.inc" +template +__device__ __noinline__ void dense_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + PDL_TRIGGER(); + constexpr int BK = 16 * KT; + constexpr int NT = NWARPS * 32; + constexpr int WTM = BM / WM; + constexpr int WTN = BN / WN; + constexpr int FM = WTM / 16, FN = WTN / 16; + constexpr int LDA = BK + 8, LDB = BN + 8; + constexpr int GEMM_SM = BM * LDA + BK * LDB; + constexpr int EPI_SM = NWARPS * WTM * WTN * 2; + constexpr int SM_H = GEMM_SM > EPI_SM ? GEMM_SM : EPI_SM; + __half *smem=reinterpret_cast<__half*>(native_arena()); + char *colok=reinterpret_cast(native_aux()); + __half *As = smem; + __half *Bs = smem + BM * LDA; + float *Cw = reinterpret_cast(smem); + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int wm = warp / WN, wn = warp % WN; + const int m0 = blockIdx.x * BM; + const int n0 = blockIdx.y * BN; + const int H = KT - 1; + const int Kg = Cin * KT; + if (tid < BN) { + const int n = n0 + tid; + colok[tid] = (n < N && (n % L) < T) ? 1 : 0; + } + wmma::fragment acc[FM][FN]; +#pragma unroll + for (int i = 0; i < FM; ++i) +#pragma unroll + for (int j = 0; j < FN; ++j) wmma::fill_fragment(acc[i][j], 0.f); + constexpr int ATOT = BM * (BK / 8); + constexpr int AREG = (ATOT + NT - 1) / NT; + constexpr int BTOT = 16 * (BN / 8); + constexpr int BREG = (BTOT + NT - 1) / NT; + int4 ar[AREG]; + int4 br0[BREG], br1[BREG]; +#define CODEC_LOADG(kg0) { const int _ci0 = (kg0) / KT; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * NT; ar[r] = make_int4(0, 0, 0, 0); if (ATOT % NT == 0 || idx < ATOT) { const int m = idx / (BK / 8), v = idx % (BK / 8); const int co = m0 + m, kg = (kg0) + v * 8; if (co < Cout && kg < Kg) ar[r] = *reinterpret_cast(w + (size_t)co * Kg + kg); } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * NT; br0[r] = make_int4(0, 0, 0, 0); br1[r] = make_int4(0, 0, 0, 0); if (BTOT % NT == 0 || idx < BTOT) { const int cil = idx / (BN / 8), nc = idx % (BN / 8); const int ci = _ci0 + cil; if (ci < Cin) { const int4 *src = reinterpret_cast( in + (size_t)ci * SCin + n0 + nc * 8); br0[r] = src[0]; br1[r] = src[1]; } } } } + CODEC_LOADG(0) + for (int kg0 = 0; kg0 < Kg; kg0 += BK) { + __syncthreads(); +#pragma unroll + for (int r = 0; r < AREG; ++r) { + const int idx = tid + r * NT; + if (ATOT % NT == 0 || idx < ATOT) + *reinterpret_cast(As + (idx / (BK / 8)) * LDA + + (idx % (BK / 8)) * 8) = ar[r]; + } +#pragma unroll + for (int r = 0; r < BREG; ++r) { + const int idx = tid + r * NT; + if (BTOT % NT != 0 && idx >= BTOT) continue; + const int cil = idx / (BN / 8), nc = idx % (BN / 8); + unsigned int u[8]; + u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; + u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; +#pragma unroll + for (int k = 0; k < KT; ++k) { + const int kh = k >> 1; + int4 o; + if (k & 1) { + o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); + o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); + o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); + o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); + } else { + o.x = u[kh + 0]; o.y = u[kh + 1]; + o.z = u[kh + 2]; o.w = u[kh + 3]; + } + *reinterpret_cast(Bs + (cil * KT + k) * LDB + nc * 8) = o; + } + } + __syncthreads(); + if (kg0 + BK < Kg) CODEC_LOADG(kg0 + BK) +#pragma unroll + for (int kf = 0; kf < BK; kf += 16) { + wmma::fragment af[FM]; + wmma::fragment bf[FN]; +#pragma unroll + for (int i = 0; i < FM; ++i) + wmma::load_matrix_sync(af[i], As + (wm * WTM + i * 16) * LDA + kf, LDA); +#pragma unroll + for (int j = 0; j < FN; ++j) + wmma::load_matrix_sync(bf[j], Bs + kf * LDB + wn * WTN + j * 16, LDB); +#pragma unroll + for (int i = 0; i < FM; ++i) +#pragma unroll + for (int j = 0; j < FN; ++j) + wmma::mma_sync(acc[i][j], af[i], bf[j], acc[i][j]); + } + } +#undef CODEC_LOADG + __syncthreads(); + float *Cme = Cw + warp * WTM * WTN; +#pragma unroll + for (int i = 0; i < FM; ++i) +#pragma unroll + for (int j = 0; j < FN; ++j) + wmma::store_matrix_sync(Cme + (i * 16) * WTN + j * 16, acc[i][j], WTN, + wmma::mem_row_major); + __syncwarp(); + const int half = Cout >> 1; + const int lane = tid & 31; + for (int e = lane; e < WTM * WTN; e += 32) { + const int mm = e / WTN, nn = e % WTN; + const int co = m0 + wm * WTM + mm; + const int nl = wn * WTN + nn; + if (co >= Cout || !colok[nl]) continue; + const int n = n0 + nl; + float v = Cme[e] + bias[co]; + if (resid != nullptr) + v += __half2float(resid[(size_t)co * SCout + n + H]); + if (do_snake) v = half_snake(v, co, half, alpha); + out[(size_t)co * SCout + n + H] = __float2half(v); + } + if (blockIdx.y == 0) { + const int nbx = gridDim.x; + if (ns != nullptr) { + const int B = N / L; + const int total = Cin * B * H; + for (int idx = blockIdx.x * NT + tid; idx < total; idx += nbx * NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int B = N / L; + const int total = Cout * B * pad_len; + for (int idx = blockIdx.x * NT + tid; idx < total; idx += nbx * NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} +template +__device__ __noinline__ void transpose_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + __half *__restrict__ out, const __half *__restrict__ pad_src, + __half *__restrict__ ns, int Cin, int Cout, int Tin, + int Lin, int SCin, int jbase, int Lout, int SCout, int Hout, int pad_off, + int pad_len, int B) { + if constexpr (PDL) PDL_ENTER(); + constexpr int Ktaps = 2 * STRIDE; + constexpr int QPT = QPT_; + constexpr int VPT = QPT * STRIDE; + static_assert(QPT % 2 == 0, "half2 input loads need an even frame count"); + static_assert(VPT % 2 == 0, "half2 staging needs an even slot"); + float *ws=reinterpret_cast(native_aux()); + __half *sm=reinterpret_cast<__half*>(native_arena()); + const int co = blockIdx.y; + const int b = blockIdx.z; + const int tid = threadIdx.x, BLK = blockDim.x; + for (int i = tid; i < G * Ktaps; i += BLK) + ws[i] = __half2float(w[(size_t)(co * G + i / Ktaps) * Ktaps + (i % Ktaps)]); + __syncthreads(); + const float bs = bias[0 + co]; + const int half = Cout >> 1; + const bool snake = co < half; + const float av = snake ? alpha[co] : 0.f; + const float ri = snake ? __frcp_rn(av + 1e-9f) : 0.f; + const __half *ic = in + (size_t)(co * G) * SCin + b * Lin + jbase; + __half *oc = out + (size_t)co * SCout + b * Lout + Hout; + constexpr bool VEC16 = (VPT % 8) == 0; + const int ophase = + VEC16 ? 0 : (int)((((uintptr_t)oc) >> 1) & 7u); + __half *const sbase = sm + ophase; + const int hd = (8 - ophase) & 7; + const int lane = tid & 31; + for (int qb = 1 + blockIdx.x * BLK * QPT; qb <= Tin; + qb += gridDim.x * BLK * QPT) { + const int q = qb + QPT * tid; + const bool a0 = (q <= Tin); + float jq[QPT][G], jp0[G]; +#pragma unroll + for (int g = 0; g < G; ++g) { +#pragma unroll + for (int e2 = 0; e2 < QPT / 2; ++e2) { + const __half2 v = a0 ? *reinterpret_cast( + ic + (size_t)g * SCin + q + 2 * e2) + : __floats2half2_rn(0.f, 0.f); + jq[2 * e2][g] = __half2float(v.x); + jq[2 * e2 + 1][g] = __half2float(v.y); + } + } +#pragma unroll + for (int g = 0; g < G; ++g) { + const float up = __shfl_up_sync(0xffffffffu, jq[QPT - 1][g], 1); + jp0[g] = (lane == 0) + ? (a0 ? __half2float(ic[(size_t)g * SCin + q - 1]) : 0.f) + : up; + } + __half *const sp = sbase + tid * VPT; + __align__(16) __half vreg[VPT]; +#pragma unroll + for (int r = 0; r < STRIDE; ++r) { + float wc[G], wp[G]; +#pragma unroll + for (int g = 0; g < G; ++g) { + wc[g] = ws[g * Ktaps + r]; + wp[g] = ws[g * Ktaps + r + STRIDE]; + } +#pragma unroll + for (int e = 0; e < QPT; ++e) { + float v = bs; +#pragma unroll + for (int g = 0; g < G; ++g) { + const float cur = jq[e][g]; + const float prv = (e == 0) ? jp0[g] : jq[e - 1][g]; + v = __fmaf_rn(wc[g], cur, __fmaf_rn(wp[g], prv, v)); + } + if (snake) { + const float sn = __sinf(av * v); + v = __fmaf_rn(sn * sn, ri, v); + } else { + v = v > 0.f ? v : 0.01f * v; + } + vreg[e * STRIDE + r] = __float2half(v); + } + } + if constexpr (VEC16) { +#pragma unroll + for (int m = 0; m < VPT / 8; ++m) + *reinterpret_cast(sp + m * 8) = + *reinterpret_cast(vreg + m * 8); + } else { +#pragma unroll + for (int i = 0; i < VPT; ++i) sp[i] = vreg[i]; + } + __syncthreads(); + const int nval = min(BLK * QPT, Tin - qb + 1); + const int tot = nval * STRIDE; + __half *dst = oc + (size_t)(qb - 1) * STRIDE; + if constexpr (VEC16) { + const int tot2 = tot >> 1; + for (int i = tid; i < tot2; i += BLK) + reinterpret_cast<__half2 *>(dst)[i] = + reinterpret_cast(sbase)[i]; + if ((tot & 1) && tid == 0) dst[tot - 1] = sbase[tot - 1]; + } else { + const int head = hd < tot ? hd : tot; + for (int i = tid; i < head; i += BLK) dst[i] = sbase[i]; + const int nb4 = (tot - head) >> 3; + const int4 *s4 = reinterpret_cast(sbase + head); + int4 *d4 = reinterpret_cast(dst + head); + for (int i = tid; i < nb4; i += BLK) d4[i] = s4[i]; + for (int i = head + (nb4 << 3) + tid; i < tot; i += BLK) + dst[i] = sbase[i]; + } + __syncthreads(); + } + if (blockIdx.x == 0) { + if (ns != nullptr) { + for (int ci = co + tid * Cout; ci < Cin; ci += BLK * Cout) + ns[(size_t)b * Cin + ci] = + in[(size_t)ci * SCin + b * Lin + jbase + native_true_tokens(b,Tin)]; + } + if (pad_src != nullptr) { + for (int j = tid; j < pad_len; j += BLK) + out[(size_t)co * SCout + b * Lout + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } +} +template +__device__ __noinline__ void transpose_conv_early_bpack_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + __half *__restrict__ out, const __half *__restrict__ pad_src, + __half *__restrict__ ns, int Cin, int Cout, int Tin, + int Lin, int SCin, int jbase, int Lout, int SCout, int Hout, int pad_off, + int pad_len, int B) { + if constexpr (PDL) PDL_ENTER(); + static_assert(BPACK <= 32); + static_assert(LPB == 8 || LPB == 16 || LPB == 32); + constexpr int Ktaps = 2 * STRIDE; + constexpr int WARP = 32; + constexpr int NT = WARP * BPACK; + constexpr int BPB = WARP / LPB; + float *ws=reinterpret_cast(native_aux()); + __half *sm=reinterpret_cast<__half*>(native_arena()); + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int lane = tid & 31; + const int sub = lane % LPB; + const int bsub = lane / LPB; + const int co = blockIdx.y; + const int b = blockIdx.z * (BPACK * BPB) + warp * BPB + bsub; + const bool active = b < B; + for (int i = tid; i < G * Ktaps; i += NT) + ws[i] = __half2float(w[(size_t)(co * G + i / Ktaps) * Ktaps + + (i % Ktaps)]); + __syncthreads(); + const int qb = 1 + blockIdx.x * LPB; + const int q = qb + sub; + __half *sp = sm + (warp * WARP + lane) * STRIDE; + const bool live = active && q <= Tin; + { + const float bs = bias[co]; + const int half = Cout >> 1; + const bool snake = co < half; + const float av = snake ? alpha[co] : 0.f; + const float ri = snake ? __frcp_rn(av + 1e-9f) : 0.f; + const __half *ic = in + (size_t)(co * G) * SCin + b * Lin + jbase; + float jq[G], jp[G]; +#pragma unroll + for (int g = 0; g < G; ++g) + jq[g] = live ? __half2float(ic[(size_t)g * SCin + q]) : 0.f; +#pragma unroll + for (int g = 0; g < G; ++g) { + const float up = __shfl_up_sync(0xffffffffu, jq[g], 1, LPB); + jp[g] = (sub == 0) + ? (live ? __half2float(ic[(size_t)g * SCin + q - 1]) : 0.f) + : up; + } + if (live) +#pragma unroll + for (int r = 0; r < STRIDE; ++r) { + float v = bs; +#pragma unroll + for (int g = 0; g < G; ++g) + v = __fmaf_rn(ws[g * Ktaps + r], jq[g], + __fmaf_rn(ws[g * Ktaps + r + STRIDE], jp[g], v)); + if (snake) { + const float sn = __sinf(av * v); + v = __fmaf_rn(sn * sn, ri, v); + } else { + v = v > 0.f ? v : 0.01f * v; + } + sp[r] = __float2half(v); + } + } + __syncwarp(); + if (active) { + const int nval = min(LPB, Tin - (qb - 1)); + const int tot2 = (nval * STRIDE) >> 1; + const __half *src = sm + (warp * WARP + bsub * LPB) * STRIDE; + __half *dst = out + (size_t)co * SCout + b * Lout + Hout + + (size_t)(qb - 1) * STRIDE; + for (int i = sub; i < tot2; i += LPB) + reinterpret_cast<__half2 *>(dst)[i] = + reinterpret_cast(src)[i]; + if ((nval * STRIDE) & 1) + if (sub == 0) dst[nval * STRIDE - 1] = src[nval * STRIDE - 1]; + } + if (blockIdx.x == 0 && active) { + if (ns != nullptr) { + for (int ci = co + sub * Cout; ci < Cin; ci += LPB * Cout) + ns[(size_t)b * Cin + ci] = + in[(size_t)ci * SCin + b * Lin + jbase + native_true_tokens(b,Tin)]; + } + if (pad_src != nullptr) { + for (int j = sub; j < pad_len; j += LPB) + out[(size_t)co * SCout + b * Lout + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } +} +template +__device__ __noinline__ void transpose_conv_bpack2_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + __half *__restrict__ out, const __half *__restrict__ pad_src, + __half *__restrict__ ns, int Cin, int Cout, int Tin, + int Lin, int SCin, int jbase, int Lout, int SCout, int Hout, int pad_off, + int pad_len, int B) { + if constexpr (PDL) PDL_ENTER(); + static_assert(BPACK <= 32); + static_assert(LPB == 8 || LPB == 16 || LPB == 32); + constexpr int Ktaps = 2 * STRIDE; + constexpr int WARP = 32; + constexpr int NT = WARP * BPACK; + constexpr int BPB = WARP / LPB; + constexpr int QPT = 2; + constexpr int VPT = QPT * STRIDE; + float *ws=reinterpret_cast(native_aux()); + __half *sm=reinterpret_cast<__half*>(native_arena()); + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int lane = tid & 31; + const int sub = lane % LPB; + const int bsub = lane / LPB; + const int co = blockIdx.y; + const int b = blockIdx.z * (BPACK * BPB) + warp * BPB + bsub; + const bool active = b < B; + for (int i = tid; i < G * Ktaps; i += NT) + ws[i] = __half2float(w[(size_t)(co * G + i / Ktaps) * Ktaps + + (i % Ktaps)]); + __syncthreads(); + const int qb = 1 + blockIdx.x * (LPB * QPT); + const int q = qb + sub * QPT; + __half *sp = sm + (warp * WARP + lane) * VPT; + const bool a0 = active && (q <= Tin); + { + const float bs = bias[co]; + const int half = Cout >> 1; + const bool snake = co < half; + const float av = snake ? alpha[co] : 0.f; + const float ri = snake ? __frcp_rn(av + 1e-9f) : 0.f; + const __half *ic = in + (size_t)(co * G) * SCin + b * Lin + jbase; + float jq0[G], jq1[G], jp0[G]; +#pragma unroll + for (int g = 0; g < G; ++g) { + const __half2 v = a0 ? *reinterpret_cast( + ic + (size_t)g * SCin + q) + : __floats2half2_rn(0.f, 0.f); + jq0[g] = __half2float(v.x); + jq1[g] = __half2float(v.y); + } +#pragma unroll + for (int g = 0; g < G; ++g) { + const float up = __shfl_up_sync(0xffffffffu, jq1[g], 1, LPB); + jp0[g] = (sub == 0) + ? (a0 ? __half2float(ic[(size_t)g * SCin + q - 1]) + : 0.f) + : up; + } + if (a0) { + float vv[VPT]; +#pragma unroll + for (int r = 0; r < STRIDE; ++r) { + float wc[G], wp[G]; +#pragma unroll + for (int g = 0; g < G; ++g) { + wc[g] = ws[g * Ktaps + r]; + wp[g] = ws[g * Ktaps + r + STRIDE]; + } +#pragma unroll + for (int e = 0; e < QPT; ++e) { + float v = bs; +#pragma unroll + for (int g = 0; g < G; ++g) { + const float cur = (e == 0) ? jq0[g] : jq1[g]; + const float prv = (e == 0) ? jp0[g] : jq0[g]; + v = __fmaf_rn(wc[g], cur, __fmaf_rn(wp[g], prv, v)); + } + if (snake) { + const float sn = __sinf(av * v); + v = __fmaf_rn(sn * sn, ri, v); + } else { + v = v > 0.f ? v : 0.01f * v; + } + vv[e * STRIDE + r] = v; + } + } + __half2 *sp2 = reinterpret_cast<__half2 *>(sp); +#pragma unroll + for (int i = 0; i < VPT / 2; ++i) + sp2[i] = __floats2half2_rn(vv[2 * i], vv[2 * i + 1]); + } + } + __syncwarp(); + if (active) { + const int nval = min(LPB * QPT, Tin - (qb - 1)); + const int tot = nval * STRIDE; + const __half *src = sm + (warp * WARP + bsub * LPB) * VPT; + __half *dst = out + (size_t)co * SCout + b * Lout + Hout + + (size_t)(qb - 1) * STRIDE; + for (int i = sub; i < (tot >> 1); i += LPB) + reinterpret_cast<__half2 *>(dst)[i] = + reinterpret_cast(src)[i]; + if (tot & 1) + if (sub == 0) dst[tot - 1] = src[tot - 1]; + } + if (blockIdx.x == 0 && active) { + if (ns != nullptr) { + for (int ci = co + sub * Cout; ci < Cin; ci += LPB * Cout) + ns[(size_t)b * Cin + ci] = + in[(size_t)ci * SCin + b * Lin + jbase + native_true_tokens(b,Tin)]; + } + if (pad_src != nullptr) { + for (int j = sub; j < pad_len; j += LPB) + out[(size_t)co * SCout + b * Lout + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } +} +#define TCONV_BPACK(G,STRIDE,BPACK,LPB,TIN,COUT,BB,...) pdl_launch(transpose_conv_early_bpack_kernel, dim3(DIVUP(TIN, LPB), COUT, DIVUP(BB, BPACK * (32 / LPB))), dim3(32 * BPACK), (size_t)(32 * (BPACK)) * (STRIDE) * 2 + 64, s, __VA_ARGS__) +#define TCONV_BP2(G,STRIDE,LPB,TIN,COUT,BB,...) pdl_launch(transpose_conv_bpack2_kernel, dim3(DIVUP(TIN, (LPB) * 2), COUT, DIVUP(BB, 32)), dim3(32 * (LPB)), (size_t)(32 * (LPB)) * (4 * (STRIDE)) + 64, s, __VA_ARGS__) +#define TCONV_BPACK_DISPATCH1(G,STRIDE,BPACK,TIN,COUT,BB,...) do { if ((TIN) >= 32) TCONV_BP2(G, STRIDE, 16, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) >= 16) TCONV_BP2(G, STRIDE, 8, TIN, COUT, BB, __VA_ARGS__); else TCONV_BPACK(G, STRIDE, BPACK, 16, TIN, COUT, BB, __VA_ARGS__); } while (0) +#define TCONV_BPACK_DISPATCH(G,STRIDE,BPACK,TIN,COUT,BB,...) do { if ((TIN) >= 32) TCONV_BP2(G, STRIDE, 16, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) >= 16) TCONV_BP2(G, STRIDE, 8, TIN, COUT, BB, __VA_ARGS__); else TCONV_BPACK(G, STRIDE, BPACK, 16, TIN, COUT, BB, __VA_ARGS__); } while (0) +#define TCONV_BPACK_DISPATCH8(G,STRIDE,BPACK,TIN,COUT,BB,...) do { if ((TIN) >= 32) TCONV_BP2(G, STRIDE, 16, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) >= 16) TCONV_BP2(G, STRIDE, 8, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) <= 8) TCONV_BPACK(G, STRIDE, BPACK, 8, TIN, COUT, BB, __VA_ARGS__); else TCONV_BPACK(G, STRIDE, BPACK, 16, TIN, COUT, BB, __VA_ARGS__); } while (0) +__device__ __noinline__ void pack_latent_kernel(const __half *__restrict__ latent, + const __half *__restrict__ state, + __half *__restrict__ out, int Cin, int T, + int L, int SC, int H, int B) { + const int total = Cin * B * L; + for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; + idx += gridDim.x * blockDim.x) { + const int j = idx % L, r = idx / L; + const int b = r % B, ci = r / B; + out[(size_t)ci * SC + b * L + j] = + (j < H) ? state[((size_t)b * Cin + ci) * H + j] + : latent[((size_t)b * Cin + ci) * T + (j - H)]; + } +} +#define PC_VEC 8 +__device__ __noinline__ void post_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, __half *__restrict__ y, + __half *__restrict__ ns, int Cin, int T, int L, int SC, int KT, int B, + int N) { + float *pw=reinterpret_cast(native_arena()); + for (int i = threadIdx.x; i < Cin * KT; i += 256) pw[i] = __half2float(w[i]); + __syncthreads(); + const int H = KT - 1; + const float bs = bias[0]; + for (int n = (blockIdx.x * 256 + threadIdx.x) * PC_VEC; n < N; + n += gridDim.x * 256 * PC_VEC) { + float v[PC_VEC]; +#pragma unroll + for (int i = 0; i < PC_VEC; ++i) v[i] = bs; + for (int ci = 0; ci < Cin; ++ci) { + const __half *xr = in + (size_t)ci * SC + n; + const float *wr = pw + ci * KT; + const int4 lo = *reinterpret_cast(xr); + const int4 hi = *reinterpret_cast(xr + 8); + unsigned u[8]; + u[0] = lo.x; u[1] = lo.y; u[2] = lo.z; u[3] = lo.w; + u[4] = hi.x; u[5] = hi.y; u[6] = hi.z; u[7] = hi.w; + float x[PC_VEC + 6]; +#pragma unroll + for (int j = 0; j < (PC_VEC + 6) / 2; ++j) { + const float2 f = __half22float2( + *reinterpret_cast(&u[j])); + x[2 * j] = f.x; + x[2 * j + 1] = f.y; + } +#pragma unroll + for (int i = 0; i < PC_VEC; ++i) + for (int k = 0; k < 7; ++k) + if (k < KT) v[i] = __fmaf_rn(wr[k], x[i + k], v[i]); + } + const int b0 = n / L, t0 = n - (n / L) * L; +#pragma unroll + for (int i = 0; i < PC_VEC; ++i) { + int t = t0 + i, b = b0; + if (t >= L) { t -= L; b = b0 + 1; } + if (t < T && b < B) + y[(size_t)b * T + t] = + __float2half(fminf(fmaxf(v[i], -1.f), 1.f)); + } + } + if (ns != nullptr) { + const int total = B * Cin * H; + for (int idx = blockIdx.x * 256 + threadIdx.x; idx < total; + idx += gridDim.x * 256) { + const int j = idx % H, r = idx / H; + const int ci = r % Cin, b = r / Cin; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SC + b * L + native_true_tokens(b,T) + j]; + } + } +} +#define RP_JOBS 6 +struct RepackSet { + const __half *src[RP_JOBS]; + __half *dst[RP_JOBS]; + int Cout[RP_JOBS], CoutP[RP_JOBS], Cin[RP_JOBS]; + int base[RP_JOBS], nch[RP_JOBS], rs[RP_JOBS]; + int k3base; + const __half *k3src[2]; + __half *k3dst[2]; +}; +__host__ __device__ inline int native_RP_COUT(int j) {const int v[6]={768, 384, 256, 128, 64, 32};return v[j];} +__host__ __device__ inline int native_RP_CIN(int j) {const int v[6]={384, 768, 128, 256, 32, 64};return v[j];} +__host__ __device__ inline int native_RP_COUTP(int j) {const int v[6]={768, 384, 256, 128, 128, 128};return v[j];} +template +__device__ __noinline__ void repack_kernel(RepackSet js) { + const int bx = blockIdx.x; + int job, job_bx; + if (bx < 1152) { + job = 0; job_bx = 0; + } else if (bx < 2304) { + job = 1; job_bx = 1152; + } else if (bx < 2432) { + job = 2; job_bx = 2304; + } else if (bx < 2560) { + job = 3; job_bx = 2432; + } else if (bx < 2576) { + job = 4; job_bx = 2560; + } else { + job = 5; job_bx = 2576; + } + const int Cout = js.Cout[job], CoutP = js.CoutP[job], Cin = js.Cin[job]; + const int idx = (bx - job_bx) * blockDim.x + threadIdx.x; + const int co = idx / Cin, ci = idx - co * Cin; + const int c = ci / CIB, cl = ci - c * CIB; + __half *d = js.dst[job] + (size_t)c * (KT * CoutP * CIB) + + (size_t)(co >> 3) * (CIB * 8) + (cl >> 3) * 64 + (co & 7) * 8 + + (cl & 7); + const __half *sp = js.src[job] + (size_t)(co * Cin + ci) * KT; + const bool live = co < Cout; +#pragma unroll + for (int k = 0; k < KT; ++k) + d[(size_t)k * CoutP * CIB] = live ? sp[k] : __float2half(0.f); +} +template +__device__ __noinline__ void repack_rs_kernel( + const RepackSet &js, const __half *__restrict__ latent, const __half *__restrict__ latent_state, + __half *__restrict__ latent_out, int latent_C, int latent_T, int latent_L, + int latent_SC, int latent_H, int latent_B) { + PDL_TRIGGER(); + constexpr int CO_TILE = 16; + constexpr int CI_TILE = 64; + constexpr int ROW_H = CI_TILE * KT + 8; + constexpr int ROW4 = CI_TILE * KT / 8; + __half *tile=reinterpret_cast<__half*>(native_arena()); + const int bx = blockIdx.x; + { + const int idx = bx + threadIdx.x * gridDim.x; + const int total = latent_C * latent_B * latent_L; + if (idx < total) { + const int j = idx % latent_L, r = idx / latent_L; + const int b = r % latent_B, ci = r / latent_B; + latent_out[(size_t)ci * latent_SC + b * latent_L + j] = + (j < latent_H) + ? latent_state[((size_t)b * latent_C + ci) * latent_H + j] + : latent[((size_t)b * latent_C + ci) * latent_T + + (j - latent_H)]; + } + } + if (bx >= js.k3base) { + constexpr int K3 = 3; + constexpr int R3_H = CI_TILE * K3 + 8; + constexpr int R3_4 = CI_TILE * K3 / 8; + const int r = bx - js.k3base; + const int conv = r / 1152; + const int loc = r - conv * 1152; + const int Cout3 = conv ? 768 : 1536; + const int Cin3 = conv ? 1536 : 768; + const int nci3 = Cin3 / CI_TILE; + const int cot3 = loc / nci3; + if (cot3 * CO_TILE >= Cout3) return; + const int co03 = cot3 * CO_TILE; + const int ci03 = (loc - cot3 * nci3) * CI_TILE; + const int NCH3 = Cin3 >> 4; + const __half *const s3 = js.k3src[conv]; + __half *const d3 = js.k3dst[conv]; + const int t3 = threadIdx.x; +#pragma unroll + for (int n = 0; n < CO_TILE * R3_4 / 128; ++n) { + const int i = t3 + n * 128; + const int co = i / R3_4, q = i - co * R3_4; + reinterpret_cast(tile + co * R3_H)[q] = + *reinterpret_cast( + s3 + ((size_t)(co03 + co) * Cin3 + ci03) * K3 + q * 8); + } + __syncthreads(); + const int w3 = t3 >> 5, l3 = t3 & 31; + const int rA3 = l3 >> 2, q43 = l3 & 3; + const int cb3 = co03 >> 4; + __half *const db3 = d3 + ((size_t)(cb3 >> 2) * NCH3 * K3) * 1024 + + ((cb3 & 3) * 32 + l3) * 8; +#pragma unroll + for (int u = w3; u < (CI_TILE / 16) * K3; u += 4) { + const int cl = u / K3, tap = u - cl * K3; + const int cgi = (ci03 >> 4) + cl; + const __half *sa = tile + rA3 * R3_H + (cl * 16 + 2 * q43) * K3 + tap; + const __half *sb = sa + 8 * R3_H; + uint32_t o[4]; + o[0] = __half_as_ushort(sa[0]) | + ((uint32_t)__half_as_ushort(sa[K3]) << 16); + o[1] = __half_as_ushort(sb[0]) | + ((uint32_t)__half_as_ushort(sb[K3]) << 16); + o[2] = __half_as_ushort(sa[8 * K3]) | + ((uint32_t)__half_as_ushort(sa[9 * K3]) << 16); + o[3] = __half_as_ushort(sb[8 * K3]) | + ((uint32_t)__half_as_ushort(sb[9 * K3]) << 16); + *reinterpret_cast(db3 + (size_t)(cgi * K3 + tap) * 1024) = + make_int4(o[0], o[1], o[2], o[3]); + } + return; + } + int job = 0; +#pragma unroll + for (int q = 1; q < RP_JOBS; ++q) if (bx >= js.base[q]) job = q; + const int local = bx - js.base[job]; + const int Cout = js.Cout[job], Cin = js.Cin[job]; + const int nci = (Cin + CI_TILE - 1) / CI_TILE; + const int cot = local / nci; + const int co0 = cot * CO_TILE; + const int ci0 = (local - cot * nci) * CI_TILE; + const int NCH = js.nch[job]; + const int t = threadIdx.x; + const int valid4 = (min(Cin - ci0, CI_TILE) * KT) >> 3; + for (int i = t; i < CO_TILE * ROW4; i += 128) { + const int co = i / ROW4, q = i - co * ROW4; + int4 v = make_int4(0, 0, 0, 0); + if (co0 + co < Cout && q < valid4) + v = *reinterpret_cast( + js.src[job] + ((size_t)(co0 + co) * Cin + ci0) * KT + q * 8); + reinterpret_cast(tile + co * ROW_H)[q] = v; + } + __syncthreads(); + const int warp = t >> 5, lane = t & 31; + const int rA = lane >> 2, q4 = lane & 3; + const int cb = co0 >> 4; + const int mtile = cb >> 2, wsub = cb & 3; + __half *const dbase = js.dst[job] + + ((size_t)mtile * NCH * KT) * 1024 + (wsub * 32 + lane) * 8; + if (!js.rs[job]) { + const int CoutP = js.CoutP[job]; + constexpr int UN2 = 16 * (CI_TILE / 8) * KT; +#pragma unroll + for (int n = 0; n < UN2 / 128; ++n) { + const int u = t + n * 128; + const int col = u & 15, cig = (u >> 4) & (CI_TILE / 8 - 1); + const int tap = u / (16 * (CI_TILE / 8)); + const int gci = ci0 + cig * 8; + if (gci >= Cin) continue; + const int gco = co0 + col; + uint32_t o[4]; + const __half *sp = tile + col * ROW_H; +#pragma unroll + for (int e = 0; e < 4; ++e) + o[e] = __half_as_ushort(sp[(cig * 8 + 2 * e) * KT + tap]) | + ((uint32_t)__half_as_ushort( + sp[(cig * 8 + 2 * e + 1) * KT + tap]) << 16); + __half *d = js.dst[job] + + (size_t)(gci >> 4) * ((size_t)KT * CoutP * 16) + + (size_t)tap * (CoutP * 16) + (size_t)(gco >> 3) * 128 + + ((gci >> 3) & 1) * 64 + (gco & 7) * 8; + *reinterpret_cast(d) = make_int4(o[0], o[1], o[2], o[3]); + } + return; + } + constexpr int UNITS = (CI_TILE / 16) * KT; +#pragma unroll + for (int u = warp; u < UNITS; u += 4) { + const int cl = u / KT, tap = u - cl * KT; + const int cgi = (ci0 >> 4) + cl; + if (cgi >= NCH) continue; + const __half *sa = tile + rA * ROW_H + (cl * 16 + 2 * q4) * KT + tap; + const __half *sb = sa + 8 * ROW_H; + uint32_t o[4]; + o[0] = __half_as_ushort(sa[0]) | + ((uint32_t)__half_as_ushort(sa[KT]) << 16); + o[1] = __half_as_ushort(sb[0]) | + ((uint32_t)__half_as_ushort(sb[KT]) << 16); + o[2] = __half_as_ushort(sa[8 * KT]) | + ((uint32_t)__half_as_ushort(sa[9 * KT]) << 16); + o[3] = __half_as_ushort(sb[8 * KT]) | + ((uint32_t)__half_as_ushort(sb[9 * KT]) << 16); + *reinterpret_cast(dbase + (size_t)(cgi * KT + tap) * 1024) = + make_int4(o[0], o[1], o[2], o[3]); + } +} +__device__ __noinline__ void repack3_rs_kernel( + const __half *__restrict__ win, const __half *__restrict__ wsk, + __half *__restrict__ din, __half *__restrict__ dsk, + const __half *__restrict__ asrc, __half *__restrict__ adst, int SCA, + int ncol, int ntile) { + PDL_TRIGGER(); + if (blockIdx.y == 12) { + if (adst == nullptr) return; + for (int u = blockIdx.x; u < 96 * ntile; u += gridDim.x) { + const int g = u / ntile; + const int col = (u - g * ntile) * 128 + (int)threadIdx.x; + if (col >= ncol) continue; + __half v[8]; +#pragma unroll + for (int h = 0; h < 8; ++h) + v[h] = asrc[(size_t)(g * 8 + h) * SCA + col]; + *reinterpret_cast(adst + (size_t)g * ((size_t)SCA * 8) + + (size_t)col * 8) = + *reinterpret_cast(v); + } + return; + } + constexpr int KT = 3; + constexpr int CO_TILE = 16; + constexpr int CI_TILE = 64; + constexpr int ROW_H = CI_TILE * KT + 8; + constexpr int ROW4 = CI_TILE * KT / 8; + __half *tile=reinterpret_cast<__half*>(native_arena()); + const int jy = blockIdx.y; + const int conv = jy & 1, b = 5 - (jy >> 1); + const int Cout = conv ? 768 : 1536; + const int Cin = conv ? 1536 : 768; + const __half *src = (conv ? wsk : win) + (size_t)(b + 1) * Cout * Cin * KT; + __half *dst = (conv ? dsk : din) + (size_t)b * ((size_t)Cout * Cin * KT); + const int nci = Cin / CI_TILE; + const int cot = blockIdx.x / nci; + const int co0 = cot * CO_TILE; + const int ci0 = (blockIdx.x - cot * nci) * CI_TILE; + const int NCH = Cin >> 4; + const int t = threadIdx.x; +#pragma unroll + for (int n = 0; n < CO_TILE * ROW4 / 128; ++n) { + const int i = t + n * 128; + const int co = i / ROW4, q = i - co * ROW4; + reinterpret_cast(tile + co * ROW_H)[q] = + *reinterpret_cast( + src + ((size_t)(co0 + co) * Cin + ci0) * KT + q * 8); + } + __syncthreads(); + const int warp = t >> 5, lane = t & 31; + const int rA = lane >> 2, q4 = lane & 3; + const int cb = co0 >> 4; + __half *const dbase = dst + ((size_t)(cb >> 2) * NCH * KT) * 1024 + + ((cb & 3) * 32 + lane) * 8; +#pragma unroll + for (int u = warp; u < (CI_TILE / 16) * KT; u += 4) { + const int cl = u / KT, tap = u - cl * KT; + const int cgi = (ci0 >> 4) + cl; + const __half *sa = tile + rA * ROW_H + (cl * 16 + 2 * q4) * KT + tap; + const __half *sb = sa + 8 * ROW_H; + uint32_t o[4]; + o[0] = __half_as_ushort(sa[0]) | + ((uint32_t)__half_as_ushort(sa[KT]) << 16); + o[1] = __half_as_ushort(sb[0]) | + ((uint32_t)__half_as_ushort(sb[KT]) << 16); + o[2] = __half_as_ushort(sa[8 * KT]) | + ((uint32_t)__half_as_ushort(sa[9 * KT]) << 16); + o[3] = __half_as_ushort(sb[8 * KT]) | + ((uint32_t)__half_as_ushort(sb[9 * KT]) << 16); + *reinterpret_cast(dbase + (size_t)(cgi * KT + tap) * 1024) = + make_int4(o[0], o[1], o[2], o[3]); + } +} +__device__ __noinline__ void repack_activation_only(const __half *asrc,__half *adst,int SCA,int ncol,int ntile){ + if (adst == nullptr) return; + for (int u = blockIdx.x; u < 96 * ntile; u += gridDim.x) { + const int g = u / ntile; + const int col = (u - g * ntile) * 128 + (int)threadIdx.x; + if (col >= ncol) continue; + __half v[8]; +#pragma unroll + for (int h = 0; h < 8; ++h) + v[h] = asrc[(size_t)(g * 8 + h) * SCA + col]; + *reinterpret_cast(adst + (size_t)g * ((size_t)SCA * 8) + + (size_t)col * 8) = + *reinterpret_cast(v); + } + return; +} +namespace { +struct Stack { + int T, L, N, SC, H; +}; +inline __host__ __device__ int tc_blk(int Tin, int qpt = 2) { + int best = 256, bw = 1 << 30; + for (int b = 256; b >= 32; b >>= 1) { + const int per = qpt * b; + const int waste = (Tin + per - 1) / per * per - Tin; + if (waste < bw) { bw = waste; best = b; } + } + return best; +} +inline __host__ __device__ int roundup(int a, int b) { return (a + b - 1) / b * b; } +__host__ __device__ Stack mk(int B, int T, int K) { + Stack s; + s.T = T; + s.H = K - 1; + s.L = s.H + T; + s.N = B * s.L; + s.SC = roundup(s.N, 128) + 128; + return s; +} +struct Plan { + Stack s0, s1, s2, s3, s4; + __half *LB, *X0, *M0, *Y0; + __half *A, *Ac, *Bb, *M1; + __half *P, *Q, *M2; + __half *R, *S, *M3; + __half *U; + __half *W7[RP_JOBS]; + __half *W3i, *W3s; + size_t total; +}; +static constexpr int kTuneB = 32; +inline __host__ __device__ int nref(const Stack &st) { return kTuneB * st.L; } +inline __host__ __device__ bool use_rp3i(int, int F) { return (size_t)kTuneB * (4 * F + 2) > 400; } +inline __host__ __device__ bool use_rp3s(int, int F) { return (size_t)kTuneB * (4 * F + 2) > 400; } +__host__ __device__ Plan build_plan(int B, int F, __half *base) { + Plan p; + p.s0 = mk(B, 2 * F, 3); + p.s1 = mk(B, 4 * F, 3); + p.s2 = mk(B, 36 * F, 7); + p.s3 = mk(B, 252 * F, 7); + p.s4 = mk(B, 1764 * F, 7); + size_t off = 0; + auto alloc = [&](int ch, const Stack &s) -> __half * { + __half *r = base ? base + off : nullptr; + off += (size_t)ch * s.SC; + return r; + }; + p.LB = alloc(40, p.s0); + p.X0 = alloc(768, p.s0); + p.M0 = alloc(1536, p.s0); + p.Y0 = alloc(768, p.s0); + p.A = alloc(768, p.s1); + p.Ac = use_rp3i(B, F) ? alloc(768, p.s1) : nullptr; + p.Bb = alloc(768, p.s1); + p.M1 = alloc(1536, p.s1); + p.P = alloc(384, p.s2); + p.Q = alloc(384, p.s2); + p.M2 = alloc(768, p.s2); + p.R = alloc(128, p.s3); + p.S = alloc(128, p.s3); + p.M3 = alloc(256, p.s3); + p.U = alloc(32, p.s4); + for (int j = 0; j < RP_JOBS; ++j) { + p.W7[j] = base ? base + off : nullptr; + off += (size_t)native_RP_COUTP(j) * native_RP_CIN(j) * 7; + } + p.W3i = nullptr; + if (use_rp3i(B, F)) { + p.W3i = base ? base + off + (size_t)3 * 1536 * 768 : nullptr; + off += (size_t)7 * 3 * 1536 * 768; + } + p.W3s = nullptr; + if (use_rp3s(B, F)) { + p.W3s = base ? base + off + (size_t)3 * 768 * 1536 : nullptr; + off += (size_t)7 * 3 * 768 * 1536; + } + p.total = off; + return p; +} +__host__ __device__ const __half *H16(const void *p) { return static_cast(p); } +__host__ __device__ const float *F32(const void *p) { return static_cast(p); } +} +size_t codec_scratch_bytes(int B, int F) { + return build_plan(B, F, nullptr).total * sizeof(__half); +} +template +__device__ __forceinline__ void wg_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, bool prep = false) { + constexpr int WG_BM = 64 * NWG, WG_NT = 128 * NWG, WG_CLD = WG_BN + 8; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int GEMM_B = STAGES * (A_E + B_E) * 2; + constexpr int EPI_B = WG_BM * WG_CLD * 4; + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute(wg_conv_kernel, SH); + return; + } + native_invoke(wg_conv_kernel,dim3(dim3(DIVUP(Cout, WG_BM), DIVUP(st.N, WG_BN))),dim3(WG_NT),SH, + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg2_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, + bool prep = false) { + constexpr int BK = CIB * KT; + constexpr int WG_BM = 64 * MSPLIT; + constexpr int WG_NT = 128 * MSPLIT * KSPLIT * NSPLIT; + constexpr int WG_BNT = WG_BN * NSPLIT; + constexpr int NCRAW = WG_BNT / 4; + constexpr int EPN = + (WG_NT % NCRAW == 0) ? WG_BNT : (((WG_BNT + 127) / 128) * 128); + constexpr int WG_CLD = EPN + 8; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int GEMM_B = KSPLIT * (SA * A_E + SB * NSPLIT * B_E) * 2; + constexpr int EPI_B = KSPLIT * WG_BM * WG_CLD * 4; + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute( + wg2_conv_kernel, + SH); + return; + } + native_invoke_pf(wg2_conv_kernel,dim3(dim3(DIVUP(Cout, WG_BM), DIVUP(st.N, WG_BNT))),dim3(WG_NT),SH, + w, (size_t)KT * Cin * Cout * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg3_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, + bool prep = false) { + constexpr int WG_BM = 64 * NWG, WG_NT = 128 * NWG * NSG; + constexpr int WG_BNT = WG_BN * NSG, WG_CLD = WG_BNT + 8; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int GEMM_B = STAGES * (A_E + NSG * B_E) * 2; + constexpr int EPI_B = WG_BM * WG_CLD * 4; + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute(wg3_conv_kernel, SH); + return; + } + native_invoke(wg3_conv_kernel,dim3(dim3(DIVUP(Cout, WG_BM), DIVUP(st.N, WG_BNT))),dim3(WG_NT),SH, + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg4_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, int CoutP, + const Stack &st, int SCin, int do_snake, int pad_off, + int pad_len, bool prep = false) { + constexpr int NT = 128 * MSPLIT * NSPLIT; + constexpr int BM = 64 * MSPLIT, BNT = WG_BN * NSPLIT; + constexpr int NBR = BNT + KT - 1; + constexpr int NBRP = BC8 ? ((NBR + 7) & ~7) : ((NBR + 31) & ~31); + constexpr int NCRAW = BNT / 4; + constexpr int EPN = + (NT % NCRAW == 0) ? BNT : (((BNT + 127) / 128) * 128); + constexpr int B_E = CIB * NBRP; + constexpr int GEMM_B = STAGES * B_E * 2; + constexpr int EPI_B = BM * (EPN + 8) * (PREACT ? 2 : 4); + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute( + wg4_conv_kernel, + SH); + return; + } + pdl_launch_pf(wg4_conv_kernel, + dim3(DIVUP(Cout, BM), DIVUP(st.N, BNT)), dim3(NT), SH, s, + w, (size_t)KT * CoutP * Cin * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, + st.T, st.L, SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg5_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, + bool prep = false) { + constexpr int BM = 64 * MSPLIT; + constexpr int NT = 128 * MSPLIT * KSPLIT; + constexpr int NCRAW = WG_BN / 4; + constexpr int EPN = (NT % NCRAW == 0) ? WG_BN : (((WG_BN + 127) / 128) * 128); + constexpr int NBRP = BC8 ? (((WG_BN + 2) + 7) & ~7) + : (((WG_BN + 2) + 31) & ~31); + constexpr int B_E = CIB * NBRP; + constexpr int GEMM_B = KSPLIT * SB * B_E * 2; + constexpr int EPI_COLS = OC8 ? WG_BN : EPN; + constexpr int EPI_CLD_R = EPI_COLS + 8; + constexpr int EPI_CLD = (EPI_CLD_R % 16 == 0) ? EPI_CLD_R + 4 : EPI_CLD_R; + constexpr int EPI_B = KSPLIT * BM * EPI_CLD * (OC8 ? 2 : 4); + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + static_assert(SH <= 227 * 1024, "wg5 shared over budget"); + if (prep) { + prepare_attribute( + wg5_conv_kernel, + SH); + return; + } + if constexpr (PDL) { + cudaLaunchConfig_t cfg = {}; + cudaLaunchAttribute attr[1] = {}; + native_config_grid(cfg) = dim3(DIVUP(Cout, BM), DIVUP(st.N, WG_BN)); + native_config_threads(cfg) = dim3(NT); + cfg.dynamicSmemBytes = SH; + cfg.stream = s; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = 1; + cfg.attrs = attr; + cfg.numAttrs = 1; + native_cudaLaunchKernelEx_pf( + &cfg, + wg5_conv_kernel, + w, (size_t)3 * Cin * Cout * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); + } else { + native_invoke_pf(wg5_conv_kernel,dim3(dim3(DIVUP(Cout, BM), DIVUP(st.N, WG_BN))),dim3(NT),SH, + w, (size_t)3 * Cin * Cout * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, + st.L, SCin, st.SC, st.N, do_snake, pad_off, pad_len); + } +} +#ifndef D7W +#define D7W 96 +#endif +#define D7_M2_C8I(BN,LATE) wg4_launch<7, 16, 2, 1, D7W, 4, 1, false, 2, true, false, LATE>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_M2P_C8O(BN,LATE) wg4_launch<7, 16, 2, 1, D7W, 4, 1, true, 2, false, true, LATE>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_M2N2(BN) wg4_launch<7, 16, 2, 2, BN, 4, 1, false, 2>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_M1(BN) wg4_launch<7, 16, 1, 1, BN, 4, 2, false, 2>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_C8I(W) \ + wg4_launch<7, 16, 2, 1, W, 4, 1, false, 2, true, false, false>( \ + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, \ + SCin, 1, pad_off, pad_len, prep) +#define D7_C8O(W) \ + wg4_launch<7, 16, 2, 1, W, 4, 1, true, 2, false, true, false>( \ + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, \ + SCin, 1, pad_off, pad_len, prep) +template +__device__ __forceinline__ int d7_band_cost(int N, int mtiles, int resident) { + const int v = mtiles * DIVUP(N, BN); + const int sub = (DIVUP(v, resident) >= 2) ? 2 : 1; + return DIVUP(DIVUP(v, sub), resident) * sub * (128 + BN + 6); +} +__device__ __forceinline__ int d7_band(int N, int mtiles, int resident) { + int best = 64, bc = d7_band_cost<64>(N, mtiles, resident); + const int c96 = d7_band_cost<96>(N, mtiles, resident); + if (c96 < bc) { bc = c96; best = 96; } + if (d7_band_cost<112>(N, mtiles, resident) < bc) best = 112; + return best; +} +__device__ __forceinline__ void dense7(cudaStream_t s, int job, const __half *in, const __half *w, + const float *bias, const float *alpha, const __half *resid, + __half *out, const __half *pad_src, __half *ns, + const Stack &st, int SCin, int pad_off, int pad_len, + bool prep = false) { + const int Cin = native_RP_CIN(job), Cout = native_RP_COUT(job), CoutP = native_RP_COUTP(job); + const int b_ = d7_band(st.N, DIVUP(Cout, 128), native_resident()); + if ((job & 1) == 0) { + if (b_ == 64) D7_C8O(64); + else if (b_ == 96) D7_C8O(96); + else D7_C8O(112); + } else { + if (b_ == 64) D7_C8I(64); + else if (b_ == 96) D7_C8I(96); + else D7_C8I(112); + } +} +__device__ __forceinline__ int wg5_band_skip(int N, int mtiles, int resident) { + const int wmin = DIVUP(mtiles * DIVUP(N, 104), resident); + if (DIVUP(mtiles * DIVUP(N, 24), resident) == wmin) return 24; + if (DIVUP(mtiles * DIVUP(N, 40), resident) == wmin) return 40; + if (DIVUP(mtiles * DIVUP(N, 48), resident) == wmin) return 48; + if (DIVUP(mtiles * DIVUP(N, 64), resident) == wmin) return 64; + return 104; +} +__device__ __forceinline__ int wg5_band_main_long(int N, int mtiles, + int resident) { + const int wmin = DIVUP(mtiles * DIVUP(N, 104), resident); + if (DIVUP(mtiles * DIVUP(N, 64), resident) == wmin) return 64; + return 104; +} +__device__ __forceinline__ int wg5_band_main_short(int N, int mtiles, + int resident) { + const int wmin = DIVUP(mtiles * DIVUP(N, 96), resident); + if (DIVUP(mtiles * DIVUP(N, 56), resident) == wmin) return 56; + if (DIVUP(mtiles * DIVUP(N, 80), resident) == wmin) return 80; + return 96; +} +#define DENSE_LAUNCH(KT,BM,BN,NW,WM,WN) native_invoke(dense_conv_kernel,dim3(dim3(DIVUP(Cout, BM), DIVUP(st.N, BN))),dim3(NW * 32),0, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, SCin, st.SC, st.N, do_snake, pad_off, pad_len) +__device__ __forceinline__ void dense(cudaStream_t s, int KT, const __half *in, const __half *w, + const float *bias, const float *alpha, const __half *resid, + __half *out, const __half *pad_src, __half *ns, int Cin, + int Cout, const Stack &st, int SCin, int do_snake, + int pad_off, int pad_len, const __half *w3 = nullptr, + int layout_mode = 0, bool prep = false) { + if (KT == 3 && w3 != nullptr) { + if (Cin == 1536) { +#define WG5S(BN,O,R,PDL) wg5_launch<16, 1, 4, BN, 4, 3, 1, 1, true, O, R, 1536, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5SW(BN,O,R,PDL) wg5_launch<32, 1, 2, BN, 4, 3, 1, 1, true, O, R, 1536, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5S3(O,R) do { const int n_ = nref(st); if (n_ > 500 && n_ <= 640) { WG5SW(56, O, R, false); } else { const int b_ = wg5_band_skip(st.N, DIVUP(Cout, 64), native_resident()); if (b_ == 24) WG5S(24, O, R, true); else if (b_ == 40) WG5S(40, O, R, true); else if (b_ == 48) WG5S(48, O, R, true); else if (b_ == 64) WG5S(64, O, R, true); else WG5S(104, O, R, true); } } while (0) + if (layout_mode == 1) WG5S3(true, false); + else if (layout_mode == 2) WG5S3(true, true ); + else if (layout_mode == 3) WG5S3(false, true ); + else WG5S3(false, false); +#undef WG5S3 +#undef WG5SW +#undef WG5S + } else { +#define WG5I(BN,BC,MS,PDL) wg5_launch<16, MS, 2, BN, 4, 3, 1, 1, BC, true, false, 768, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5I4(BN,BC,PDL) wg5_launch<16, 1, 4, BN, 4, 3, 1, 1, BC, true, false, 768, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5I2W(BN,BC,PDL) wg5_launch<32, 1, 2, BN, 4, 3, 1, 1, BC, true, false, 768, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5I3(BC) do { const int n_ = nref(st); if (n_ > 640) { const int b_ = wg5_band_main_long(st.N, DIVUP(Cout, 128), native_resident()); if (b_ == 64) WG5I(64, BC, 2, true); else WG5I(104, BC, 2, true); } else if (n_ > 500) { WG5I2W(64, BC, false); } else { const int b_ = wg5_band_main_short(st.N, DIVUP(Cout, 64), native_resident()); if (b_ == 56) WG5I4(56, BC, true); else if (b_ == 80) WG5I4(80, BC, true); else WG5I4(96, BC, true); } } while (0) + if (layout_mode == 1) WG5I3(true); + else WG5I3(false); +#undef WG5I3 +#undef WG5I2W +#undef WG5I4 +#undef WG5I + } + } else if (KT == 3) { + const int n_ = nref(st); + if (Cin == 1536 && Cout == 768) { + if (n_ <= 704) + wg2_launch<3, 16, 1, 4, 1, 64, 2, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else + wg2_launch<3, 16, 1, 4, 1, 128, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + } else if (Cin == 768 && Cout == 1536) { + if (n_ > 1000) + wg2_launch<3, 16, 1, 2, 1, 128, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ > 640 && n_ <= 768) + wg2_launch<3, 16, 1, 2, 1, 64, 2, 2, 2>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ <= 320) + wg2_launch<3, 16, 1, 4, 1, 64, 2, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ > 384 && n_ <= 512) + wg2_launch<3, 16, 1, 4, 1, 112, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ > 512 && n_ <= 640) + wg2_launch<3, 16, 1, 4, 1, 120, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else + wg2_launch<3, 16, 1, 4, 1, 128, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + } else { + wg2_launch<3, 16, 1, 1, 1, 64, 3, 2, 2>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, + SCin, do_snake, pad_off, pad_len, prep); + } + } else if (Cout >= 128 && + ((Cin >= 384 && st.T == 288) || + (Cin < 384 && st.T == 2016))) { + wg3_launch<7, 112, 2, 2, 1, 128, 2>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, + SCin, do_snake, pad_off, pad_len); + } else { + wg3_launch<7, 112, 1, 4, 1, 128, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, + SCin, do_snake, pad_off, pad_len); + } +} +__device__ void codec_program(const CodecArgs &a, const Plan &p, const RepackSet &js, + int repack_grid, + const __half *const *state_override = nullptr) { + cudaStream_t s = nullptr; + const int B = a.B; + const void *const *W = a.inp; + const __half *const *ST = state_override ? state_override : + reinterpret_cast(a.st); + __half *const *NS = reinterpret_cast<__half *const *>(a.ns); + dense(s, 3, p.LB, H16(W[1]), F32(W[2]), nullptr, nullptr, p.X0, ST[1], NS[0], + 40, 768, p.s0, p.s0.SC, 0, 0, 2); + const __half *k3iw = H16(W[3]); + const float *k3ib = F32(W[4]); + const float *k3ia = F32(W[5]); + const __half *k3sw = H16(W[6]); + const float *k3sb = F32(W[7]); + const float *k3sa = F32(W[8]); + const size_t IW = (size_t)1536 * 768 * 3, SW = (size_t)768 * 1536 * 3; + dense(s, 3, p.X0, k3iw, k3ib, k3ia, nullptr, p.M0, ST[2], NS[1], 768, 1536, + p.s0, p.s0.SC, 1, 0, 2, p.W3i ? p.W3i - IW : nullptr, 0); + dense(s, 3, p.M0, k3sw, k3sb, k3sa, p.X0, p.Y0, ST[3], NS[2], 1536, 768, + p.s0, p.s0.SC, 1, 1, 1, p.W3s ? p.W3s - SW : nullptr, 0); + TCONV_BPACK_DISPATCH8(1, 2, 8, p.s0.T, 768, B, + p.Y0, H16(W[9]), F32(W[10]), F32(W[11]), p.A, ST[4], NS[3], 768, 768, + p.s0.T, p.s0.L, p.s0.SC, 1, p.s1.L, p.s1.SC, 2, 0, 2, B); + if (p.W3i != nullptr) + native_invoke(repack_activation_only,dim3(min(1152,96*DIVUP(p.s1.SC,128))),dim3(128),0, + p.A, p.Ac, p.s1.SC, p.s1.SC, DIVUP(p.s1.SC,128)); + __half *cur = p.A, *nxt = p.Bb; + for (int i = 0; i < 6; ++i) { + const int bi = i + 1; + const __half *ain = (i == 0 && p.Ac) ? p.Ac : cur; + dense(s, 3, ain, k3iw + bi * IW, k3ib + bi * 1536, k3ia + bi * 768, + nullptr, p.M1, ST[5 + 2 * i], NS[4 + 2 * i], 768, 1536, p.s1, + p.s1.SC, 1, 0, 2, p.W3i ? p.W3i + (bi - 1) * IW : nullptr, + (i > 0 || p.Ac) ? 1 : 0); + const __half *ps = (i < 5) ? ST[6 + 2 * i] : ST[16]; + const int po = (i < 5) ? 0 : 1, pl = 1 + (i < 5); + dense(s, 3, p.M1, k3sw + bi * SW, k3sb + bi * 768, k3sa + bi * 384, cur, + nxt, ps, NS[5 + 2 * i], 1536, 768, p.s1, p.s1.SC, 1, po, pl, + p.W3s ? p.W3s + (bi - 1) * SW : nullptr, + i == 0 ? 1 : (i < 5 ? 2 : 3)); + __half *t = cur; cur = nxt; nxt = t; + } + TCONV_BPACK_DISPATCH1(2, 9, 16, p.s1.T, 384, B, + cur, H16(W[12]), F32(W[13]), F32(W[14]), p.P, ST[17], NS[16], 768, 384, + p.s1.T, p.s1.L, p.s1.SC, 1, p.s2.L, p.s2.SC, 6, 0, 6, B); + dense7(s, 0, p.P, p.W7[0], F32(W[22]), F32(W[23]), nullptr, p.M2, ST[18], + NS[17], p.s2, p.s2.SC, 0, 6); + dense7(s, 1, p.M2, p.W7[1], F32(W[25]), F32(W[26]), p.P, p.Q, ST[19], + NS[18], p.s2, p.s2.SC, 5, 1); + TCONV_BPACK_DISPATCH(3, 7, 16, p.s2.T, 128, B, + p.Q, H16(W[15]), F32(W[16]), F32(W[17]), p.R, ST[20], NS[19], 384, 128, + p.s2.T, p.s2.L, p.s2.SC, 5, p.s3.L, p.s3.SC, 6, 0, 6, B); + dense7(s, 2, p.R, p.W7[2], F32(W[28]), F32(W[29]), nullptr, p.M3, ST[21], + NS[20], p.s3, p.s3.SC, 0, 6); + dense7(s, 3, p.M3, p.W7[3], F32(W[31]), F32(W[32]), p.R, p.S, ST[22], + NS[21], p.s3, p.s3.SC, 5, 1); + if (p.s3.T >= 1500) + pdl_launch(transpose_conv_kernel<4, 7, 8, true>, dim3(1, 32, B), + dim3(tc_blk(p.s3.T, 8)), + (size_t)tc_blk(p.s3.T, 8) * (8 * 7) * 2 + 64, s, + p.S, H16(W[18]), F32(W[19]), F32(W[20]), p.U, ST[23], + NS[22], 128, 32, p.s3.T, p.s3.L, p.s3.SC, 5, p.s4.L, + p.s4.SC, 6, 0, 6, B); + else + pdl_launch(transpose_conv_kernel<4, 7, 4, true>, dim3(1, 32, B), + dim3(tc_blk(p.s3.T, 4)), + (size_t)tc_blk(p.s3.T, 4) * (4 * 7) * 2 + 64, s, + p.S, H16(W[18]), F32(W[19]), F32(W[20]), p.U, ST[23], + NS[22], 128, 32, p.s3.T, p.s3.L, p.s3.SC, 5, p.s4.L, + p.s4.SC, 6, 0, 6, B); + { + const int ntile = DIVUP(p.s4.N, TL_BN); + const int grid = ntile < 132 ? ntile : 132; + pdl_launch_pf(tail_fused_kernel, dim3(grid), dim3(TL_NT), TL_SHARED, s, + p.U, (size_t)32 * p.s4.SC * sizeof(__half), + p.U, p.W7[4], p.W7[5], H16(W[39]), F32(W[34]), F32(W[35]), + F32(W[37]), F32(W[38]), F32(W[40]), ST[24], ST[25], + static_cast<__half *>(a.waveform), NS[23], NS[24], NS[25], B, + p.s4.T, p.s4.L, p.s4.SC, p.s4.N, ntile); + } +} diff --git a/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/native_stage1.cu b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/native_stage1.cu new file mode 100644 index 0000000..cca3e11 --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/native_stage1.cu @@ -0,0 +1,357 @@ +#include "native_stage1.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#ifndef NATIVE_CODEC_PREFETCH +#define NATIVE_CODEC_PREFETCH 1 +#endif +#ifndef NATIVE_CODEC_PACK_SUBCTAS +#define NATIVE_CODEC_PACK_SUBCTAS 1 +#endif +namespace cg=cooperative_groups; +constexpr int PHYSICAL_THREADS=512,MAIN_ARENA=227840,AUX_BYTES=4096,META_OFFSET=231936,TOTAL_SHARED=232192; +constexpr int SLICE_BYTES=MAIN_ARENA+AUX_BYTES; +constexpr int MAX_SUBCTA=8; +struct NativeMeta { + dim3 grid,threads; + uint3 block[MAX_SUBCTA]; + NativeStage1Status *status;const int32_t *lengths;int padded_frames; + int submask,tshift,tmask,slice; + int resident; + int identity_slots; +}; +static_assert(sizeof(NativeMeta)<=TOTAL_SHARED-META_OFFSET); +__device__ __forceinline__ char *native_base(){extern __shared__ __align__(16) char arena[];return arena;} +__device__ __forceinline__ NativeMeta *native_meta(){return reinterpret_cast(native_base()+META_OFFSET);} +__device__ __forceinline__ int native_subid(){const NativeMeta *m=native_meta();return (int)((threadIdx.x>>m->tshift)&(unsigned)m->submask);} +__device__ __forceinline__ char *native_arena(){const NativeMeta *m=native_meta();return native_base()+native_subid()*m->slice;} +__device__ __forceinline__ char *native_aux(){const NativeMeta *m=native_meta();return native_base()+native_subid()*m->slice+(m->slice-AUX_BYTES);} +__device__ __forceinline__ uint3 native_tid(){uint3 r;r.x=threadIdx.x&(unsigned)native_meta()->tmask;r.y=0u;r.z=0u;return r;} +__device__ __forceinline__ dim3 native_block(){const uint3 b=native_meta()->block[native_subid()];return dim3(b.x,b.y,b.z);} +__device__ __forceinline__ dim3 native_grid(){return native_meta()->grid;} +__device__ __forceinline__ dim3 native_threads(){return native_meta()->threads;} +__device__ __forceinline__ int native_resident(){return native_meta()->resident;} +__device__ __forceinline__ void native_local_sync(){ + const NativeMeta *m=native_meta();unsigned n=m->threads.x,id=15u-(unsigned)native_subid(); + asm volatile("barrier.sync %0, %1;"::"r"(id),"r"(n):"memory"); +} +__device__ __forceinline__ dim3 &native_config_grid(cudaLaunchConfig_t &c){return c.gridDim;} +__device__ __forceinline__ dim3 &native_config_threads(cudaLaunchConfig_t &c){return c.blockDim;} +__device__ __forceinline__ int native_true_tokens(int row,int padded_tokens){ + const NativeMeta *m=native_meta(); + return (padded_tokens/m->padded_frames)*m->lengths[row]; +} +__device__ __forceinline__ void native_warm(const char *p,size_t n,int k,int nk,int me){ + if(p==nullptr||n<16||nk<=0)return; + const size_t a=(size_t)p&15u;p-=a;n=(n+a)&~(size_t)15; + const uint4 *q=reinterpret_cast(p); + const size_t cnt=n>>4,step=(size_t)nk*PHYSICAL_THREADS; + uint32_t acc=0u; + for(size_t i=(size_t)k*PHYSICAL_THREADS+(size_t)me;i __device__ __forceinline__ void native_invoke_pf(F fn,dim3 g,dim3 t,size_t shared,const void *pf,size_t pfb,A&&... args){ + const int me=threadIdx.x,owner=blockIdx.x,owners=gridDim.x; + const cg::grid_group all=cg::this_grid();NativeMeta *meta=native_meta(); + if(t.x>512||t.y!=1||t.z!=1||t.x%32||shared>MAIN_ARENA){ + if(me==0&&owner==0)meta->status->code=10; + all.sync();return; + } + int sub=1,tshift=9; + const unsigned tx=t.x; + const int vtotal=(int)(g.x*g.y*g.z); + if(NATIVE_CODEC_PACK_SUBCTAS && (tx&(tx-1))==0u){ + int s=0;while((1u<MAX_SUBCTA)cand=MAX_SUBCTA; + const int want=(vtotal+owners-1)/owners; + while(cand>1&&cand>want)cand>>=1; + while(cand>1&&shared>(size_t)(SLICE_BYTES/cand-AUX_BYTES))cand>>=1; + sub=cand; + } + const int slice=(SLICE_BYTES/sub)&~127; + if(me==0){meta->grid=g;meta->threads=t;meta->submask=sub-1;meta->tshift=tshift; + meta->tmask=(int)((1u<slice=slice;} + const int total=vtotal,grp=me>>tshift; + const bool holder=me=busy&&pf!=nullptr)native_warm((const char*)pf,pfb,owner-busy,owners-busy,me); + for(int base=owner*sub;baseblock[grp]=(g.y==1u&&g.z==1u)?make_uint3((unsigned)i,0u,0u) + :make_uint3(i%g.x,(i/g.x)%g.y,i/(g.x*g.y)); + __syncthreads(); + if(holder&&i(args)...); + asm volatile("cp.async.wait_all;":::"memory"); + if(t.x>=128)asm volatile("wgmma.wait_group.sync.aligned 0;":::"memory"); + } + __syncthreads(); + } + all.sync(); +} +template __device__ __forceinline__ void native_invoke(F fn,dim3 g,dim3 t,size_t shared,A&&... args){ + native_invoke_pf(fn,g,t,shared,nullptr,(size_t)0,static_cast(args)...); +} +template __device__ __forceinline__ void native_cudaLaunchKernelEx(cudaLaunchConfig_t *c,F fn,A... args){ + native_invoke(fn,c->gridDim,c->blockDim,c->dynamicSmemBytes,args...); +} +template __device__ __forceinline__ void native_cudaLaunchKernelEx_pf(cudaLaunchConfig_t *c,F fn,const void *pf,size_t pfb,A... args){ + native_invoke_pf(fn,c->gridDim,c->blockDim,c->dynamicSmemBytes,pf,pfb,args...); +} +#define blockIdx native_block() +#define gridDim native_grid() +#define blockDim native_threads() +#define threadIdx native_tid() +#define __syncthreads() native_local_sync() +#include "legacy_codec.cuh" +#undef __syncthreads +#undef threadIdx +#undef blockDim +#undef gridDim +#undef blockIdx +namespace { +struct FastArgs { + CodecArgs codec;Plan plan;RepackSet repack;int repack_grid;int input_frames;int blocks; + const int64_t *codes;const int32_t *lengths,*slots; + const __half *state[26];__half *next[26];float *pcm;int32_t *samples; + NativeStage1Status *status;int nslots; + int channels[26],histories[26]; +}; +static_assert(sizeof(FastArgs)<=4096,"precomputed kernel argument must fit the conservative parameter budget"); +struct Handle {FastArgs args{};Plan plans[17]{};RepackSet repacks[17]{};int repack_grids[17]{};void *scratch=nullptr;size_t bytes=0;void *immutable=nullptr;size_t immutable_bytes=0;int setup_launches=0;std::vector packed;int max_frames=0,blocks=0,device=-1;}; +thread_local std::string error; +bool check(cudaError_t e,const char *where){if(e==cudaSuccess)return true;error=std::string(where)+": "+cudaGetErrorString(e);return false;} +__device__ void reject(NativeStage1Status *s,int code,int row,int frame,int detail){ + if(atomicCAS(&s->code,0,code)==0){s->row=row;s->frame=frame;s->detail=detail;} +} +__global__ __launch_bounds__(512,1) void stage1_fast_persistent(const __grid_constant__ FastArgs a){ + { + const auto all=cg::this_grid(); + const size_t me=(size_t)blockIdx.x*512+threadIdx.x,step=(size_t)gridDim.x*512; + if(threadIdx.x==0){native_meta()->status=a.status;native_meta()->lengths=a.lengths;native_meta()->padded_frames=a.codec.F;native_meta()->resident=a.blocks;native_meta()->identity_slots=(a.nslots==a.codec.B);} + if(me==0)*a.status=NativeStage1Status{0,-1,-1,0}; + all.sync();const int B=a.codec.B,F=a.codec.F,inputF=a.input_frames; + for(int row=threadIdx.x;rowidentity_slots,0); + __syncthreads(); + for(int row=me;rowinputF)reject(a.status,8,row,-1,n); + if(slot<0||slot>=a.nslots)reject(a.status,2,row,-1,slot); + for(int j=0;j=0&&n<=inputF)?n*1764:0; + } + for(size_t i=me;i<(size_t)B*inputF;i+=step){ + int row=i/inputF,t=i%inputF;bool any=false; + for(int c=0;c<16;++c){int64_t v=a.codes[i*16+c];if(v<0||v>=1024)reject(a.status,4,row,t,c);any|=v!=0;} + if(t=a.lengths[row]&&any)reject(a.status,6,row,t,0); + } + if(a.nslots!=B) + for(int p=0;p<26;++p){ + const size_t vectors=(size_t)a.nslots*a.channels[p]*a.histories[p]/8; + const uint4 *src=reinterpret_cast(a.state[p]); + uint4 *dst=reinterpret_cast(a.next[p]); + for(size_t i=me;icode)return; + if(!native_meta()->identity_slots) for(int p=1;p<26;++p){ + const size_t stride=(size_t)a.channels[p]*a.histories[p]; + __half *dest=(__half*)a.codec.st[p]; + for(int row=blockIdx.x;row(a.state[p]+(size_t)slot*stride); + uint4 *dst=reinterpret_cast(dest+(size_t)row*stride); + for(size_t offset=threadIdx.x;offset>(2*(c%5)))&3)-2)*.5f; + } + v=__float2half_rn(x); + } + packed[(size_t)c*SC+(size_t)row*L+j]=v; + } + all.sync(); + } + codec_program(a.codec,a.plan,a.repack,a.repack_grid, + native_meta()->identity_slots?a.state:nullptr); + if(a.status->code)return; + const size_t me=(size_t)blockIdx.x*512+threadIdx.x,step=(size_t)gridDim.x*512; + const int B=a.codec.B,F=a.codec.F,inputF=a.input_frames; + for(int p=0;p<26;++p){ + const size_t stride=(size_t)a.channels[p]*a.histories[p];const __half *src=(const __half*)a.codec.ns[p]; + for(int row=blockIdx.x;row(src+(size_t)row*stride); + const uint4 *prior=reinterpret_cast(a.state[p]+(size_t)slot*stride); + uint4 *dest=reinterpret_cast(a.next[p]+(size_t)slot*stride); + for(size_t offset=threadIdx.x;offset(a.codec.waveform); + float2 *pcm=reinterpret_cast(a.pcm); + const int row_pairs=inputF*882; + for(size_t i=me;i<(size_t)B*row_pairs;i+=step){ + const int row=i/row_pairs,t=i%(size_t)row_pairs; + pcm[i]=tstatus=a.status;native_meta()->resident=a.blocks;} + if(blockIdx.x==0&&threadIdx.x==0)*a.status=NativeStage1Status{0,-1,-1,0}; + all.sync(); + native_invoke(repack_rs_kernel<7>,dim3(a.repack_grid),dim3(128),MAIN_ARENA, + a.repack,H16(a.codec.inp[0]),H16(a.codec.st[0]),a.plan.LB, + 40,a.plan.s0.T,a.plan.s0.L,a.plan.s0.SC,2,a.codec.B); + native_invoke(repack3_rs_kernel,dim3(1152,12),dim3(128),MAIN_ARENA, + H16(a.codec.inp[3]),H16(a.codec.inp[6]),a.plan.W3i,a.plan.W3s, + a.plan.A,a.plan.Ac,a.plan.s1.SC,a.plan.s1.SC,DIVUP(a.plan.s1.SC,128)); +} +bool stack(Handle *h,int slot,const NativeStage1Weights *w,int first,int field,size_t elements){ + const int ids[7]={first,first+3,first+5,first+7,first+9,first+11,first+13}; + const size_t stride=elements*(field==0?sizeof(__half):sizeof(float));void *ptr=nullptr; + if(!check(cudaMalloc(&ptr,7*stride),"packed-weight allocation"))return false; + h->packed.push_back(ptr);h->args.codec.inp[slot]=ptr; + for(int i=0;i<7;++i){const void *src=field==0?w[ids[i]].weight:field==1?(const void*)w[ids[i]].bias:(const void*)w[ids[i]].alpha; + if(!check(cudaMemcpy((char*)ptr+i*stride,src,stride,cudaMemcpyDeviceToDevice),"immutable weight stack"))return false;} + return true; +} +void release(Handle *h){if(h->immutable)cudaFree(h->immutable);if(h->scratch)cudaFree(h->scratch);for(void *p:h->packed)cudaFree(p);delete h;} +} +extern "C" void *native_stage1_create(int b,int maxf,int slots,int blocks,const NativeStage1Weights *w){ + error.clear();if((b!=1&&b!=16&&b!=32&&b!=64&&b!=128&&b!=256&&b!=512&&b!=1024&&b!=2048&&b!=4096)||maxf<1||maxf>16||slots4096||blocks<0||!w){error="invalid create arguments";return nullptr;} + for(int p=0;p<26;++p)if(!w[p].weight||!w[p].bias||(p!=0&&p!=25&&!w[p].alpha)){error="missing weight pointer";return nullptr;} + Handle *h=new(std::nothrow)Handle;if(!h){error="host allocation";return nullptr;} + h->max_frames=maxf;h->args.nslots=slots;h->args.codec.B=b; + const int cs[26]={40,768,1536,768,768,1536,768,1536,768,1536,768,1536,768,1536,768,1536,768,384,768,384,128,256,128,32,64,32}; + const int hs[26]={2,2,2,1,2,2,2,2,2,2,2,2,2,2,2,2,1,6,6,1,6,6,1,6,6,6}; + cudaDeviceProp prop{};int occupancy=0,coop=0; + if(!check(cudaGetDevice(&h->device),"device")||!check(cudaGetDeviceProperties(&prop,h->device),"properties")|| + !check(cudaDeviceGetAttribute(&coop,cudaDevAttrCooperativeLaunch,h->device),"cooperative support")|| + !check(cudaFuncSetAttribute(stage1_fast_persistent,cudaFuncAttributeMaxDynamicSharedMemorySize,TOTAL_SHARED),"shared-memory optin")|| + !check(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&occupancy,stage1_fast_persistent,512,TOTAL_SHARED),"occupancy")) {release(h);return nullptr;} + if(prop.major!=9||prop.minor!=0||!coop||occupancy<1){error="requires resident SM90 cooperative kernel";release(h);return nullptr;} + h->blocks=blocks?blocks:occupancy*prop.multiProcessorCount; + if(h->blocks>occupancy*prop.multiProcessorCount){error="grid exceeds cooperative occupancy";release(h);return nullptr;} + h->args.blocks=h->blocks; + const int capacityF=maxf<=3?3:(maxf<=5?5:std::max(8,maxf)); + size_t decoder=codec_scratch_bytes(b,capacityF),halfs=(size_t)b*40*2*capacityF+(size_t)b*1764*capacityF; + for(int p=0;p<26;++p){h->args.channels[p]=cs[p];h->args.histories[p]=hs[p];halfs+=2*(size_t)b*cs[p]*hs[p];} + h->bytes=decoder+halfs*sizeof(__half); + if(!check(cudaMalloc(&h->scratch,h->bytes),"scratch allocation")){release(h);return nullptr;} + if(!check(cudaMemset(h->scratch,0,h->bytes),"initial scratch padding")){release(h);return nullptr;} + h->args.codec.scratch=h->scratch;__half *q=(__half*)((char*)h->scratch+decoder); + h->args.codec.inp[0]=q;q+=(size_t)b*40*2*capacityF;h->args.codec.waveform=q;q+=(size_t)b*1764*capacityF; + for(int p=0;p<26;++p){h->args.codec.st[p]=q;q+=(size_t)b*cs[p]*hs[p];h->args.codec.ns[p]=q;q+=(size_t)b*cs[p]*hs[p];} + const void **in=h->args.codec.inp; + in[1]=w[0].weight;in[2]=w[0].bias; + if(!stack(h,3,w,1,0,(size_t)1536*768*3)||!stack(h,4,w,1,1,1536)||!stack(h,5,w,1,2,768)|| + !stack(h,6,w,2,0,(size_t)768*1536*3)||!stack(h,7,w,2,1,768)||!stack(h,8,w,2,2,384)){release(h);return nullptr;} + const int upops[4]={3,16,19,22};for(int i=0;i<4;++i){in[9+3*i]=w[upops[i]].weight;in[10+3*i]=w[upops[i]].bias;in[11+3*i]=w[upops[i]].alpha;} + for(int i=0;i<3;++i){int p=17+3*i;in[21+6*i]=w[p].weight;in[22+6*i]=w[p].bias;in[23+6*i]=w[p].alpha; + in[24+6*i]=w[p+1].weight;in[25+6*i]=w[p+1].bias;in[26+6*i]=w[p+1].alpha;} + in[39]=w[25].weight;in[40]=w[25].bias; + __half *immutable_w7[6]{};__half *immutable_w3i=nullptr,*immutable_w3s=nullptr; + if(capacityF>=3){ + size_t wh=0;for(int j=0;j<6;++j)wh+=(size_t)native_RP_COUTP(j)*native_RP_CIN(j)*7; + const size_t IW=(size_t)3*1536*768;wh+=14*IW; + h->immutable_bytes=wh*sizeof(__half); + if(!check(cudaMalloc(&h->immutable,h->immutable_bytes+sizeof(NativeStage1Status)),"immutable weight allocation")){release(h);return nullptr;} + if(!check(cudaMemset(h->immutable,0,h->immutable_bytes+sizeof(NativeStage1Status)),"immutable padding initialization")){release(h);return nullptr;} + __half *q=(__half*)h->immutable; + for(int j=0;j<6;++j){immutable_w7[j]=q;q+=(size_t)native_RP_COUTP(j)*native_RP_CIN(j)*7;} + immutable_w3i=q+IW;q+=7*IW;immutable_w3s=q+IW;q+=7*IW; + if((char*)q!=(char*)h->immutable+h->immutable_bytes){error="immutable layout size mismatch";release(h);return nullptr;} + } + for(int f:{3,5,8}) if(f<=capacityF){ + Plan p=build_plan(b,f,static_cast<__half*>(h->scratch)); + for(int j=0;j<6;++j)p.W7[j]=immutable_w7[j];p.W3i=immutable_w3i;p.W3s=immutable_w3s; + const void *const *W=h->args.codec.inp; + RepackSet js{}; + const void *SRC[RP_JOBS] = {W[21], W[24], W[27], W[30], W[33], W[36]}; + for (int j = 0; j < RP_JOBS; ++j) { + js.src[j] = H16(SRC[j]); + js.dst[j] = p.W7[j]; + js.Cout[j] = native_RP_COUT(j); + js.CoutP[j] = native_RP_COUTP(j); + js.Cin[j] = native_RP_CIN(j); + } + int acc = 0; + for (int j = 0; j < RP_JOBS; ++j) { + js.base[j] = acc; + js.nch[j] = native_RP_CIN(j) / 16; + js.rs[j] = (j < 4) ? 1 : 0; + acc += (native_RP_COUTP(j) / 16) * DIVUP(native_RP_CIN(j), 64); + } + js.k3base = acc; + js.k3src[0] = H16(W[3]); + js.k3src[1] = H16(W[6]); + js.k3dst[0] = p.W3i ? p.W3i - (size_t)3 * 1536 * 768 : nullptr; + js.k3dst[1] = p.W3s ? p.W3s - (size_t)3 * 768 * 1536 : nullptr; + const int k3n = p.W3i ? 2304 : 0; + h->plans[f]=p;h->repacks[f]=js;h->repack_grids[f]=acc+k3n; + } + if(capacityF>=3){ + FastArgs a=h->args;a.codec.F=3;a.plan=h->plans[3];a.repack=h->repacks[3];a.repack_grid=h->repack_grids[3]; + a.status=(NativeStage1Status*)((char*)h->immutable+h->immutable_bytes); + int setup_occupancy=0; + if(!check(cudaFuncSetAttribute(stage1_prepare_immutable,cudaFuncAttributeMaxDynamicSharedMemorySize,TOTAL_SHARED),"setup shared-memory optin")|| + !check(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&setup_occupancy,stage1_prepare_immutable,512,TOTAL_SHARED),"setup occupancy")|| + setup_occupancy<1||h->blocks>setup_occupancy*prop.multiProcessorCount){ + if(error.empty())error="setup grid exceeds cooperative occupancy";release(h);return nullptr;} + void *params[]={&a}; + if(!check(cudaLaunchCooperativeKernel((void*)stage1_prepare_immutable,dim3(h->blocks),dim3(512),params,TOTAL_SHARED,nullptr),"immutable setup launch")|| + !check(cudaStreamSynchronize(nullptr),"immutable setup completion")){release(h);return nullptr;} + NativeStage1Status status{}; + if(!check(cudaMemcpy(&status,a.status,sizeof(status),cudaMemcpyDeviceToHost),"immutable setup status")){release(h);return nullptr;} + if(status.code){error="immutable setup device status "+std::to_string(status.code);release(h);return nullptr;} + h->setup_launches=1; + } + return h; +} +extern "C" int native_stage1_run(void *handle,const int64_t *codes,const int32_t *lengths,const int32_t *slots, + const void *const *states,void *const *next,float *pcm,int32_t *samples,NativeStage1Status *status,int f,void *stream){ + error.clear();Handle *h=(Handle*)handle; + if(!h||!codes||!lengths||!slots||!states||!next||!pcm||!samples||!status){error="null run argument";return 1;} + if(f<1||f>8||f>h->max_frames){error="true-state variant supports public F1..8";return 1;} + const int paddedF=f<=3?3:(f<=5?5:8); + int device=-1;if(!check(cudaGetDevice(&device),"device")||device!=h->device){error="device mismatch";return 1;} + FastArgs a=h->args;a.input_frames=f;a.plan=h->plans[paddedF];a.repack=h->repacks[paddedF];a.repack_grid=h->repack_grids[paddedF];a.codec.F=paddedF;a.codes=codes;a.lengths=lengths;a.slots=slots;a.pcm=pcm;a.samples=samples;a.status=status; + for(int p=0;p<26;++p){if(!states[p]||!next[p]||states[p]==next[p]){error="invalid/aliased state";return 1;}a.state[p]=(const __half*)states[p];a.next[p]=(__half*)next[p];} + void *params[]={&a};if(!check(cudaLaunchCooperativeKernel((void*)stage1_fast_persistent,dim3(h->blocks),dim3(512),params,TOTAL_SHARED,(cudaStream_t)stream),"fast Stage1 launch"))return 1; + return check(cudaPeekAtLastError(),"immediate launch status")?0:1; +} +extern "C" int native_stage1_destroy(void *p){error.clear();if(!p)return 0;release((Handle*)p);return check(cudaPeekAtLastError(),"destroy status")?0:1;} +extern "C" const char *native_stage1_last_error(){return error.c_str();} +extern "C" size_t native_stage1_scratch_bytes(void *p){return p?((Handle*)p)->bytes:0;} +extern "C" int native_stage1_resident_blocks(void *p){return p?((Handle*)p)->blocks:0;} +extern "C" int native_stage1_setup_kernel_launches(void *p){return p?((Handle*)p)->setup_launches:0;} +extern "C" size_t native_stage1_immutable_weight_bytes(void *p){return p?((Handle*)p)->immutable_bytes:0;} diff --git a/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/native_stage1.h b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/native_stage1.h new file mode 100644 index 0000000..a95b5f0 --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/native_stage1.h @@ -0,0 +1,26 @@ +#pragma once +#include +#include +struct NativeStage1Weights { const void *weight; const float *bias; const float *alpha; }; +struct NativeStage1Status { int32_t code, row, frame, detail; }; +#ifdef __cplusplus +extern "C" { +#endif +void *native_stage1_create(int batch, int max_frames, int state_slots, + int resident_blocks, const NativeStage1Weights *weights26); +int native_stage1_run(void *handle, const int64_t *codes, const int32_t *lengths, + const int32_t *slots, const void *const *states26, + void *const *next_states26, float *pcm, int32_t *samples, + NativeStage1Status *status, int frames, void *stream); +int native_stage1_run_profile(void *handle, const int64_t *codes, const int32_t *lengths, + const int32_t *slots, const void *const *states26, + void *const *next_states26, float *pcm, int32_t *samples, + NativeStage1Status *status, int frames, void *stream, + uint64_t *phase_timestamps30); +int native_stage1_destroy(void *handle); +const char *native_stage1_last_error(); +size_t native_stage1_scratch_bytes(void *handle); +int native_stage1_resident_blocks(void *handle); +#ifdef __cplusplus +} +#endif diff --git a/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/repack3.inc b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/repack3.inc new file mode 100644 index 0000000..2f61c93 --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/repack3.inc @@ -0,0 +1,66 @@ +__device__ __noinline__ void repack3_kernel( + const __half *__restrict__ win, const __half *__restrict__ wsk, + __half *__restrict__ din, __half *__restrict__ dsk) { + int4 (*sm)[49]=reinterpret_cast(native_arena()); + const int job = blockIdx.y; + const int Cout = job ? 768 : 1536; + const int Cin = job ? 1536 : 768; + const __half *src = job ? wsk : win; + __half *dst = job ? dsk : din; + const int GPB = Cin >> 3; + const int Kg = Cin * 3; + const int GT = GPB >> 4; + const int RT = Cout >> 5; + const int t = threadIdx.x; + for (int tile = blockIdx.x; tile < 6 * RT * GT; tile += gridDim.x) { + const int gblk = tile % GT; + const int rest = tile / GT; + const int corow = rest % RT; + const int b = rest / RT; + const int co0 = corow * 32; + const int g0 = gblk * 16; + const __half *sbase = src + (size_t)(b + 1) * Cout * Kg + + (size_t)co0 * Kg + g0 * 24; +#pragma unroll + for (int j = 0; j < 6; ++j) { + const int idx = t + j * 256; + const int r = idx / 48, c = idx - (idx / 48) * 48; + sm[r][c] = *reinterpret_cast(sbase + (size_t)r * Kg + + c * 8); + } + __syncthreads(); + __half *dbase = dst + (size_t)b * (3 * (size_t)Cout * Cin); +#pragma unroll + for (int p = 0; p < 2; ++p) { + const int u = t + p * 256; + const int col = (u & 7) + 8 * (u >> 7); + const int g = g0 + ((u >> 3) & 15); + const int co = co0 + col; + const int c = g >> 1, v = g & 1; + const int4 a = sm[col][((u >> 3) & 15) * 3]; + const int4 e = sm[col][((u >> 3) & 15) * 3 + 1]; + const int4 f = sm[col][((u >> 3) & 15) * 3 + 2]; + uint32_t h[12]; + h[0] = a.x; h[1] = a.y; h[2] = a.z; h[3] = a.w; + h[4] = e.x; h[5] = e.y; h[6] = e.z; h[7] = e.w; + h[8] = f.x; h[9] = f.y; h[10] = f.z; h[11] = f.w; + __half *d = dbase + (size_t)c * (48 * Cout) + + (size_t)(co >> 3) * 128 + v * 64 + (co & 7) * 8; +#pragma unroll + for (int tp = 0; tp < 3; ++tp) { + uint32_t o[4]; +#pragma unroll + for (int jj = 0; jj < 4; ++jj) { + const uint32_t g0w = h[3 * jj], g1w = h[3 * jj + 1], + g2w = h[3 * jj + 2]; + o[jj] = (tp == 0) ? __byte_perm(g0w, g1w, 0x7610) + : (tp == 1) ? __byte_perm(g0w, g2w, 0x5432) + : __byte_perm(g1w, g2w, 0x7610); + } + *reinterpret_cast(d + (size_t)tp * (16 * Cout)) = + make_int4(o[0], o[1], o[2], o[3]); + } + } + __syncthreads(); + } +} diff --git a/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/tail_kernel.inc b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/tail_kernel.inc new file mode 100644 index 0000000..8794f42 --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/tail_kernel.inc @@ -0,0 +1,492 @@ +#define TL_NT 512 +#define TL_MW 512 +#define TL_MLD TL_MW +#define TL_VW 504 +#define TL_UW 520 +#define TL_BN 496 +#define TL_VLD 520 +#define TL_LBO_U (TL_UW * 8) +#define TL_LBO_M (TL_MLD * 8) +#define TL_LBO_V (TL_VLD * 8) +#define TL_SBO_A 128 +#define TL_PU_E (4 * TL_UW * 8) +#define TL_PM_E (8 * TL_MLD * 8) +#define TL_PV_E (32 * TL_VLD) +#define TL_AI_E (14 * 1024) +#define TL_AS_E (32 * 512 + 512) +#define TL_SHARED ((2 * TL_PU_E + TL_PM_E + TL_PV_E + TL_AI_E + TL_AS_E) * 2) +__device__ __forceinline__ void tl_stmatrix4(void *dst, const uint32_t r[4], + bool trans) { + const uint32_t a = static_cast(__cvta_generic_to_shared(dst)); + if (trans) + asm volatile( + "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1,%2,%3,%4};\n" ::"r"( + a), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3])); + else + asm volatile( + "stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1,%2,%3,%4};\n" ::"r"( + a), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3])); +} +__device__ __forceinline__ void tl_ldmatrix4t(uint32_t r[4], const void *src) { + const uint32_t a = static_cast(__cvta_generic_to_shared(src)); + asm volatile( + "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(a)); +} +__device__ __forceinline__ void tl_stmatrix1t(void *dst, uint32_t r) { + const uint32_t a = static_cast(__cvta_generic_to_shared(dst)); + asm volatile( + "stmatrix.sync.aligned.m8n8.x1.trans.shared.b16 [%0], {%1};\n" :: + "r"(a), "r"(r)); +} +template +__device__ __forceinline__ void tl_skip_band( + const __half *PU, __half *PM, const __half *AS, + __half *PV, int c0, int warp, int lane, int er0, const float *sh_ask, + const float *sh_ari_v, float b0v, float b1v) { + constexpr int NB8 = BW / 8; + float acc[BW / 2]; +#pragma unroll + for (int i = 0; i < BW / 2; ++i) acc[i] = ((i & 2) == 0) ? b0v : b1v; + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int c = 0; c < 4; ++c) { + const uint64_t da0 = + wg_desc(AS + c * 8 * 512, 128u, TL_SBO_A * 2u); + const uint64_t db0 = + wg_desc(PM + (2 * c) * TL_LBO_M + (c0 + 2) * 8, + TL_LBO_M * 2u, 128u); +#pragma unroll + for (int p = 0; p < 4; ++p) { + const uint64_t da = da0 + 64 * (2 * p); + const uint64_t db = db0 + 2 * p; + if constexpr (BW == 120) { + WGMMA_M64N120K16_T(acc, da, db, 0, 0); + } else { + WGMMA_M64N128K16_T(acc, da, db, 0, 0); + } + } + } + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + const int m_addr = lane >> 3, row = lane & 7; + const int cg = (warp & 1) * 2; + if (er0 >= 32) { +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = __floats2half2_rn( + acc[jj * 4 + h * 2 + 0], acc[jj * 4 + h * 2 + 1]); + r[m] = *reinterpret_cast(&hv); + } else { + r[m] = 0; + } + } + const int jj = j + (m_addr >> 1), h = m_addr & 1; + tl_stmatrix4(PV + (cg + h) * TL_LBO_V + (c0 + jj * 8 + row) * 8, + r, true); + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + if (er0 < 32) { +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; + const int jja = j + (m_addr >> 1), ha = m_addr & 1; + tl_ldmatrix4t(r, PV + (cg + ha) * TL_LBO_V + + (c0 + jja * 8 + row + 1) * 8); +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = *reinterpret_cast<__half2 *>(&r[m]); + acc[jj * 4 + h * 2 + 0] += __half2float(hv.x); + acc[jj * 4 + h * 2 + 1] += __half2float(hv.y); + } + } + } +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; + const int jja = j + (m_addr >> 1), ha = m_addr & 1; + tl_ldmatrix4t(r, PU + (cg + ha) * TL_LBO_U + + (c0 + jja * 8 + row + 16) * 8); +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = *reinterpret_cast<__half2 *>(&r[m]); + acc[jj * 4 + h * 2 + 0] += __half2float(hv.x); + acc[jj * 4 + h * 2 + 1] += __half2float(hv.y); + } + } + } + } + if (er0 >= 32) return; + if (er0 < 16) { + const float a0 = sh_ask[er0], r0 = sh_ari_v[er0]; + const float a1 = sh_ask[er0 + 8], r1 = sh_ari_v[er0 + 8]; +#pragma unroll + for (int i = 0; i < BW / 2; ++i) { + const int h = (i & 2) >> 1; + const float v = acc[i]; + const float sn = __sinf((h ? a1 : a0) * v); + acc[i] = __fmaf_rn(sn * sn, h ? r1 : r0, v); + } + } else { +#pragma unroll + for (int i = 0; i < BW / 2; ++i) { + const float v = acc[i]; + acc[i] = v > 0.f ? v : 0.01f * v; + } + } +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = __floats2half2_rn( + acc[jj * 4 + h * 2 + 0], acc[jj * 4 + h * 2 + 1]); + r[m] = *reinterpret_cast(&hv); + } else { + r[m] = 0; + } + } + const int jj = j + (m_addr >> 1), h = m_addr & 1; + tl_stmatrix4(PM + (cg + h) * TL_LBO_V + + (c0 + jj * 8 + row) * 8, + r, true); + } +} +__device__ __noinline__ void tail_fused_kernel( + const __half *__restrict__ U, const __half *__restrict__ win, + const __half *__restrict__ wsk, const __half *__restrict__ wpo, + const float *__restrict__ bin, const float *__restrict__ ain, + const float *__restrict__ bsk, const float *__restrict__ ask, + const float *__restrict__ bpo, const __half *__restrict__ st_m4, + const __half *__restrict__ st_v, __half *__restrict__ wave, + __half *__restrict__ ns_u, __half *__restrict__ ns_m4, + __half *__restrict__ ns_v, int B, int T4, int L4, int SC4, int N4, + int NTILE) { + char *tl_arena=reinterpret_cast(native_arena()); + __half *const PU = reinterpret_cast<__half *>(tl_arena); + __half *const PM = PU + 2 * TL_PU_E; + __half *const PV = PM + TL_PM_E; + __half *const AI = PV + TL_PV_E; + __half *const AS = AI + TL_AI_E; + __half *const PB = AS + 32 * 512; + float *sh_ain=(float*)native_aux(), *sh_ari=sh_ain+32, *sh_ask=sh_ari+32, *sh_ari_v=sh_ask+16; + const int tid = threadIdx.x; + const int wg = tid >> 7; + const int warp = (tid & 127) >> 5; + const int lane = tid & 31; + const int er0 = warp * 16 + (lane >> 2); + const int cb2 = (lane & 3) * 2; + for (int i = tid; i < 32; i += TL_NT) { + const float a = ain[i]; + sh_ain[i] = a; + sh_ari[i] = __frcp_rn(a + 1e-9f); + } + for (int i = tid; i < 16; i += TL_NT) { + const float a = ask[i]; + sh_ask[i] = a; + sh_ari_v[i] = __frcp_rn(a + 1e-9f); + } + for (int i = tid; i < TL_AI_E / 8; i += TL_NT) { + const int blk = i / 128, off = (i - blk * 128) * 8; + const int c = blk / 7, k = blk - c * 7; + *reinterpret_cast(AI + blk * 1024 + off) = + *reinterpret_cast(win + (size_t)c * 14336 + + (size_t)k * 2048 + off); + } + for (int i = tid; i < 32 * 64; i += TL_NT) { + const int blk = i / 64, off = (i - blk * 64) * 8; + const int c = blk >> 3, k = blk & 7; + *reinterpret_cast(AS + blk * 512 + off) = + (k < 7) ? *reinterpret_cast( + wsk + (size_t)c * 14336 + (size_t)k * 2048 + off) + : make_int4(0, 0, 0, 0); + } + for (int i = tid; i < 4 * (TL_VLD - TL_VW) * 8; i += TL_NT) { + const int g = i / ((TL_VLD - TL_VW) * 8); + PV[g * TL_LBO_V + TL_VW * 8 + (i - g * ((TL_VLD - TL_VW) * 8))] = + __float2half(0.f); + } + for (int i = tid; i < 512; i += TL_NT) PB[i] = __float2half(0.f); + for (int i = tid; i < 2 * 16 * 8; i += TL_NT) { + const int c = i / (16 * 8); + const int r = i - c * 16 * 8; + const int ki = r / 8, k = r & 7; + PB[c * 128 + (ki >> 3) * 64 + (ki & 7) * 8 + k] = + (k < 7) ? wpo[(c * 16 + ki) * 7 + k] : __float2half(0.f); + } + const float bias_i0 = bin[er0], bias_i1 = bin[er0 + 8]; + const float bias_s0 = (er0 < 32) ? bsk[er0] : 0.f; + const float bias_s1 = (er0 + 8 < 32) ? bsk[er0 + 8] : 0.f; + const float bpo0 = bpo[0]; + constexpr int UCOL = TL_UW / 32; + constexpr int UUNIT = 4 * UCOL; + constexpr int UWARP = TL_NT / 32; + constexpr int UITER = (UUNIT + UWARP - 1) / UWARP; + const int uwarp = tid >> 5; + const int uq = lane & 3; + int ugc[UITER], uoc[UITER]; +#pragma unroll + for (int t = 0; t < UITER; ++t) { + const int bi = uwarp + t * UWARP; + const int gc = bi / UCOL; + ugc[t] = gc; + uoc[t] = bi - gc * UCOL; + } + const __half *const urow = + U + (size_t)(lane >> 2) * SC4 + 8 * uq; + __syncthreads(); + PDL_ENTER(); +#define TL_STAGE_U(PUD,UB) { _Pragma("unroll") for (int t2 = 0; t2 < UITER; t2 += 2) { uint32_t u[2][4]; _Pragma("unroll") for (int d = 0; d < 2; ++d) { const int t = t2 + d; const int n = (UB) + uoc[t] * 32; const int gn = n + 8 * uq; const bool ok = gn >= 0 && gn < N4; int4 v = make_int4(0, 0, 0, 0); if (ok) v = *reinterpret_cast( urow + (size_t)ugc[t] * (8 * (size_t)SC4) + n); u[d][0] = v.x; u[d][1] = v.y; u[d][2] = v.z; u[d][3] = v.w; } _Pragma("unroll") for (int d = 0; d < 2; ++d) { const int t = t2 + d; uint32_t *q = u[d]; uint32_t x, y; x = (uq & 1) ? q[0] : q[1]; y = (uq & 1) ? q[2] : q[3]; x = __shfl_xor_sync(0xffffffffu, x, 1); y = __shfl_xor_sync(0xffffffffu, y, 1); if (uq & 1) { q[0] = x; q[2] = y; } else { q[1] = x; q[3] = y; } x = (uq & 2) ? q[0] : q[2]; y = (uq & 2) ? q[1] : q[3]; x = __shfl_xor_sync(0xffffffffu, x, 2); y = __shfl_xor_sync(0xffffffffu, y, 2); if (uq & 2) { q[0] = x; q[1] = y; } else { q[2] = x; q[3] = y; } tl_stmatrix4((PUD) + ugc[t] * TL_LBO_U + (uoc[t] * 32 + lane) * 8, q, true); } } if (uwarp < 4) { const int n = (UB) + 512; const int gn = n + 2 * (lane & 3); uint32_t q = 0; if (gn >= 0 && gn + 1 < N4) q = *reinterpret_cast( U + (size_t)(uwarp * 8 + (lane >> 2)) * SC4 + gn); tl_stmatrix1t((PUD) + uwarp * TL_LBO_U + (512 + (lane & 7)) * 8, q); } } + TL_STAGE_U(PU, (int)blockIdx.x * TL_BN - 16) + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + int ubuf = 0; + for (int tile = blockIdx.x; tile < NTILE; tile += gridDim.x, ubuf ^= 1) { + const int g0 = tile * TL_BN; + __half *const PUc = PU + ubuf * TL_PU_E; + __half *const PUn = PU + (ubuf ^ 1) * TL_PU_E; + { + const int b0 = g0 / L4; + for (int s = 0; s < 2; ++s) { + const int sb = b0 + s; + if (sb >= B) break; + const int cv = (int)((long long)sb * L4 + native_true_tokens(sb,T4) - g0); + if (cv < 0 || cv >= TL_BN) continue; + for (int i = tid; i < 32 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + ns_u[((size_t)sb * 32 + co) * 6 + h] = + PUc[(co >> 3) * TL_LBO_U + + (cv + h + 16) * 8 + (co & 7)]; + } + } + } + { + float acc[64]; +#pragma unroll + for (int i = 0; i < 64; ++i) + acc[i] = ((i & 2) == 0) ? bias_i0 : bias_i1; + const int c0 = wg * 128; + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int c = 0; c < 2; ++c) { + const uint64_t da0 = + wg_desc(AI + c * 7 * 1024, 128u, TL_SBO_A * 2u); + const uint64_t db0 = + wg_desc(PUc + (2 * c) * TL_LBO_U + (c0 + 2) * 8, + TL_LBO_U * 2u, 128u); +#pragma unroll + for (int k = 0; k < 7; ++k) { + const uint64_t da = da0 + 128 * k; + const uint64_t db = db0 + k; + WGMMA_M64N128K16_T(acc, da, db, 0, 0); + } + } + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); + if (tile + (int)gridDim.x < NTILE) + TL_STAGE_U(PUn, (tile + (int)gridDim.x) * TL_BN - 16) + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + if (er0 < 32) { + const float a0 = sh_ain[er0], r0 = sh_ari[er0]; + const float a1 = sh_ain[er0 + 8], r1 = sh_ari[er0 + 8]; +#pragma unroll + for (int i = 0; i < 64; ++i) { + const int h = (i & 2) >> 1; + const float v = acc[i]; + const float sn = __sinf((h ? a1 : a0) * v); + acc[i] = __fmaf_rn(sn * sn, h ? r1 : r0, v); + } + } else { +#pragma unroll + for (int i = 0; i < 64; ++i) { + const float v = acc[i]; + acc[i] = v > 0.f ? v : 0.01f * v; + } + } +#pragma unroll + for (int j = 0; j < 16; j += 2) { + uint32_t r[4]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + const __half2 hv = __floats2half2_rn( + acc[jj * 4 + h * 2 + 0], acc[jj * 4 + h * 2 + 1]); + r[m] = *reinterpret_cast(&hv); + } + const int m_addr = lane >> 3, row = lane & 7; + const int jj = j + (m_addr >> 1), h = m_addr & 1; + tl_stmatrix4(PM + (warp * 2 + h) * TL_LBO_M + + (c0 + jj * 8 + row) * 8, + r, true); + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + { + const int base = g0 - 8; + const int b0 = (base >= 0) ? base / L4 : 0; + const int cm0 = (int)((long long)b0 * L4 - base); + const int cm1 = cm0 + L4; + const bool hit0 = (b0 < B) && (cm0 + 6 > 0) && (cm0 < TL_MW); + const bool hit1 = (b0 + 1 < B) && (cm1 + 6 > 0) && (cm1 < TL_MW); + if (hit0 || hit1) { + for (int s = 0; s < 2; ++s) { + if (!(s ? hit1 : hit0)) continue; + const int sb = b0 + s, cm = s ? cm1 : cm0; + for (int i = tid; i < 64 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + const int c = cm + h; + if (c >= 0 && c < TL_MW) + PM[(co >> 3) * TL_LBO_M + c * 8 + (co & 7)] = + st_m4[((size_t)sb * 64 + co) * 6 + h]; + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + } + } + { + const int b0 = g0 / L4; + for (int s = 0; s < 2; ++s) { + const int sb = b0 + s; + if (sb >= B) break; + const int cv = (int)((long long)sb * L4 + native_true_tokens(sb,T4) - g0); + if (cv < 0 || cv >= TL_BN) continue; + for (int i = tid; i < 64 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + ns_m4[((size_t)sb * 64 + co) * 6 + h] = + PM[(co >> 3) * TL_LBO_M + + (cv + h + 8) * 8 + (co & 7)]; + } + } + } + { + const int c0 = wg * 128; + if (wg == 3) + tl_skip_band<120>(PUc, PM, AS, PV, c0, warp, lane, er0, sh_ask, + sh_ari_v, bias_s0, bias_s1); + else + tl_skip_band<128>(PUc, PM, AS, PV, c0, warp, lane, er0, sh_ask, + sh_ari_v, bias_s0, bias_s1); + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + { + const int b0 = g0 / L4; + const int cv0 = (int)((long long)b0 * L4 - g0), cv1 = cv0 + L4; + const bool hv0 = (b0 < B) && (cv0 + 6 > 0) && (cv0 < TL_VW); + const bool hv1 = (b0 + 1 < B) && (cv1 + 6 > 0) && (cv1 < TL_VW); + if (hv0 || hv1) { + for (int s = 0; s < 2; ++s) { + if (!(s ? hv1 : hv0)) continue; + const int sb = b0 + s, cv = s ? cv1 : cv0; + for (int i = tid; i < 32 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + const int c = cv + h; + if (c >= 0 && c < TL_VW) + PM[(co >> 3) * TL_LBO_V + c * 8 + (co & 7)] = + st_v[((size_t)sb * 32 + co) * 6 + h]; + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + } + } + { + const int b0 = g0 / L4; + for (int s = 0; s < 2; ++s) { + const int sb = b0 + s; + if (sb >= B) break; + const int cv = (int)((long long)sb * L4 + native_true_tokens(sb,T4) - g0); + if (cv < 0 || cv >= TL_BN) continue; + for (int i = tid; i < 32 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + ns_v[((size_t)sb * 32 + co) * 6 + h] = + PM[(co >> 3) * TL_LBO_V + (cv + h) * 8 + (co & 7)]; + } + } + } + __syncthreads(); + constexpr int PC_LD = 10; + { + float *const PC = reinterpret_cast(PV); +#pragma unroll + for (int round = 0; round < 2; ++round) { + float acc[4]; + const int p0 = round * 256 + wg * 64; + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); + const uint64_t da0 = + wg_desc(PM + p0 * 8, TL_LBO_V * 2u, 128u); + const uint64_t db0 = wg_desc(PB, 128u, 128u); + const uint64_t da1 = + wg_desc(PM + 2 * TL_LBO_V + p0 * 8, + TL_LBO_V * 2u, 128u); + const uint64_t db1 = wg_desc(PB + 128, 128u, 128u); + WGMMA_M64N8K16_ZERO_T(acc, da0, db0, 0, 1); + WGMMA_M64N8K16_T(acc, da1, db1, 0, 1); + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + const int row0 = p0 + warp * 16 + (lane >> 2); + const int col = (lane & 3) * 2; + *reinterpret_cast(PC + row0 * PC_LD + col) = + make_float2(acc[0], acc[1]); + *reinterpret_cast(PC + (row0 + 8) * PC_LD + col) = + make_float2(acc[2], acc[3]); + } + } + __syncthreads(); + { + const float *const PC = reinterpret_cast(PV); + const int p0 = tid; + const bool live = p0 < TL_BN; + const unsigned live_mask = __ballot_sync(0xffffffffu, live); + if (live) { + const float s0 = __fadd_rn(PC[(p0 + 0) * PC_LD + 0], + PC[(p0 + 1) * PC_LD + 1]); + const float s1 = __fadd_rn(PC[(p0 + 2) * PC_LD + 2], + PC[(p0 + 3) * PC_LD + 3]); + const float s2 = __fadd_rn(PC[(p0 + 4) * PC_LD + 4], + PC[(p0 + 5) * PC_LD + 5]); + const float s3 = PC[(p0 + 6) * PC_LD + 6]; + float y0 = __fadd_rn(__fadd_rn(s0, s1), + __fadd_rn(s2, s3)); + y0 = fminf(fmaxf(y0 + bpo0, -1.f), 1.f); + const float y1 = __shfl_down_sync(live_mask, y0, 1); + if ((p0 & 1) == 0) { + const int b0 = g0 / L4; + const int rr = g0 - b0 * L4; + int t0 = rr + p0, bo0 = b0; + int t1 = t0 + 1, bo1 = bo0; + if (t0 >= L4) { t0 -= L4; ++bo0; } + if (t1 >= L4) { t1 -= L4; ++bo1; } + if (bo0 == bo1 && bo0 < B && t1 < T4) { + *reinterpret_cast<__half2 *>( + wave + (size_t)bo0 * T4 + t0) = + __floats2half2_rn(y0, y1); + } else { + if (bo0 < B && t0 < T4) + wave[(size_t)bo0 * T4 + t0] = __float2half(y0); + if (bo1 < B && t1 < T4) + wave[(size_t)bo1 * T4 + t1] = __float2half(y1); + } + } + } + } + } +#undef TL_STAGE_U +} diff --git a/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wg2_kernel.inc b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wg2_kernel.inc new file mode 100644 index 0000000..4ffa7ad --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wg2_kernel.inc @@ -0,0 +1,280 @@ +__device__ __forceinline__ void cpa16(void *dsm, const void *src, bool pred) { + unsigned a = static_cast(__cvta_generic_to_shared(dsm)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(a), + "l"(src), "r"(pred ? 16 : 0) + : "memory"); +} +template +__device__ __forceinline__ void cpa16o(unsigned a, const void *src, bool pred) { + asm volatile("cp.async.cg.shared.global [%0+%3], [%1], 16, %2;\n" ::"r"(a), + "l"(src), "r"(pred ? 16 : 0), "n"(OFF) + : "memory"); +} +template +__device__ __noinline__ +void wg2_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int BK = CIB * KT; + constexpr int WG_BM = 64 * MSPLIT; + constexpr int WG_BNT = WG_BN * NSPLIT; + constexpr int NTS = 128 * MSPLIT * NSPLIT; + constexpr int WG_NT = NTS * KSPLIT; + constexpr int NCRAW = WG_BNT / 4; + constexpr int EPN = + (WG_NT % NCRAW == 0) ? WG_BNT : (((WG_BNT + 127) / 128) * 128); + constexpr int WG_CLD = EPN + 8; + constexpr bool NB8_P2 = ((WG_BN / 8) & (WG_BN / 8 - 1)) == 0; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int ST_E = SA * A_E + SB * NSPLIT * B_E; + constexpr int SBO_AE = KG * WG_SLOT; + constexpr int LBO_BE = (WG_BN / 8) * WG_SLOT; + constexpr int NB8 = WG_BN / 8; + constexpr bool ALDG = (WG_BN <= 64); + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *smem = reinterpret_cast<__half *>(arena); + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int kgrp = tid / NTS; + const int stid = tid - kgrp * NTS; + const int wgl = stid >> 7; + const int mg = wgl % MSPLIT; + const int ngid = wgl / MSPLIT; + const int warp = (stid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * WG_BM; + const int nb = blockIdx.y * WG_BNT; + const int n0 = nb + ngid * WG_BN; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + const int NCH = (Kg + BK - 1) / BK; + if (tid < WG_BM) { + const int co = m0 + tid; + if (do_snake && co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + for (int i = tid; i < EPN; i += WG_NT) { + const int n = nb + i; + colok[i] = (i < WG_BNT && n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mg * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (kgrp == 0 && eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (kgrp == 0 && eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + constexpr int KGX = ALDG ? (KG <= 8 ? 8 : (KG <= 16 ? 16 : 32)) : KG; + constexpr int ATOT = WG_BM * KGX; + constexpr int AREG = (ATOT + NTS - 1) / NTS; + constexpr int BTOT = NSPLIT * CIB * NB8; + constexpr int BREG = (BTOT + NTS - 1) / NTS; + const __half *aptr[AREG]; + int aoff[AREG]; + int akmax[AREG]; +#pragma unroll + for (int r = 0; r < AREG; ++r) { + const int idx = stid + r * NTS; + const int m = idx / KGX, v = idx - (idx / KGX) * KGX; + const int co = m0 + m; + const bool live = (ATOT % NTS == 0 || idx < ATOT) && (co < Cout) && + v < KG; + aoff[r] = (m >> 3) * SBO_AE + (v < KG ? v : 0) * WG_SLOT + (m & 7) * 8; + akmax[r] = live ? (Kg - v * 8) : 0; + aptr[r] = w + (size_t)(live ? co : 0) * Kg + v * 8; + } + const __half *bptr[BREG]; + int bcimax[BREG], bnc[BREG], bcil[BREG], bgrp[BREG]; +#pragma unroll + for (int r = 0; r < BREG; ++r) { + const int idx = stid + r * NTS; + const int gg = idx / (CIB * NB8); + const int rr = idx - gg * (CIB * NB8); + const int cil = rr / NB8, nc = rr - (rr / NB8) * NB8; + bcil[r] = cil; + bnc[r] = nc; + bgrp[r] = gg; + const bool live = (BTOT % NTS == 0 || idx < BTOT); + bcimax[r] = live ? (Cin - cil) : 0; + bptr[r] = in + (size_t)cil * SCin + nb + gg * WG_BN + nc * 8; + } + int4 br0[BREG], br1[BREG]; + int4 ar[ALDG ? AREG : 1]; + uint32_t edge[BREG]; + __half *const kaB = smem + kgrp * ST_E; + __half *const kbB = kaB + SA * A_E; + const unsigned barid = 1u + kgrp + (unsigned)(KSPLIT * native_subid()); +#define WG2_BAR() asm volatile("bar.sync %0, %1;\n" ::"r"(barid), "r"(NTS)) +#define WG2_LOADB(c) { const int _ci0 = (c) * CIB; _Pragma("unroll") for (int r = 0; r < BREG; ++r) { br0[r] = make_int4(0, 0, 0, 0); if constexpr (KT == 3) edge[r] = 0; else br1[r] = make_int4(0, 0, 0, 0); if (_ci0 < bcimax[r]) { const int4 *src = reinterpret_cast( bptr[r] + (size_t)_ci0 * SCin); br0[r] = src[0]; if constexpr (KT == 3) { if (!NB8_P2 || bnc[r] == NB8 - 1) edge[r] = src[1].x; } else { br1[r] = src[1]; } } } } +#define WG2_LOADA(c) if constexpr (ALDG) { const int _kg0 = (c) * BK; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { ar[r] = make_int4(0, 0, 0, 0); if (_kg0 < akmax[r]) ar[r] = *reinterpret_cast(aptr[r] + _kg0); } } +#define WG2_PUTA(c,AD) { _Pragma("unroll") for (int r = 0; r < AREG; ++r) { if constexpr (ALDG) { if (akmax[r] > 0) *reinterpret_cast((AD) + aoff[r]) = ar[r]; } else { const int _kg0 = (c) * BK; if (akmax[r] > 0) cpa16((AD) + aoff[r], aptr[r] + _kg0, _kg0 < akmax[r]); } } if constexpr (!ALDG) asm volatile("cp.async.commit_group;\n" ::: "memory"); } +#define WG2_STSB(BD) { _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int cil = bcil[r], nc = bnc[r]; unsigned int u[8]; u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; if constexpr (KT == 3) { if constexpr (NB8_P2) { u[4] = __shfl_down_sync(0xffffffffu, u[0], 1, NB8); if (nc == NB8 - 1) u[4] = edge[r]; } else { u[4] = edge[r]; } u[5] = u[6] = u[7] = 0; } else { u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; } if (BTOT % NTS == 0 || (stid + r * NTS) < BTOT) { _Pragma("unroll") for (int k = 0; k < KT; ++k) { const int kh = k >> 1; int4 o; if (k & 1) { o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); } else { o.x = u[kh + 0]; o.y = u[kh + 1]; o.z = u[kh + 2]; o.w = u[kh + 3]; } const int row = cil * KT + k; *reinterpret_cast((BD) + bgrp[r] * B_E + nc * WG_SLOT + (row >> 3) * LBO_BE + (row & 7) * 8) = o; } } } } +#define WG2_MMA(AD,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t da0 = wg_desc((AD) + mg * 8 * SBO_AE, WG_SLOT * 2, SBO_AE * 2); const uint64_t db0 = wg_desc((BD) + ngid * B_E, LBO_BE * 2, WG_SLOT * 2); _Pragma("unroll") for (int kk = 0; kk < BK; kk += 16) { const uint64_t da = da0 + (kk >> 3) * (WG_SLOT / 8); const uint64_t db = db0 + (kk >> 3) * (LBO_BE / 8); if constexpr (WG_BN == 64) { WGMMA_M64N64K16(acc, da, db); } else if constexpr (WG_BN == 256) { WGMMA_M64N256K16(acc, da, db); } else if constexpr (WG_BN == 104) { WGMMA_M64N104K16(acc, da, db); } else if constexpr (WG_BN == 112) { WGMMA_M64N112K16(acc, da, db); } else if constexpr (WG_BN == 120) { WGMMA_M64N120K16(acc, da, db); } else if constexpr (WG_BN == 224) { WGMMA_M64N224K16(acc, da, db); } else { static_assert(WG_BN == 128, "no wgmma macro for WG_BN"); WGMMA_M64N128K16(acc, da, db); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + constexpr int SMAX = SA > SB ? SA : SB; + WG2_LOADA(kgrp) + WG2_LOADB(kgrp) +#pragma unroll 1 + for (int s = 0; s < SMAX; ++s) { + const int c = s * KSPLIT + kgrp; + if (s < SA) { + if (c < NCH) WG2_PUTA(c, kaB + s * A_E) + if (c + KSPLIT < NCH) WG2_LOADA(c + KSPLIT) + } + if (s < SB) { + if (c < NCH) WG2_STSB(kbB + s * NSPLIT * B_E) + if (c + KSPLIT < NCH) WG2_LOADB(c + KSPLIT) + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + const int NIT = (NCH - kgrp + KSPLIT - 1) / KSPLIT; + for (int i = 0; i < NIT; ++i) { + if constexpr (!ALDG) + asm volatile("cp.async.wait_group %0;\n" ::"n"(SA - 2) : "memory"); + WG2_BAR(); + WG2_MMA(kaB + (i % SA) * A_E, kbB + (i % SB) * NSPLIT * B_E) + asm volatile("wgmma.wait_group.sync.aligned 1;\n" ::: "memory"); + if constexpr (NTS > 128) WG2_BAR(); + if (i >= 1) { + const int ca = (i + SA - 1) * KSPLIT + kgrp; + if (ca < NCH) WG2_PUTA(ca, kaB + ((i - 1) % SA) * A_E) + if (ALDG && ca + KSPLIT < NCH) WG2_LOADA(ca + KSPLIT) + const int cb = (i + SB - 1) * KSPLIT + kgrp; + if (cb < NCH) { + WG2_STSB(kbB + ((i - 1) % SB) * NSPLIT * B_E) + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if (cb + KSPLIT < NCH) WG2_LOADB(cb + KSPLIT) + } else if constexpr (ALDG) { + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + } + } + } + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); +#undef WG2_LOADB +#undef WG2_LOADA +#undef WG2_PUTA +#undef WG2_STSB +#undef WG2_MMA +#undef WG2_BAR + __syncthreads(); + float *Ck = Cs + kgrp * WG_BM * WG_CLD; + const int cb = (lane & 3) * 2; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Ck + (er0 + h * 8) * WG_CLD + + ngid * WG_BN + j * 8 + cb) = v; + } + } + __syncthreads(); + constexpr int NC = EPN / 4; + constexpr int MSTEP = WG_NT / NC; + constexpr int MITER = WG_BM / MSTEP; + constexpr int UNR = (MITER % 4 == 0) ? 4 : ((MITER % 2 == 0) ? 2 : 1); + static_assert(MITER % UNR == 0, "row batch must divide the row count"); + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + const int n = nb + lc; +#pragma unroll 1 + for (int t0 = 0; t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + v[u] = *reinterpret_cast(Cs + m * WG_CLD + lc); + if (resid != nullptr) { + const int co = m0 + m; + const int cc = co < Cout ? co : 0; + const __half2 *rp = reinterpret_cast( + resid + (size_t)cc * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + const int co = m0 + m; + if (co >= Cout) continue; + float4 vv = v[u]; +#pragma unroll + for (int g = 1; g < KSPLIT; ++g) { + float4 p = *reinterpret_cast( + Cs + g * WG_BM * WG_CLD + m * WG_CLD + lc); + vv.x += p.x; vv.y += p.y; vv.z += p.z; vv.w += p.w; + } + if (resid != nullptr) { + vv.x += __half2float(r0[u].x); vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (do_snake) { + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + } + __half *op = out + (size_t)co * SCout + n + H; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B * H; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int total = Cout * B * pad_len; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} diff --git a/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wg3_kernel.inc b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wg3_kernel.inc new file mode 100644 index 0000000..fa55b32 --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wg3_kernel.inc @@ -0,0 +1,177 @@ +template +__device__ __noinline__ void wg3_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int WG_BNT = WG_BN * NSG; + constexpr int WG_CLD = WG_BNT + 8; + constexpr int WG_NT = 128 * NWG * NSG; + constexpr int WG_BM = 64 * NWG; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int SBO_AE = KG * WG_SLOT; + constexpr int LBO_BE = (WG_BN / 8) * WG_SLOT; + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + constexpr int ST_E = A_E + NSG * B_E; + __half *As0 = reinterpret_cast<__half *>(arena); + __half *Bs0 = As0 + A_E; + __half *As1 = (STAGES == 2) ? As0 + ST_E : As0; + __half *Bs1 = (STAGES == 2) ? Bs0 + ST_E : Bs0; + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int wgid = tid >> 7; + const int mgid = wgid % NWG, ngid = wgid / NWG; + const int warp = (tid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * WG_BM; + constexpr int WGROWS = 8; + const int nb = blockIdx.y * WG_BNT; + const int n0 = nb + ngid * WG_BN; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + if (tid < WG_BM) { + const int co = m0 + tid; + if (do_snake && co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + if (tid < WG_BNT) { + const int n = nb + tid; + colok[tid] = (n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mgid * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + constexpr int ATOT = WG_BM * KG; + constexpr int AREG = (ATOT + WG_NT - 1) / WG_NT; + constexpr int CIB = BK / KT; + constexpr int BTOT = NSG * CIB * (WG_BN / 8); + constexpr int BREG = (BTOT + WG_NT - 1) / WG_NT; + int4 ar[AREG]; + int4 br0[BREG], br1[BREG]; + uint32_t edge[BREG]; +#define WG_LOADG(kg0) { const int _ci0 = (kg0) / KT; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; ar[r] = make_int4(0, 0, 0, 0); if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; const int co = m0 + m, kg = (kg0) + v * 8; if (co < Cout && kg < Kg) ar[r] = *reinterpret_cast( w + (size_t)co * Kg + kg); } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; br0[r] = make_int4(0, 0, 0, 0); if constexpr (KT == 3) edge[r] = 0; else br1[r] = make_int4(0, 0, 0, 0); if (BTOT % WG_NT == 0 || idx < BTOT) { const int _g = idx / (CIB * (WG_BN / 8)); const int _r = idx - _g * (CIB * (WG_BN / 8)); const int ci = _ci0 + _r / (WG_BN / 8); const int nc = _r % (WG_BN / 8); if (ci < Cin) { const int4 *src = reinterpret_cast( in + (size_t)ci * SCin + nb + _g * WG_BN + nc * 8); br0[r] = src[0]; if constexpr (KT == 3) { if (nc == WG_BN / 8 - 1) edge[r] = src[1].x; } else { br1[r] = src[1]; } } } } } + WG_LOADG(0) +#define WG_STS(AD,BD) { _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; *reinterpret_cast( (AD) + (m >> 3) * SBO_AE + v * WG_SLOT + (m & 7) * 8) = ar[r]; } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; if (BTOT % WG_NT == 0 || idx < BTOT) { const int _g = idx / (CIB * (WG_BN / 8)); const int _r = idx - _g * (CIB * (WG_BN / 8)); const int cil = _r / (WG_BN / 8), nc = _r % (WG_BN / 8); unsigned int u[8]; u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; if constexpr (KT == 3) { u[4] = __shfl_down_sync(0xffffffffu, u[0], 1, WG_BN / 8); if (nc == WG_BN / 8 - 1) u[4] = edge[r]; u[5] = u[6] = u[7] = 0; } else { u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; } _Pragma("unroll") for (int k = 0; k < KT; ++k) { const int kh = k >> 1; int4 o; if (k & 1) { o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); } else { o.x = u[kh + 0]; o.y = u[kh + 1]; o.z = u[kh + 2]; o.w = u[kh + 3]; } const int row = cil * KT + k; *reinterpret_cast((BD) + _g * B_E + nc * WG_SLOT + (row >> 3) * LBO_BE + (row & 7) * 8) = o; } } } if constexpr (WG_BN == 64) { asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); } } +#define WG_MMA(AD,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t da0 = wg_desc((AD) + mgid * WGROWS * SBO_AE, WG_SLOT * 2, SBO_AE * 2); const uint64_t db0 = wg_desc((BD) + ngid * B_E, LBO_BE * 2, WG_SLOT * 2); _Pragma("unroll") for (int kk = 0; kk < BK; kk += 16) { const uint64_t da = da0 + (kk >> 3) * (WG_SLOT / 8); const uint64_t db = db0 + (kk >> 3) * (LBO_BE / 8); if constexpr (WG_BN == 64) { WGMMA_M64N64K16(acc, da, db); } else { WGMMA_M64N128K16(acc, da, db); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + WG_STS(As0, Bs0) + if (BK < Kg) WG_LOADG(BK) + __syncthreads(); + int cur = 0; + for (int kg0 = 0; kg0 < Kg; kg0 += BK) { + if (cur == 0) { WG_MMA(As0, Bs0) } else { WG_MMA(As1, Bs1) } + if (kg0 + BK < Kg) { + if (STAGES == 2) + asm volatile("wgmma.wait_group.sync.aligned 1;\n" ::: "memory"); + else + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + __syncthreads(); + if (cur == 0) { WG_STS(As1, Bs1) } else { WG_STS(As0, Bs0) } + if (kg0 + 2 * BK < Kg) WG_LOADG(kg0 + 2 * BK) + __syncthreads(); + } else { + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + } + if (STAGES == 2) cur ^= 1; + } +#undef WG_LOADG +#undef WG_STS +#undef WG_MMA + __syncthreads(); + const int r0 = er0; + const int cb = (lane & 3) * 2; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Cs + (r0 + h * 8) * WG_CLD + + ngid * WG_BN + j * 8 + cb) = v; + } + } + __syncthreads(); + constexpr int NC = WG_BNT / 4; + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + for (int m = lr; m < WG_BM; m += WG_NT / NC) { + const int co = m0 + m; + if (co >= Cout) continue; + const int n = nb + lc; + float4 v = *reinterpret_cast(Cs + m * WG_CLD + lc); + __half *op = out + (size_t)co * SCout + n + H; + if (resid != nullptr) { + const __half2 *rp = reinterpret_cast( + resid + (size_t)co * SCout + n + H); + __half2 r0h = rp[0], r1h = rp[1]; + v.x += __half2float(r0h.x); v.y += __half2float(r0h.y); + v.z += __half2float(r1h.x); v.w += __half2float(r1h.y); + } + if (do_snake) { + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + } + if (ok4) { + __half2 o0 = __floats2half2_rn(v.x, v.y); + __half2 o1 = __floats2half2_rn(v.z, v.w); + reinterpret_cast<__half2 *>(op)[0] = o0; + reinterpret_cast<__half2 *>(op)[1] = o1; + } else { + const float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B * H; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int total = Cout * B * pad_len; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} diff --git a/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wg4_step_widths.inc b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wg4_step_widths.inc new file mode 100644 index 0000000..9772bda --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wg4_step_widths.inc @@ -0,0 +1,459 @@ +template +__device__ __noinline__ void wg4_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int CoutP, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + if constexpr (PDL && !LATE_PDL) PDL_ENTER(); + constexpr int NT = 128 * MSPLIT * NSPLIT; + constexpr int BM = 64 * MSPLIT; + constexpr int BNT = WG_BN * NSPLIT; + constexpr int NCRAW = BNT / 4; + constexpr int EPN = (NT % NCRAW == 0) ? BNT : (((BNT + 127) / 128) * 128); + constexpr int WG_CLD = EPN + 8; + constexpr int CG = CIB / 8; + constexpr int NBR = BNT + KT - 1; + constexpr int NBRP = BC8 ? ((NBR + 7) & ~7) : ((NBR + 31) & ~31); + constexpr int UCOL = NBRP / 32; + constexpr int NUNIT = CG * UCOL; + constexpr int NWARP = NT / 32; + constexpr int BITER = (NUNIT + NWARP - 1) / NWARP; + constexpr int B_E = CIB * NBRP; + constexpr int ST_E = B_E; + constexpr int LBO_B = NBRP * 8; + constexpr int ABUF = ABUF_; + static_assert(!OC8 || PREACT, + "K=7 C8 output is only used for residual-free pair heads"); + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *smem = reinterpret_cast<__half *>(arena); + __half *Hs = smem; + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int wgid = tid >> 7; + const int mg = wgid % MSPLIT, ng = wgid / MSPLIT; + const int warp = (tid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * BM; + const int nb = blockIdx.y * BNT; + const int H = KT - 1; + const int half = Cout >> 1; + const int NCH = (Cin + CIB - 1) / CIB; + for (int i = tid; i < BM; i += NT) { + const int co = m0 + i; + if (co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[i] = a; + row_ri[i] = __frcp_rn(a + 1e-9f); + } + } + for (int i = tid; i < EPN; i += NT) { + const int n = nb + i; + colok[i] = (i < BNT && n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mg * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + const int mtile = blockIdx.x * MSPLIT + mg; + const __half *const apt = w + ((size_t)mtile * NCH * KT) * 1024 + + (size_t)(tid & 127) * 8; + uint32_t areg[ABUF][KT][4]; + const int wi = tid >> 5; + int bdst[BITER], boff[BITER]; + bool blive[BITER]; + const __half *const bbase = in + (size_t)(lane >> 2) * SCin + nb + + 8 * (lane & 3); +#pragma unroll + for (int t = 0; t < BITER; ++t) { + const int u = wi + t * NWARP; + const int gc = (u < NUNIT) ? u / UCOL : 0; + const int cu = (u < NUNIT) ? u - gc * UCOL : 0; + blive[t] = u < NUNIT; + bdst[t] = gc * LBO_B + (cu * 32 + lane) * 8; + boff[t] = gc * 8 * SCin + cu * 32; + } + const size_t SC8i = (size_t)SCin * 8; + constexpr int BQ = B_E / 8; + constexpr int BCREG = (BQ + NT - 1) / NT; + int cdst[BCREG]; + const __half *csrc[BCREG]; + bool clive[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) { + const int q = tid + r * NT; + const int gc = q / NBRP, rem = q - gc * NBRP; + clive[r] = ((BQ % NT == 0) || q < BQ) && gc < CG; + cdst[r] = clive[r] ? gc * LBO_B + rem * 8 : 0; + csrc[r] = in + (clive[r] ? (size_t)gc * SC8i + + (size_t)nb * 8 + rem * 8 : 0); + } + uint32_t br[BITER][4]; + constexpr int CPA_STAGES = 4; + constexpr int BSTAGES = BC8 ? CPA_STAGES : STAGES; + static_assert(STAGES >= CPA_STAGES, "C8 B-panel pipeline allocation"); + static_assert(ABUF == 2 && BSTAGES == 4, "wg4 unrolled pipeline shape"); + uint64_t dslot[BSTAGES]; +#pragma unroll + for (int s = 0; s < BSTAGES; ++s) + dslot[s] = wg_desc(smem + s * ST_E + ng * WG_BN * 8, LBO_B * 2u, 128u); + uint32_t csm0[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) + csm0[r] = static_cast(__cvta_generic_to_shared(smem + cdst[r])); + uint32_t bsm0[BITER]; +#pragma unroll + for (int t = 0; t < BITER; ++t) + bsm0[t] = static_cast(__cvta_generic_to_shared(smem + bdst[t])); + const size_t CSTEP = (size_t)CG * SC8i; + const __half *crol[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) crol[r] = csrc[r] + CPA_STAGES * CSTEP; + const __half *arol = apt + (size_t)ABUF * (KT * 1024); +#define WG4_CPB(c,BD) { const size_t _o = (size_t)(c) * CG * SC8i; _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16((BD) + cdst[r], csrc[r] + _o, true); if ((c) + 1 == CPA_STAGES) { _Pragma("unroll") for (int r = 0; r < BCREG; ++r) csrc[r] += (size_t)CPA_STAGES * CSTEP; } } +#define WG4_CPB_R(SLOT) { _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16o<(SLOT) * ST_E * 2>(csm0[r], csrc[r], true); _Pragma("unroll") for (int r = 0; r < BCREG; ++r) csrc[r] += CSTEP; } +#define WG4_LOADB(c) { const size_t _o = (size_t)(c) * CIB * SCin; _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { const int4 _v = *reinterpret_cast( bbase + _o + boff[t]); br[t][0] = _v.x; br[t][1] = _v.y; br[t][2] = _v.z; br[t][3] = _v.w; } } } +#define WG4_XPOSE(u) { const int _q = lane & 3; uint32_t _x, _y; _x = (_q & 1) ? u[0] : u[1]; _y = (_q & 1) ? u[2] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 1); _y = __shfl_xor_sync(0xffffffffu, _y, 1); if (_q & 1) { u[0] = _x; u[2] = _y; } else { u[1] = _x; u[3] = _y; } _x = (_q & 2) ? u[0] : u[2]; _y = (_q & 2) ? u[1] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 2); _y = __shfl_xor_sync(0xffffffffu, _y, 2); if (_q & 2) { u[0] = _x; u[1] = _y; } else { u[2] = _x; u[3] = _y; } } +#define WG4_STOREB(BD) { _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { WG4_XPOSE(br[t]) const uint32_t _a = static_cast( __cvta_generic_to_shared((BD) + bdst[t])); asm volatile( "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 " "[%0], {%1,%2,%3,%4};\n" ::"r"(_a), "r"(br[t][0]), "r"(br[t][1]), "r"(br[t][2]), "r"(br[t][3])); } } } +#define WG4_STOREB_R(SLOT) { _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { WG4_XPOSE(br[t]) asm volatile( "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 " "[%0+%5], {%1,%2,%3,%4};\n" ::"r"(bsm0[t]), "r"(br[t][0]), "r"(br[t][1]), "r"(br[t][2]), "r"(br[t][3]), "n"((SLOT) * ST_E * 2)); } } } +#define WG4_LOADA_R(BUF) { _Pragma("unroll") for (int kk = 0; kk < 4; ++kk) { const int4 _v = *reinterpret_cast(arol + kk * 1024); areg[0][kk][0] = _v.x; areg[0][kk][1] = _v.y; areg[0][kk][2] = _v.z; areg[0][kk][3] = _v.w; } } +#define WG4_MMA_R(BUF,SLOT) { const uint64_t db0 = dslot[0] + (uint64_t)(SLOT) * (ST_E / 8); asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); _Pragma("unroll") for (int kk = 0; kk < 4; ++kk) { WGMMA_RS_ANY_T(WG_BN, acc, areg[0][kk], db0 + kk, 0); } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); _Pragma("unroll") for (int kk = 0; kk < 4; ++kk) asm volatile("" : "+r"(areg[0][kk][0]), "+r"(areg[0][kk][1]), "+r"(areg[0][kk][2]), "+r"(areg[0][kk][3])); _Pragma("unroll") for (int kk = 4; kk < KT; ++kk) { const int4 _v = *reinterpret_cast(arol + kk * 1024); areg[0][kk - 4][0] = _v.x; areg[0][kk - 4][1] = _v.y; areg[0][kk - 4][2] = _v.z; areg[0][kk - 4][3] = _v.w; } asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); _Pragma("unroll") for (int kk = 4; kk < KT; ++kk) { WGMMA_RS_ANY_T(WG_BN, acc, areg[0][kk - 4], db0 + kk, 0); } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); arol += KT * 1024; } +#define WG4_LOADA(BUF,c) { if constexpr ((BUF) == 0) { const __half *_p = apt + (size_t)(c) * (KT * 1024); _Pragma("unroll") for (int kk = 0; kk < 4; ++kk) { const int4 _v = *reinterpret_cast(_p + kk * 1024); areg[0][kk][0] = _v.x; areg[0][kk][1] = _v.y; areg[0][kk][2] = _v.z; areg[0][kk][3] = _v.w; } arol = _p; } } +#define WG4_KEEPA(BUF) { _Pragma("unroll") for (int kk = 0; kk < KT - 4; ++kk) asm volatile("" : "+r"(areg[0][kk][0]), "+r"(areg[0][kk][1]), "+r"(areg[0][kk][2]), "+r"(areg[0][kk][3])); } +#define WG4_MMA(BUF,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t db0 = wg_desc((BD) + ng * WG_BN * 8, LBO_B * 2u, 128u); _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { WGMMA_RS_ANY_T(WG_BN, acc, areg[BUF][kk], db0 + kk, 0); } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + if constexpr (PDL && LATE_PDL) PDL_ENTER(); + constexpr int WGD = 1; + if constexpr (BC8) { +#pragma unroll 1 + for (int st = 0; st < CPA_STAGES; ++st) { + if (st < NCH) WG4_CPB(st, smem + st * ST_E) + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } + WG4_LOADA(0, 0) + if (ABUF > 1 && 1 < NCH) WG4_LOADA(1, 1) + if (ABUF > 2 && 2 < NCH) WG4_LOADA(2, 2) + } else { + WG4_LOADB(0) + WG4_LOADA(0, 0) + if (ABUF > 1 && 1 < NCH) WG4_LOADA(1, 1) + if (ABUF > 2 && 2 < NCH) WG4_LOADA(2, 2) + WG4_STOREB(smem) + if (1 < NCH) WG4_LOADB(1) + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + } +#define WG4_BODY(BUF,i,SLOT) { if constexpr (BC8) { asm volatile("cp.async.wait_group %0;\n" : : "n"(CPA_STAGES - 3) : "memory"); __syncthreads(); if ((i) >= 2) { if ((i) + 2 < NCH) { WG4_CPB_R(((SLOT) + 2) % CPA_STAGES) } asm volatile("cp.async.commit_group;\n" ::: "memory"); } } else { if (i) __syncthreads(); } WG4_MMA_R(BUF, SLOT) if constexpr (NSPLIT > 1 || BC8) { asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); WG4_KEEPA(0) if constexpr (NSPLIT > 1) __syncthreads(); } if constexpr (!BC8) { const int cb = (i) + 1; if (cb < NCH) { WG4_STOREB_R(((SLOT) + 1) % BSTAGES) asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); if (cb + 1 < NCH) WG4_LOADB(cb + 1) } } if ((i) + 1 < NCH) { if constexpr (NSPLIT == 1 && !BC8) { asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); WG4_KEEPA(0) } WG4_LOADA_R(0) } } + for (int i = 0; i < NCH; i += 4) { + WG4_BODY(0, i, 0) + if (i + 1 < NCH) WG4_BODY(1, i + 1, 1) + if (i + 2 < NCH) WG4_BODY(0, i + 2, 2) + if (i + 3 < NCH) WG4_BODY(1, i + 3, 3) + } + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + if constexpr (BC8) + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); +#undef WG4_CPB +#undef WG4_CPB_R +#undef WG4_STOREB_R +#undef WG4_LOADA_R +#undef WG4_MMA_R +#undef WG4_LOADB +#undef WG4_XPOSE +#undef WG4_STOREB +#undef WG4_LOADA +#undef WG4_KEEPA +#undef WG4_MMA +#undef WG4_BODY +#define TL4_SW(c,r) ((c) * BM + (((r) & ~7) ^ (((c) & 7) << 3)) + ((r) & 7)) + __syncthreads(); + const int cb2 = (lane & 3) * 2; + const bool snake = PREACT && (m0 + er0) < half; + const float sn_a0 = snake ? row_alpha[er0] : 0.f; + const float sn_r0 = snake ? row_ri[er0] : 0.f; + const float sn_a1 = snake ? row_alpha[er0 + 8] : 0.f; + const float sn_r1 = snake ? row_ri[er0 + 8] : 0.f; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + const int row = er0 + h * 8; + if constexpr (PREACT) { + if (snake) { + const float a = h ? sn_a1 : sn_a0; + const float ri = h ? sn_r1 : sn_r0; + const float s0 = __sinf(a * v.x); + const float s1 = __sinf(a * v.y); + v.x = __fmaf_rn(s0 * s0, ri, v.x); + v.y = __fmaf_rn(s1 * s1, ri, v.y); + } else { + v.x = v.x > 0.f ? v.x : 0.01f * v.x; + v.y = v.y > 0.f ? v.y : 0.01f * v.y; + } + if constexpr (OC8) { + const int c0 = ng * WG_BN + j * 8 + cb2; + Hs[TL4_SW(c0, row)] = __float2half(v.x); + Hs[TL4_SW(c0 + 1, row)] = __float2half(v.y); + } else { + *reinterpret_cast<__half2 *>( + Hs + row * WG_CLD + ng * WG_BN + j * 8 + cb2) = + __floats2half2_rn(v.x, v.y); + } + } else { + *reinterpret_cast( + Cs + row * WG_CLD + ng * WG_BN + j * 8 + cb2) = v; + } + } + } + __syncthreads(); + if constexpr (PREACT) { + if constexpr (OC8) { + constexpr int GRP = BM / 8; + constexpr int TOTG = GRP * BNT; + constexpr int GITER = (TOTG + NT - 1) / NT; + const size_t SC8o = (size_t)SCout * 8; +#pragma unroll + for (int t = 0; t < GITER; ++t) { + const int idx = tid + t * NT; + if (TOTG % NT != 0 && idx >= TOTG) continue; + const int g = idx / BNT; + const int j = idx - g * BNT; + if (!colok[j]) continue; + const int m0g = g * 8; + if (m0 + m0g >= Cout) continue; + *reinterpret_cast( + out + (size_t)((m0 + m0g) >> 3) * SC8o + + (size_t)(nb + j + H) * 8) = + *reinterpret_cast(Hs + TL4_SW(j, m0g)); + } + } else { + static_assert((BNT & 3) == 0, "four-column preactivated epilogue"); + constexpr int NC = BNT / 4; + constexpr int TOT = BM * NC; + for (int q = tid; q < TOT; q += NT) { + const int m = q / NC; + const int lc = (q - m * NC) * 4; + const int co = m0 + m; + const bool any4 = (colok[lc] | colok[lc + 1] | + colok[lc + 2] | colok[lc + 3]) != 0; + if (co >= Cout || !any4) continue; + const __half2 *sp = reinterpret_cast( + Hs + m * WG_CLD + lc); + const __half2 v0 = sp[0], v1 = sp[1]; + __half *op = out + (size_t)co * SCout + nb + lc + H; + const bool ok4 = (colok[lc] & colok[lc + 1] & + colok[lc + 2] & colok[lc + 3]) != 0; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = v0; + reinterpret_cast<__half2 *>(op)[1] = v1; + } else { + if (colok[lc + 0]) op[0] = v0.x; + if (colok[lc + 1]) op[1] = v0.y; + if (colok[lc + 2]) op[2] = v1.x; + if (colok[lc + 3]) op[3] = v1.y; + } + } + } + } else { + constexpr bool COMPACT_EPI = + (WG_BN == 136 || WG_BN == 160 || WG_BN == 208); + if constexpr (COMPACT_EPI) { + static_assert((BNT & 3) == 0, "4-column epilogue groups"); + constexpr int NC = BNT / 4; + constexpr int TOT = BM * NC; + constexpr int MITER = (TOT + NT - 1) / NT; + constexpr int UNR = (MITER >= 8) ? 8 : ((MITER >= 4) ? 4 : + ((MITER >= 2) ? 2 : 1)); +#pragma unroll 1 + for (int t0 = 0; t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int q = tid + (t0 + u) * NT; + if (q >= TOT) continue; + const int m = q / NC; + const int lc = (q - m * NC) * 4; + const int co = m0 + m; + const bool any4 = (colok[lc] | colok[lc + 1] | + colok[lc + 2] | colok[lc + 3]) != 0; + if (co >= Cout || !any4) continue; + v[u] = *reinterpret_cast(Cs + m * WG_CLD + lc); + if constexpr (!PREACT) { + const int n = nb + lc; + const __half2 *rp = reinterpret_cast( + resid + (size_t)co * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int q = tid + (t0 + u) * NT; + if (q >= TOT) continue; + const int m = q / NC; + const int lc = (q - m * NC) * 4; + const int co = m0 + m; + const bool any4 = (colok[lc] | colok[lc + 1] | + colok[lc + 2] | colok[lc + 3]) != 0; + if (co >= Cout || !any4) continue; + const int n = nb + lc; + float4 vv = v[u]; + if constexpr (!PREACT) { + vv.x += __half2float(r0[u].x); vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (co < half) { + const float a = row_alpha[m], ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + __half *op = out + (size_t)co * SCout + n + H; + const bool ok4 = (colok[lc] & colok[lc + 1] & + colok[lc + 2] & colok[lc + 3]) != 0; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + } else { + constexpr int NC = EPN / 4; + constexpr int MSTEP = NT / NC; + constexpr int MITER = BM / MSTEP; + constexpr int UNR = (MITER % 8 == 0) ? 8 + : ((MITER % 4 == 0) ? 4 + : ((MITER % 2 == 0) ? 2 : 1)); + static_assert(MITER % UNR == 0, "row batch must divide row count"); + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && + colok[lc + 2] && colok[lc + 3]; + const bool live4 = colok[lc] || colok[lc + 1] || + colok[lc + 2] || colok[lc + 3]; + const int n = nb + lc; +#pragma unroll 1 + for (int t0 = 0; live4 && t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + v[u] = *reinterpret_cast( + Cs + m * WG_CLD + lc); + if constexpr (!PREACT) { + const int co = m0 + m; + const int cc = co < Cout ? co : 0; + const __half2 *rp = reinterpret_cast( + resid + (size_t)cc * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + const int co = m0 + m; + if (co >= Cout) continue; + float4 vv = v[u]; + if constexpr (!PREACT) { + vv.x += __half2float(r0[u].x); + vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); + vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (co < half) { + const float a = row_alpha[m], ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + __half *op = out + (size_t)co * SCout + n + H; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B; + for (int idx = bxy * NT + tid; idx < total; idx += nbx * NT) { + const int b = idx / Cin, ci = idx - b * Cin; + const int x = b * L + native_true_tokens(b,T); + __half *const dst = ns + ((size_t)b * Cin + ci) * H; + if constexpr (BC8) { + const __half *const src = + in + (size_t)(ci >> 3) * ((size_t)SCin * 8) + + (size_t)x * 8 + (ci & 7); +#pragma unroll + for (int j = 0; j < H; ++j) dst[j] = src[j * 8]; + } else { + const __half *const src = in + (size_t)ci * SCin + x; +#pragma unroll + for (int j = 0; j < H; j += 2) + *reinterpret_cast(dst + j) = + *reinterpret_cast(src + j); + } + } + } + if (pad_src != nullptr) { + const int total = Cout * B; + for (int idx = bxy * NT + tid; idx < total; idx += nbx * NT) { + const int b = idx / Cout, co = idx - b * Cout; + const int x = b * L + pad_off; + const __half *const src = + pad_src + ((size_t)b * Cout + co) * pad_len; + if constexpr (OC8) { + __half *const dst = + out + (size_t)(co >> 3) * ((size_t)SCout * 8) + + (size_t)x * 8 + (co & 7); + for (int j = 0; j < pad_len; ++j) dst[j * 8] = src[j]; + } else { + __half *const dst = out + (size_t)co * SCout + x; + for (int j = 0; j < pad_len; ++j) dst[j] = src[j]; + } + } + } + } +} diff --git a/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wg5_step_widths.inc b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wg5_step_widths.inc new file mode 100644 index 0000000..5b362d9 --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wg5_step_widths.inc @@ -0,0 +1,400 @@ +template +__device__ __noinline__ +void wg5_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int KT = 3; + constexpr int BM = 64 * MSPLIT; + constexpr int BNT = WG_BN; + constexpr int NTS = 128 * MSPLIT; + constexpr int WG_NT = NTS * KSPLIT; + constexpr int NCRAW = BNT / 4; + constexpr int EPN = + (WG_NT % NCRAW == 0) ? BNT : (((BNT + 127) / 128) * 128); + constexpr int WG_CLD_R = (OC8 ? WG_BN : EPN) + 8; + constexpr int WG_CLD = (WG_CLD_R % 16 == 0) ? WG_CLD_R + 4 : WG_CLD_R; + constexpr int CG = CIB / 8; + constexpr int CSUB = CIB / 16; + constexpr int AK = KT * CSUB; + constexpr int NBRP = BC8 ? (((BNT + KT - 1) + 7) & ~7) + : (((BNT + KT - 1) + 31) & ~31); + constexpr int UCOL = NBRP / 32; + constexpr int NUNIT = CG * UCOL; + constexpr int NWARP = NTS / 32; + constexpr int BITER = (NUNIT + NWARP - 1) / NWARP; + constexpr int B_E = CIB * NBRP; + constexpr int ST_E = SB * B_E; + constexpr int LBO_B = NBRP * 8; + constexpr int ABUF = 3; + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *smem = reinterpret_cast<__half *>(arena); + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int kgrp = tid / NTS; + const int stid = tid - kgrp * NTS; + const int mg = stid >> 7; + const int warp = (stid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * BM; + const int nb = blockIdx.y * BNT; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + constexpr int NCH = (CIN + CIB - 1) / CIB; + static_assert(NCH % KSPLIT == 0, "wg5 K-split must divide the chunk count"); + if (tid < BM) { + const int co = m0 + tid; + if (co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + for (int i = tid; i < EPN; i += WG_NT) { + const int n = nb + i; + colok[i] = (i < BNT && n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mg * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (kgrp == 0 && eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (kgrp == 0 && eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + const int mtile = blockIdx.x * MSPLIT + mg; + const __half *const apt = w + ((size_t)mtile * NCH * AK) * 1024 + + (size_t)(stid & 127) * 8; + uint32_t areg[ABUF][AK][4]; + const int wi = stid >> 5; + int bdst[BITER], boff[BITER]; + bool blive[BITER]; + const __half *const bbase = in + (size_t)(lane >> 2) * SCin + nb + + 8 * (lane & 3); +#pragma unroll + for (int t = 0; t < BITER; ++t) { + const int u = wi + t * NWARP; + const int gc = (u < NUNIT) ? u / UCOL : 0; + const int cu = (u < NUNIT) ? u - gc * UCOL : 0; + blive[t] = u < NUNIT; + bdst[t] = gc * LBO_B + (cu * 32 + lane) * 8; + boff[t] = gc * 8 * SCin + cu * 32; + } + const size_t SC8i = (size_t)SCin * 8; + constexpr int BQ = B_E / 8; + constexpr int BCREG = (BQ + NTS - 1) / NTS; + int cdst[BCREG]; + const __half *csrc[BCREG]; + bool clive[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) { + const int q = stid + r * NTS; + const int gc = q / NBRP, rem = q - gc * NBRP; + clive[r] = ((BQ % NTS == 0) || q < BQ) && gc < CG; + cdst[r] = clive[r] ? gc * LBO_B + rem * 8 : 0; + csrc[r] = in + (clive[r] ? (size_t)gc * SC8i + + (size_t)nb * 8 + rem * 8 : 0); + } + uint32_t br[BITER][4]; + __half *const kbB = smem + kgrp * ST_E; + const unsigned barid = 1u + kgrp + (unsigned)(KSPLIT * native_subid()); + uint64_t dbuf[ABUF]; +#pragma unroll + for (int s = 0; s < ABUF; ++s) + dbuf[s] = wg_desc(kbB + s * B_E, LBO_B * 2u, 128u); + uint32_t csm0[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) + csm0[r] = static_cast( + __cvta_generic_to_shared(kbB + cdst[r])); + const size_t CSTEP = (size_t)KSPLIT * CG * SC8i; + constexpr int ASTEP = KSPLIT * AK * 1024; + const __half *crol[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) + crol[r] = csrc[r] + (size_t)(ABUF * KSPLIT + kgrp) * CG * SC8i; + const __half *arol = apt + (size_t)(ABUF * KSPLIT + kgrp) * (AK * 1024); + int crolc = ABUF * KSPLIT + kgrp; +#define WG5_BAR() asm volatile("bar.sync %0, %1;\n" ::"r"(barid), "r"(NTS)) +#define WG5_CPB(c,BD) { const size_t _o = (size_t)(c) * CG * SC8i; _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16((BD) + cdst[r], csrc[r] + _o, true); if ((c) / KSPLIT == ABUF - 1) { const size_t _r = (size_t)(ABUF * KSPLIT + kgrp) * CG * SC8i; _Pragma("unroll") for (int r = 0; r < BCREG; ++r) csrc[r] += _r; } } +#define WG5_LOADA(BUF,c) { (void)(c); } +#define WG5_KEEPA(BUF) { _Pragma("unroll") for (int kk = 0; kk < AK; ++kk) asm volatile("" : "+r"(areg[BUF][kk][0]), "+r"(areg[BUF][kk][1]), "+r"(areg[BUF][kk][2]), "+r"(areg[BUF][kk][3])); } +#define WG5_LOADB(c) { const size_t _o = (size_t)(c) * CIB * SCin; _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (NUNIT == NWARP * BITER || blive[t]) { const int4 _v = *reinterpret_cast( bbase + _o + boff[t]); br[t][0] = _v.x; br[t][1] = _v.y; br[t][2] = _v.z; br[t][3] = _v.w; } } } +#define WG5_BXPOSE(u) { const int _q = lane & 3; uint32_t _x, _y; _x = (_q & 1) ? u[0] : u[1]; _y = (_q & 1) ? u[2] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 1); _y = __shfl_xor_sync(0xffffffffu, _y, 1); if (_q & 1) { u[0] = _x; u[2] = _y; } else { u[1] = _x; u[3] = _y; } _x = (_q & 2) ? u[0] : u[2]; _y = (_q & 2) ? u[1] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 2); _y = __shfl_xor_sync(0xffffffffu, _y, 2); if (_q & 2) { u[0] = _x; u[1] = _y; } else { u[2] = _x; u[3] = _y; } } +#define WG5_STOREB(BD) { _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (NUNIT == NWARP * BITER || blive[t]) { WG5_BXPOSE(br[t]) const uint32_t _a = static_cast( __cvta_generic_to_shared((BD) + bdst[t])); asm volatile( "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 " "[%0], {%1,%2,%3,%4};\n" ::"r"(_a), "r"(br[t][0]), "r"(br[t][1]), "r"(br[t][2]), "r"(br[t][3])); } } } +#define WG5_MMA(BUF) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t db0 = dbuf[0] + (uint64_t)(BUF) * (B_E / 8); _Pragma("unroll") for (int h = 0; h < CSUB; ++h) { _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { WGMMA_RS_ANY_T(WG_BN, acc, areg[BUF][h * KT + kk], db0 + kk + h * (LBO_B / 4), 0); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } +#define WG5_CPB_R(SLOT) { _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16o<(SLOT) * B_E * 2>(csm0[r], csrc[r], true); } +#define WG5_LOADA_R(BUF) { } +#define WG5_ROLL() { crolc += KSPLIT; arol += ASTEP; if constexpr (BC8) { _Pragma("unroll") for (int r = 0; r < BCREG; ++r) csrc[r] += CSTEP; } } + if constexpr (PDL) { + (void)0; + if (tid == 0) + (void)0; + } + static_assert(SB == 3 && WGD == 1 && ABUF == 3, "wg5 pipeline shape"); + if constexpr (BC8) { +#pragma unroll 1 + for (int s = 0; s < SB; ++s) { + const int c = s * KSPLIT + kgrp; + if (c < NCH) WG5_CPB(c, kbB + s * B_E) + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } + } else { + WG5_LOADB(kgrp) +#pragma unroll 1 + for (int s = 0; s < SB; ++s) { + const int c = s * KSPLIT + kgrp; + if (c < NCH) WG5_STOREB(kbB + s * B_E) + if (c + KSPLIT < NCH) WG5_LOADB(c + KSPLIT) + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + } + WG5_LOADA(0, kgrp) + if (KSPLIT + kgrp < NCH) WG5_LOADA(1, KSPLIT + kgrp) + if (2 * KSPLIT + kgrp < NCH) WG5_LOADA(2, 2 * KSPLIT + kgrp) + constexpr int NIT = NCH / KSPLIT; + static_assert(NIT % ABUF == 0, "wg5 mainloop unroll shape"); +#define WG5_BODY(BUF,i) { { const __half *_at = apt + (size_t)((i) * KSPLIT + kgrp) * (AK * 1024); _Pragma("unroll") for (int kk = 0; kk < AK; ++kk) { const int4 _v = *reinterpret_cast(_at + kk * 1024); areg[BUF][kk][0] = _v.x; areg[BUF][kk][1] = _v.y; areg[BUF][kk][2] = _v.z; areg[BUF][kk][3] = _v.w; } } if constexpr (BC8) asm volatile("cp.async.wait_group %0;\n" : : "n"(SB - 2) : "memory"); WG5_BAR(); WG5_MMA(BUF) asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); WG5_KEEPA(BUF) if constexpr (!BC8) if ((i) > 1 && crolc < NCH) WG5_LOADB(crolc) if constexpr (NTS > 128) WG5_BAR(); if ((i) >= 1) { const int nb2 = ((BUF) + 2) % ABUF; if (crolc < NCH) { if constexpr (BC8) { WG5_CPB_R(nb2) } else { WG5_STOREB(kbB + nb2 * B_E) asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); } WG5_LOADA_R(nb2) } WG5_ROLL() if constexpr (BC8) asm volatile("cp.async.commit_group;\n" ::: "memory"); } } +#pragma unroll + for (int i = 0; i < NIT; i += 3) { + WG5_BODY(0, i) + WG5_BODY(1, i + 1) + WG5_BODY(2, i + 2) + } + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + if constexpr (BC8) + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); +#undef WG5_BXPOSE +#undef WG5_CPB +#undef WG5_CPB_R +#undef WG5_LOADA_R +#undef WG5_ROLL +#undef WG5_LOADB +#undef WG5_STOREB +#undef WG5_MMA +#undef WG5_LOADA +#undef WG5_KEEPA +#undef WG5_BODY +#undef WG5_BAR + __syncthreads(); + const int cb2 = (lane & 3) * 2; + if constexpr (OC8) { + __half *Ck = smem + kgrp * BM * WG_CLD; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + *reinterpret_cast<__half2 *>(Ck + (er0 + h * 8) * WG_CLD + + j * 8 + cb2) = + __floats2half2_rn(acc[j * 4 + h * 2], + acc[j * 4 + h * 2 + 1]); + } + } + } else { + float *Ck = Cs + kgrp * BM * WG_CLD; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], + acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Ck + (er0 + h * 8) * WG_CLD + + j * 8 + cb2) = v; + } + } + } + __syncthreads(); + if constexpr (OC8) { + constexpr int GRP = BM / 8; + constexpr int TOTG = GRP * WG_BN; + constexpr int GITER = (TOTG + WG_NT - 1) / WG_NT; + const size_t SC8o = (size_t)SCout * 8; +#pragma unroll + for (int t = 0; t < GITER; ++t) { + const int idx = tid + t * WG_NT; + if (TOTG % WG_NT != 0 && idx >= TOTG) continue; + const int j = idx - (idx / WG_BN) * WG_BN, g = idx / WG_BN; + if (!colok[j]) continue; + const int m0g = g * 8; + if (m0 + m0g >= Cout) continue; + const int n = nb + j; + __half rr[8]; + if constexpr (CIN == 1536) { + if constexpr (RC8) { + *reinterpret_cast(rr) = + *reinterpret_cast( + resid + (size_t)((m0 + m0g) >> 3) * SC8o + + (size_t)(n + H) * 8); + } else { +#pragma unroll + for (int h = 0; h < 8; ++h) + rr[h] = resid[(size_t)(m0 + m0g + h) * SCout + n + H]; + } + } + uint32_t pk[4]; + const bool snake = (m0 + m0g) < half; +#pragma unroll + for (int h = 0; h < 8; h += 2) { + float f[2]; +#pragma unroll + for (int e = 0; e < 2; ++e) { + const int m = m0g + h + e; + float a = __half2float(smem[m * WG_CLD + j]); +#pragma unroll + for (int kg = 1; kg < KSPLIT; ++kg) + a += __half2float( + smem[kg * BM * WG_CLD + m * WG_CLD + j]); + if constexpr (CIN == 1536) + a += __half2float(rr[h + e]); + if (snake) { + const float sn = __sinf(row_alpha[m] * a); + a = __fmaf_rn(sn * sn, row_ri[m], a); + } else { + a = a > 0.f ? a : 0.01f * a; + } + f[e] = a; + } + const __half2 p = __floats2half2_rn(f[0], f[1]); + pk[h >> 1] = *reinterpret_cast(&p); + } + *reinterpret_cast( + out + (size_t)((m0 + m0g) >> 3) * SC8o + + (size_t)(n + H) * 8) = make_int4(pk[0], pk[1], pk[2], pk[3]); + } + } else { + constexpr int NC = EPN / 4; + constexpr int MSTEP = WG_NT / NC; + constexpr int MITER = BM / MSTEP; + constexpr int UNR = (MITER % 4 == 0) ? 4 : ((MITER % 2 == 0) ? 2 : 1); + static_assert(MITER % UNR == 0, "row batch must divide the row count"); + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + const bool live4 = colok[lc] || colok[lc + 1] || colok[lc + 2] || + colok[lc + 3]; + const int n = nb + lc; +#pragma unroll 1 + for (int t0 = 0; live4 && t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + v[u] = *reinterpret_cast(Cs + m * WG_CLD + lc); + if constexpr (CIN == 1536) { + const int co = m0 + m; + const int cc = co < Cout ? co : 0; + if constexpr (RC8) { + const __half *rp = + resid + (size_t)(cc >> 3) * ((size_t)SCout * 8) + + (size_t)(n + H) * 8 + (cc & 7); + r0[u].x = rp[0]; r0[u].y = rp[8]; + r1[u].x = rp[16]; r1[u].y = rp[24]; + } else { + const __half2 *rp = reinterpret_cast( + resid + (size_t)cc * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + const int co = m0 + m; + if (co >= Cout) continue; + float4 vv = v[u]; +#pragma unroll + for (int g = 1; g < KSPLIT; ++g) { + float4 p = *reinterpret_cast( + Cs + g * BM * WG_CLD + m * WG_CLD + lc); + vv.x += p.x; vv.y += p.y; vv.z += p.z; vv.w += p.w; + } + if constexpr (CIN == 1536) { + vv.x += __half2float(r0[u].x); vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + __half *op = out + (size_t)co * SCout + n + H; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + static_assert(CIN == 768 || CIN == 1536); + constexpr int COUT = (CIN == 768) ? 1536 : 768; + if (ns != nullptr) { + const int total = CIN * B; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int b = idx / CIN, ci = idx - b * CIN; + const int x = b * L + native_true_tokens(b,T); + __half *const dst = ns + ((size_t)b * CIN + ci) * H; + if constexpr (BC8) { + const __half *const src = + in + (size_t)(ci >> 3) * ((size_t)SCin * 8) + + (size_t)x * 8 + (ci & 7); +#pragma unroll + for (int j = 0; j < H; ++j) dst[j] = src[j * 8]; + } else { + const __half *const src = in + (size_t)ci * SCin + x; +#pragma unroll + for (int j = 0; j < H; j += 2) + *reinterpret_cast(dst + j) = + *reinterpret_cast(src + j); + } + } + } + if (pad_src != nullptr) { + const int total = COUT * B; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int b = idx / COUT, co = idx - b * COUT; + const int x = b * L + pad_off; + const __half *const src = + pad_src + ((size_t)b * COUT + co) * pad_len; + if constexpr (OC8) { + __half *const dst = + out + (size_t)(co >> 3) * ((size_t)SCout * 8) + + (size_t)x * 8 + (co & 7); + for (int j = 0; j < pad_len; ++j) dst[j * 8] = src[j]; + } else { + __half *const dst = out + (size_t)co * SCout + x; + for (int j = 0; j < pad_len; ++j) dst[j] = src[j]; + } + } + } + } +} diff --git a/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wg_kernel.inc b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wg_kernel.inc new file mode 100644 index 0000000..2d5c766 --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wg_kernel.inc @@ -0,0 +1,182 @@ +__device__ __forceinline__ uint64_t wg_desc(const void *p, uint32_t lbo_bytes, + uint32_t sbo_bytes) { + uint32_t a = static_cast(__cvta_generic_to_shared(p)); + uint64_t d = static_cast((a & 0x3FFFFu) >> 4); + d |= static_cast((lbo_bytes & 0x3FFFFu) >> 4) << 16; + d |= static_cast((sbo_bytes & 0x3FFFFu) >> 4) << 32; + return d; +} +#include "wgmma_macro.inc" +constexpr int WG_SLOT = 72; +template +__device__ __noinline__ void wg_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int WG_CLD = WG_BN + 8; + constexpr int WG_NT = 128 * NWG; + constexpr int WG_BM = 64 * NWG; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int SBO_AE = KG * WG_SLOT; + constexpr int LBO_BE = (WG_BN / 8) * WG_SLOT; + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *As0 = reinterpret_cast<__half *>(arena); + __half *Bs0 = As0 + A_E; + __half *As1 = (STAGES == 2) ? Bs0 + B_E : As0; + __half *Bs1 = (STAGES == 2) ? Bs0 + B_E + A_E : Bs0; + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int wgid = tid >> 7; + const int warp = (tid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * WG_BM; + constexpr int WGROWS = 8; + const int n0 = blockIdx.y * WG_BN; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + if (tid < WG_BM) { + const int co = m0 + tid; + if (do_snake && co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + if (tid < WG_BN) { + const int n = n0 + tid; + colok[tid] = (n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = wgid * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + constexpr int ATOT = WG_BM * KG; + constexpr int AREG = (ATOT + WG_NT - 1) / WG_NT; + constexpr int CIB = BK / KT; + constexpr int BTOT = CIB * (WG_BN / 8); + constexpr int BREG = (BTOT + WG_NT - 1) / WG_NT; + int4 ar[AREG]; + int4 br0[BREG], br1[BREG]; + uint32_t edge[BREG]; +#define WG_LOADG(kg0) { const int _ci0 = (kg0) / KT; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; ar[r] = make_int4(0, 0, 0, 0); if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; const int co = m0 + m, kg = (kg0) + v * 8; if (co < Cout && kg < Kg) ar[r] = *reinterpret_cast( w + (size_t)co * Kg + kg); } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; br0[r] = make_int4(0, 0, 0, 0); if constexpr (KT == 3) edge[r] = 0; else br1[r] = make_int4(0, 0, 0, 0); if (BTOT % WG_NT == 0 || idx < BTOT) { const int ci = _ci0 + idx / (WG_BN / 8); const int nc = idx % (WG_BN / 8); if (ci < Cin) { const int4 *src = reinterpret_cast( in + (size_t)ci * SCin + n0 + nc * 8); br0[r] = src[0]; if constexpr (KT == 3) { if (nc == WG_BN / 8 - 1) edge[r] = src[1].x; } else { br1[r] = src[1]; } } } } } + WG_LOADG(0) +#define WG_STS(AD,BD) { _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; *reinterpret_cast( (AD) + (m >> 3) * SBO_AE + v * WG_SLOT + (m & 7) * 8) = ar[r]; } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; if (BTOT % WG_NT == 0 || idx < BTOT) { const int cil = idx / (WG_BN / 8), nc = idx % (WG_BN / 8); unsigned int u[8]; u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; if constexpr (KT == 3) { u[4] = __shfl_down_sync(0xffffffffu, u[0], 1, WG_BN / 8); if (nc == WG_BN / 8 - 1) u[4] = edge[r]; u[5] = u[6] = u[7] = 0; } else { u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; } _Pragma("unroll") for (int k = 0; k < KT; ++k) { const int kh = k >> 1; int4 o; if (k & 1) { o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); } else { o.x = u[kh + 0]; o.y = u[kh + 1]; o.z = u[kh + 2]; o.w = u[kh + 3]; } const int row = cil * KT + k; *reinterpret_cast((BD) + nc * WG_SLOT + (row >> 3) * LBO_BE + (row & 7) * 8) = o; } } } if constexpr (WG_BN == 64) { asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); } } +#define WG_MMA(AD,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t da0 = wg_desc((AD) + wgid * WGROWS * SBO_AE, WG_SLOT * 2, SBO_AE * 2); const uint64_t db0 = wg_desc((BD), LBO_BE * 2, WG_SLOT * 2); _Pragma("unroll") for (int kk = 0; kk < BK; kk += 16) { const uint64_t da = da0 + (kk >> 3) * (WG_SLOT / 8); const uint64_t db = db0 + (kk >> 3) * (LBO_BE / 8); if constexpr (WG_BN == 64) { WGMMA_M64N64K16(acc, da, db); } else { WGMMA_M64N128K16(acc, da, db); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + WG_STS(As0, Bs0) + if (BK < Kg) WG_LOADG(BK) + __syncthreads(); + int cur = 0; + for (int kg0 = 0; kg0 < Kg; kg0 += BK) { + if (cur == 0) { WG_MMA(As0, Bs0) } else { WG_MMA(As1, Bs1) } + if (kg0 + BK < Kg) { + if (STAGES == 2) + asm volatile("wgmma.wait_group.sync.aligned 1;\n" ::: "memory"); + else + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + __syncthreads(); + if (cur == 0) { WG_STS(As1, Bs1) } else { WG_STS(As0, Bs0) } + if (kg0 + 2 * BK < Kg) WG_LOADG(kg0 + 2 * BK) + __syncthreads(); + } else { + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + } + cur ^= 1; + } +#undef WG_LOADG +#undef WG_STS +#undef WG_MMA + __syncthreads(); + const int r0 = er0; + const int cb = (lane & 3) * 2; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Cs + (r0 + h * 8) * WG_CLD + j * 8 + cb) = v; + } + } + __syncthreads(); + constexpr int NC = WG_BN / 4; + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + for (int m = lr; m < WG_BM; m += WG_NT / NC) { + const int co = m0 + m; + if (co >= Cout) continue; + const int n = n0 + lc; + float4 v = *reinterpret_cast(Cs + m * WG_CLD + lc); + __half *op = out + (size_t)co * SCout + n + H; + if (resid != nullptr) { + const __half2 *rp = reinterpret_cast( + resid + (size_t)co * SCout + n + H); + __half2 r0h = rp[0], r1h = rp[1]; + v.x += __half2float(r0h.x); v.y += __half2float(r0h.y); + v.z += __half2float(r1h.x); v.w += __half2float(r1h.y); + } + if (do_snake) { + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + } + if (ok4) { + __half2 o0 = __floats2half2_rn(v.x, v.y); + __half2 o1 = __floats2half2_rn(v.z, v.w); + reinterpret_cast<__half2 *>(op)[0] = o0; + reinterpret_cast<__half2 *>(op)[1] = o1; + } else { + const float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B * H; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int total = Cout * B * pad_len; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} diff --git a/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wgmma_macro.inc b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wgmma_macro.inc new file mode 100644 index 0000000..db624e9 --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wgmma_macro.inc @@ -0,0 +1,102 @@ +#define WGMMA_M64N8K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " "{" "%0,%1,%2,%3}" ", %4, %5, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N8K16(D,DA,DB) WGMMA_M64N8K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N8K16_ZERO_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " "{" "%0,%1,%2,%3}" ", %4, %5, p, 1, 1, " #TA ", " #TB ";\n}\n" : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) : "l"(DA), "l"(DB), "r"(0)) +#define WGMMA_M64N16K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %10, 0;\n" "wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7}" ", %8, %9, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N16K16(D,DA,DB) WGMMA_M64N16K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N24K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %14, 0;\n" "wgmma.mma_async.sync.aligned.m64n24k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11}" ", %12, %13, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N24K16(D,DA,DB) WGMMA_M64N24K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N32K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %18, 0;\n" "wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15}" ", %16, %17, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N32K16(D,DA,DB) WGMMA_M64N32K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N40K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %22, 0;\n" "wgmma.mma_async.sync.aligned.m64n40k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19}" ", %20, %21, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N40K16(D,DA,DB) WGMMA_M64N40K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N48K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %26, 0;\n" "wgmma.mma_async.sync.aligned.m64n48k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N48K16(D,DA,DB) WGMMA_M64N48K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N56K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %30, 0;\n" "wgmma.mma_async.sync.aligned.m64n56k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N56K16(D,DA,DB) WGMMA_M64N56K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N64K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %34, 0;\n" "wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31}" ", %32, %33, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N64K16(D,DA,DB) WGMMA_M64N64K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N72K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %38, 0;\n" "wgmma.mma_async.sync.aligned.m64n72k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35}" ", %36, %37, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N72K16(D,DA,DB) WGMMA_M64N72K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N80K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %42, 0;\n" "wgmma.mma_async.sync.aligned.m64n80k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39}" ", %40, %41, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N80K16(D,DA,DB) WGMMA_M64N80K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N88K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %46, 0;\n" "wgmma.mma_async.sync.aligned.m64n88k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43}" ", %44, %45, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N88K16(D,DA,DB) WGMMA_M64N88K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N96K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %50, 0;\n" "wgmma.mma_async.sync.aligned.m64n96k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47}" ", %48, %49, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N96K16(D,DA,DB) WGMMA_M64N96K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N104K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %54, 0;\n" "wgmma.mma_async.sync.aligned.m64n104k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51}" ", %52, %53, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N104K16(D,DA,DB) WGMMA_M64N104K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N112K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %58, 0;\n" "wgmma.mma_async.sync.aligned.m64n112k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55}" ", %56, %57, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N112K16(D,DA,DB) WGMMA_M64N112K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N120K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %62, 0;\n" "wgmma.mma_async.sync.aligned.m64n120k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59}" ", %60, %61, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N120K16(D,DA,DB) WGMMA_M64N120K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N128K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %66, 0;\n" "wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63}" ", %64, %65, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N128K16(D,DA,DB) WGMMA_M64N128K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N136K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %70, 0;\n" "wgmma.mma_async.sync.aligned.m64n136k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67}" ", %68, %69, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N136K16(D,DA,DB) WGMMA_M64N136K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N144K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %74, 0;\n" "wgmma.mma_async.sync.aligned.m64n144k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71}" ", %72, %73, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N144K16(D,DA,DB) WGMMA_M64N144K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N152K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %78, 0;\n" "wgmma.mma_async.sync.aligned.m64n152k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75}" ", %76, %77, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N152K16(D,DA,DB) WGMMA_M64N152K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N160K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %82, 0;\n" "wgmma.mma_async.sync.aligned.m64n160k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79}" ", %80, %81, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N160K16(D,DA,DB) WGMMA_M64N160K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N168K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %86, 0;\n" "wgmma.mma_async.sync.aligned.m64n168k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83}" ", %84, %85, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N168K16(D,DA,DB) WGMMA_M64N168K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N176K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %90, 0;\n" "wgmma.mma_async.sync.aligned.m64n176k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87}" ", %88, %89, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N176K16(D,DA,DB) WGMMA_M64N176K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N184K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %94, 0;\n" "wgmma.mma_async.sync.aligned.m64n184k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91}" ", %92, %93, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N184K16(D,DA,DB) WGMMA_M64N184K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N192K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %98, 0;\n" "wgmma.mma_async.sync.aligned.m64n192k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95}" ", %96, %97, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N192K16(D,DA,DB) WGMMA_M64N192K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N200K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %102, 0;\n" "wgmma.mma_async.sync.aligned.m64n200k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99}" ", %100, %101, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N200K16(D,DA,DB) WGMMA_M64N200K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N208K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %106, 0;\n" "wgmma.mma_async.sync.aligned.m64n208k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103}" ", %104, %105, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N208K16(D,DA,DB) WGMMA_M64N208K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N216K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %110, 0;\n" "wgmma.mma_async.sync.aligned.m64n216k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107}" ", %108, %109, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N216K16(D,DA,DB) WGMMA_M64N216K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N224K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %114, 0;\n" "wgmma.mma_async.sync.aligned.m64n224k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111}" ", %112, %113, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N224K16(D,DA,DB) WGMMA_M64N224K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N232K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %118, 0;\n" "wgmma.mma_async.sync.aligned.m64n232k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115}" ", %116, %117, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N232K16(D,DA,DB) WGMMA_M64N232K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N240K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %122, 0;\n" "wgmma.mma_async.sync.aligned.m64n240k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119}" ", %120, %121, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N240K16(D,DA,DB) WGMMA_M64N240K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N248K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %126, 0;\n" "wgmma.mma_async.sync.aligned.m64n248k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123}" ", %124, %125, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N248K16(D,DA,DB) WGMMA_M64N248K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N256K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %130, 0;\n" "wgmma.mma_async.sync.aligned.m64n256k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123,%124,%125,%126,%127}" ", %128, %129, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]), "+f"(D[124]), "+f"(D[125]), "+f"(D[126]), "+f"(D[127]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N256K16(D,DA,DB) WGMMA_M64N256K16_T(D, DA, DB, 0, 1) +#define WGMMA_ANY_T(NW,D,DA,DB,TA,TB) do { if constexpr (NW == 8) { WGMMA_M64N8K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 16) { WGMMA_M64N16K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 24) { WGMMA_M64N24K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 32) { WGMMA_M64N32K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 40) { WGMMA_M64N40K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 48) { WGMMA_M64N48K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 56) { WGMMA_M64N56K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 64) { WGMMA_M64N64K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 72) { WGMMA_M64N72K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 80) { WGMMA_M64N80K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 88) { WGMMA_M64N88K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 96) { WGMMA_M64N96K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 104) { WGMMA_M64N104K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 112) { WGMMA_M64N112K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 120) { WGMMA_M64N120K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 128) { WGMMA_M64N128K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 136) { WGMMA_M64N136K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 144) { WGMMA_M64N144K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 152) { WGMMA_M64N152K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 160) { WGMMA_M64N160K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 168) { WGMMA_M64N168K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 176) { WGMMA_M64N176K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 184) { WGMMA_M64N184K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 192) { WGMMA_M64N192K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 200) { WGMMA_M64N200K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 208) { WGMMA_M64N208K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 216) { WGMMA_M64N216K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 224) { WGMMA_M64N224K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 232) { WGMMA_M64N232K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 240) { WGMMA_M64N240K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 248) { WGMMA_M64N248K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 256) { WGMMA_M64N256K16_T(D, DA, DB, TA, TB); } } while (0) +#define WGMMA_ANY(NW,D,DA,DB) WGMMA_ANY_T(NW, D, DA, DB, 0, 1) +#define WGMMA_RS_M64N96K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %53, 0;\n" "wgmma.mma_async.sync.aligned.m64n96k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47}" ", {%48,%49,%50,%51}, %52, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N104K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %57, 0;\n" "wgmma.mma_async.sync.aligned.m64n104k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51}" ", {%52,%53,%54,%55}, %56, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N112K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %61, 0;\n" "wgmma.mma_async.sync.aligned.m64n112k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55}" ", {%56,%57,%58,%59}, %60, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N120K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %65, 0;\n" "wgmma.mma_async.sync.aligned.m64n120k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59}" ", {%60,%61,%62,%63}, %64, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N128K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %69, 0;\n" "wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63}" ", {%64,%65,%66,%67}, %68, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_ANY_T(NW,D,A,DB,TB) do { if constexpr (NW == 96) { WGMMA_RS_M64N96K16_T(D, A, DB, TB); } else if constexpr (NW == 104) { WGMMA_RS_M64N104K16_T(D, A, DB, TB); } else if constexpr (NW == 112) { WGMMA_RS_M64N112K16_T(D, A, DB, TB); } else if constexpr (NW == 120) { WGMMA_RS_M64N120K16_T(D, A, DB, TB); } else if constexpr (NW == 128) { WGMMA_RS_M64N128K16_T(D, A, DB, TB); } } while (0) +#define WGMMA_RS_M64N8K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %9, 0;\n" "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " "{" "%0,%1,%2,%3}" ", {%4,%5,%6,%7}, %8, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N16K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %13, 0;\n" "wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7}" ", {%8,%9,%10,%11}, %12, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N24K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %17, 0;\n" "wgmma.mma_async.sync.aligned.m64n24k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11}" ", {%12,%13,%14,%15}, %16, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N32K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %21, 0;\n" "wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15}" ", {%16,%17,%18,%19}, %20, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N40K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %25, 0;\n" "wgmma.mma_async.sync.aligned.m64n40k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N48K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %29, 0;\n" "wgmma.mma_async.sync.aligned.m64n48k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N56K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %33, 0;\n" "wgmma.mma_async.sync.aligned.m64n56k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N64K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %37, 0;\n" "wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31}" ", {%32,%33,%34,%35}, %36, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N72K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %41, 0;\n" "wgmma.mma_async.sync.aligned.m64n72k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35}" ", {%36,%37,%38,%39}, %40, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N80K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %45, 0;\n" "wgmma.mma_async.sync.aligned.m64n80k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39}" ", {%40,%41,%42,%43}, %44, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N88K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %49, 0;\n" "wgmma.mma_async.sync.aligned.m64n88k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43}" ", {%44,%45,%46,%47}, %48, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N136K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %73, 0;\n" "wgmma.mma_async.sync.aligned.m64n136k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67}" ", {%68,%69,%70,%71}, %72, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N144K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %77, 0;\n" "wgmma.mma_async.sync.aligned.m64n144k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71}" ", {%72,%73,%74,%75}, %76, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N152K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %81, 0;\n" "wgmma.mma_async.sync.aligned.m64n152k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75}" ", {%76,%77,%78,%79}, %80, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N160K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %85, 0;\n" "wgmma.mma_async.sync.aligned.m64n160k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79}" ", {%80,%81,%82,%83}, %84, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N168K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %89, 0;\n" "wgmma.mma_async.sync.aligned.m64n168k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83}" ", {%84,%85,%86,%87}, %88, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N176K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %93, 0;\n" "wgmma.mma_async.sync.aligned.m64n176k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87}" ", {%88,%89,%90,%91}, %92, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N184K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %97, 0;\n" "wgmma.mma_async.sync.aligned.m64n184k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91}" ", {%92,%93,%94,%95}, %96, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N192K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %101, 0;\n" "wgmma.mma_async.sync.aligned.m64n192k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95}" ", {%96,%97,%98,%99}, %100, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N200K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %105, 0;\n" "wgmma.mma_async.sync.aligned.m64n200k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99}" ", {%100,%101,%102,%103}, %104, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N208K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %109, 0;\n" "wgmma.mma_async.sync.aligned.m64n208k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103}" ", {%104,%105,%106,%107}, %108, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N216K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %113, 0;\n" "wgmma.mma_async.sync.aligned.m64n216k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107}" ", {%108,%109,%110,%111}, %112, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N224K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %117, 0;\n" "wgmma.mma_async.sync.aligned.m64n224k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111}" ", {%112,%113,%114,%115}, %116, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N232K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %121, 0;\n" "wgmma.mma_async.sync.aligned.m64n232k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115}" ", {%116,%117,%118,%119}, %120, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N240K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %125, 0;\n" "wgmma.mma_async.sync.aligned.m64n240k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119}" ", {%120,%121,%122,%123}, %124, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N248K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %129, 0;\n" "wgmma.mma_async.sync.aligned.m64n248k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123}" ", {%124,%125,%126,%127}, %128, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N256K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %133, 0;\n" "wgmma.mma_async.sync.aligned.m64n256k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123,%124,%125,%126,%127}" ", {%128,%129,%130,%131}, %132, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]), "+f"(D[124]), "+f"(D[125]), "+f"(D[126]), "+f"(D[127]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#undef WGMMA_RS_ANY_T +#define WGMMA_RS_ANY_T(NW,D,A,DB,TB) do { if constexpr (NW == 8) { WGMMA_RS_M64N8K16_T(D, A, DB, TB); } else if constexpr (NW == 16) { WGMMA_RS_M64N16K16_T(D, A, DB, TB); } else if constexpr (NW == 24) { WGMMA_RS_M64N24K16_T(D, A, DB, TB); } else if constexpr (NW == 32) { WGMMA_RS_M64N32K16_T(D, A, DB, TB); } else if constexpr (NW == 40) { WGMMA_RS_M64N40K16_T(D, A, DB, TB); } else if constexpr (NW == 48) { WGMMA_RS_M64N48K16_T(D, A, DB, TB); } else if constexpr (NW == 56) { WGMMA_RS_M64N56K16_T(D, A, DB, TB); } else if constexpr (NW == 64) { WGMMA_RS_M64N64K16_T(D, A, DB, TB); } else if constexpr (NW == 72) { WGMMA_RS_M64N72K16_T(D, A, DB, TB); } else if constexpr (NW == 80) { WGMMA_RS_M64N80K16_T(D, A, DB, TB); } else if constexpr (NW == 88) { WGMMA_RS_M64N88K16_T(D, A, DB, TB); } else if constexpr (NW == 96) { WGMMA_RS_M64N96K16_T(D, A, DB, TB); } else if constexpr (NW == 104) { WGMMA_RS_M64N104K16_T(D, A, DB, TB); } else if constexpr (NW == 112) { WGMMA_RS_M64N112K16_T(D, A, DB, TB); } else if constexpr (NW == 120) { WGMMA_RS_M64N120K16_T(D, A, DB, TB); } else if constexpr (NW == 128) { WGMMA_RS_M64N128K16_T(D, A, DB, TB); } else if constexpr (NW == 136) { WGMMA_RS_M64N136K16_T(D, A, DB, TB); } else if constexpr (NW == 144) { WGMMA_RS_M64N144K16_T(D, A, DB, TB); } else if constexpr (NW == 152) { WGMMA_RS_M64N152K16_T(D, A, DB, TB); } else if constexpr (NW == 160) { WGMMA_RS_M64N160K16_T(D, A, DB, TB); } else if constexpr (NW == 168) { WGMMA_RS_M64N168K16_T(D, A, DB, TB); } else if constexpr (NW == 176) { WGMMA_RS_M64N176K16_T(D, A, DB, TB); } else if constexpr (NW == 184) { WGMMA_RS_M64N184K16_T(D, A, DB, TB); } else if constexpr (NW == 192) { WGMMA_RS_M64N192K16_T(D, A, DB, TB); } else if constexpr (NW == 200) { WGMMA_RS_M64N200K16_T(D, A, DB, TB); } else if constexpr (NW == 208) { WGMMA_RS_M64N208K16_T(D, A, DB, TB); } else if constexpr (NW == 216) { WGMMA_RS_M64N216K16_T(D, A, DB, TB); } else if constexpr (NW == 224) { WGMMA_RS_M64N224K16_T(D, A, DB, TB); } else if constexpr (NW == 232) { WGMMA_RS_M64N232K16_T(D, A, DB, TB); } else if constexpr (NW == 240) { WGMMA_RS_M64N240K16_T(D, A, DB, TB); } else if constexpr (NW == 248) { WGMMA_RS_M64N248K16_T(D, A, DB, TB); } else if constexpr (NW == 256) { WGMMA_RS_M64N256K16_T(D, A, DB, TB); } } while (0) diff --git a/deploy/nemotron_native/kernels/b4096/stage0.cu b/deploy/nemotron_native/kernels/b4096/stage0.cu new file mode 100644 index 0000000..6c04cd7 --- /dev/null +++ b/deploy/nemotron_native/kernels/b4096/stage0.cu @@ -0,0 +1 @@ +#include "regions/stage0/stage0.cu" diff --git a/deploy/nemotron_native/kernels/b512/codec.cu b/deploy/nemotron_native/kernels/b512/codec.cu new file mode 100644 index 0000000..7cf7265 --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/codec.cu @@ -0,0 +1,4 @@ +#define NATIVE_CODEC_PREFETCH 1 +#define NATIVE_CODEC_PACK_SUBCTAS 1 +#define D7W 96 +#include "native/codec/native_stage1.cu" diff --git a/deploy/nemotron_native/kernels/b512/native/codec/kernel_kf_codec.h b/deploy/nemotron_native/kernels/b512/native/codec/kernel_kf_codec.h new file mode 100644 index 0000000..9a8b4d5 --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/native/codec/kernel_kf_codec.h @@ -0,0 +1,13 @@ +#pragma once +#include +struct CodecArgs { + const void *inp[41]; + const void *st[26]; + void *ns[26]; + void *waveform; + void *scratch; + int B; + int F; +}; +void codec_launch(const CodecArgs &a, void *stream); +size_t codec_scratch_bytes(int B, int F); diff --git a/deploy/nemotron_native/kernels/b512/native/codec/legacy_codec.cuh b/deploy/nemotron_native/kernels/b512/native/codec/legacy_codec.cuh new file mode 100644 index 0000000..fd1b37c --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/native/codec/legacy_codec.cuh @@ -0,0 +1,1389 @@ +#include "kernel_kf_codec.h" +#include +#include +#include +#include +using namespace nvcuda; +template __device__ void prepare_attribute(F,int) {} +#ifndef TC2_QPT +#define TC2_QPT 4 +#endif +#define DIVUP(a,b) (((a) + (b) - 1) / (b)) +__device__ __forceinline__ float half_snake(float v, int c, int half, + const float *__restrict__ alpha) { + if (c < half) { + const float a = alpha[c]; + const float s = __sinf(a * v); + return __fmaf_rn(s, s / (a + 1e-9f), v); + } + return v > 0.f ? v : 0.01f * v; +} +#ifndef WG_ST3 +#define WG_ST3 1 +#define WG_MB3 3 +#define WG_ST7 1 +#define WG_MB7 3 +#define WG_BN3 128 +#define WG_BN7 128 +#endif +#define PDL_ENTER() do { (void)0; if (threadIdx.x == 0) (void)0; } while (0) +#define PDL_TRIGGER() do { if (threadIdx.x == 0) (void)0; } while (0) +template +__device__ __forceinline__ void pdl_launch(KFn kern, dim3 grid, dim3 blk, size_t shmem, + cudaStream_t s, KArgs... args) { + native_invoke(kern, grid, blk, shmem, args...); +} +template +__device__ __forceinline__ void pdl_launch_pf(KFn kern, dim3 grid, dim3 blk, size_t shmem, + cudaStream_t s, const void *pf, size_t pfb, KArgs... args) { + native_invoke_pf(kern, grid, blk, shmem, pf, pfb, args...); +} +#include "wg_kernel.inc" +#include "wg2_kernel.inc" +#include "wg3_kernel.inc" +#include "wg4_step_widths.inc" +#include "wg5_step_widths.inc" +#include "repack3.inc" +#include "tail_kernel.inc" +template +__device__ __noinline__ void dense_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + PDL_TRIGGER(); + constexpr int BK = 16 * KT; + constexpr int NT = NWARPS * 32; + constexpr int WTM = BM / WM; + constexpr int WTN = BN / WN; + constexpr int FM = WTM / 16, FN = WTN / 16; + constexpr int LDA = BK + 8, LDB = BN + 8; + constexpr int GEMM_SM = BM * LDA + BK * LDB; + constexpr int EPI_SM = NWARPS * WTM * WTN * 2; + constexpr int SM_H = GEMM_SM > EPI_SM ? GEMM_SM : EPI_SM; + __half *smem=reinterpret_cast<__half*>(native_arena()); + char *colok=reinterpret_cast(native_aux()); + __half *As = smem; + __half *Bs = smem + BM * LDA; + float *Cw = reinterpret_cast(smem); + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int wm = warp / WN, wn = warp % WN; + const int m0 = blockIdx.x * BM; + const int n0 = blockIdx.y * BN; + const int H = KT - 1; + const int Kg = Cin * KT; + if (tid < BN) { + const int n = n0 + tid; + colok[tid] = (n < N && (n % L) < T) ? 1 : 0; + } + wmma::fragment acc[FM][FN]; +#pragma unroll + for (int i = 0; i < FM; ++i) +#pragma unroll + for (int j = 0; j < FN; ++j) wmma::fill_fragment(acc[i][j], 0.f); + constexpr int ATOT = BM * (BK / 8); + constexpr int AREG = (ATOT + NT - 1) / NT; + constexpr int BTOT = 16 * (BN / 8); + constexpr int BREG = (BTOT + NT - 1) / NT; + int4 ar[AREG]; + int4 br0[BREG], br1[BREG]; +#define CODEC_LOADG(kg0) { const int _ci0 = (kg0) / KT; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * NT; ar[r] = make_int4(0, 0, 0, 0); if (ATOT % NT == 0 || idx < ATOT) { const int m = idx / (BK / 8), v = idx % (BK / 8); const int co = m0 + m, kg = (kg0) + v * 8; if (co < Cout && kg < Kg) ar[r] = *reinterpret_cast(w + (size_t)co * Kg + kg); } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * NT; br0[r] = make_int4(0, 0, 0, 0); br1[r] = make_int4(0, 0, 0, 0); if (BTOT % NT == 0 || idx < BTOT) { const int cil = idx / (BN / 8), nc = idx % (BN / 8); const int ci = _ci0 + cil; if (ci < Cin) { const int4 *src = reinterpret_cast( in + (size_t)ci * SCin + n0 + nc * 8); br0[r] = src[0]; br1[r] = src[1]; } } } } + CODEC_LOADG(0) + for (int kg0 = 0; kg0 < Kg; kg0 += BK) { + __syncthreads(); +#pragma unroll + for (int r = 0; r < AREG; ++r) { + const int idx = tid + r * NT; + if (ATOT % NT == 0 || idx < ATOT) + *reinterpret_cast(As + (idx / (BK / 8)) * LDA + + (idx % (BK / 8)) * 8) = ar[r]; + } +#pragma unroll + for (int r = 0; r < BREG; ++r) { + const int idx = tid + r * NT; + if (BTOT % NT != 0 && idx >= BTOT) continue; + const int cil = idx / (BN / 8), nc = idx % (BN / 8); + unsigned int u[8]; + u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; + u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; +#pragma unroll + for (int k = 0; k < KT; ++k) { + const int kh = k >> 1; + int4 o; + if (k & 1) { + o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); + o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); + o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); + o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); + } else { + o.x = u[kh + 0]; o.y = u[kh + 1]; + o.z = u[kh + 2]; o.w = u[kh + 3]; + } + *reinterpret_cast(Bs + (cil * KT + k) * LDB + nc * 8) = o; + } + } + __syncthreads(); + if (kg0 + BK < Kg) CODEC_LOADG(kg0 + BK) +#pragma unroll + for (int kf = 0; kf < BK; kf += 16) { + wmma::fragment af[FM]; + wmma::fragment bf[FN]; +#pragma unroll + for (int i = 0; i < FM; ++i) + wmma::load_matrix_sync(af[i], As + (wm * WTM + i * 16) * LDA + kf, LDA); +#pragma unroll + for (int j = 0; j < FN; ++j) + wmma::load_matrix_sync(bf[j], Bs + kf * LDB + wn * WTN + j * 16, LDB); +#pragma unroll + for (int i = 0; i < FM; ++i) +#pragma unroll + for (int j = 0; j < FN; ++j) + wmma::mma_sync(acc[i][j], af[i], bf[j], acc[i][j]); + } + } +#undef CODEC_LOADG + __syncthreads(); + float *Cme = Cw + warp * WTM * WTN; +#pragma unroll + for (int i = 0; i < FM; ++i) +#pragma unroll + for (int j = 0; j < FN; ++j) + wmma::store_matrix_sync(Cme + (i * 16) * WTN + j * 16, acc[i][j], WTN, + wmma::mem_row_major); + __syncwarp(); + const int half = Cout >> 1; + const int lane = tid & 31; + for (int e = lane; e < WTM * WTN; e += 32) { + const int mm = e / WTN, nn = e % WTN; + const int co = m0 + wm * WTM + mm; + const int nl = wn * WTN + nn; + if (co >= Cout || !colok[nl]) continue; + const int n = n0 + nl; + float v = Cme[e] + bias[co]; + if (resid != nullptr) + v += __half2float(resid[(size_t)co * SCout + n + H]); + if (do_snake) v = half_snake(v, co, half, alpha); + out[(size_t)co * SCout + n + H] = __float2half(v); + } + if (blockIdx.y == 0) { + const int nbx = gridDim.x; + if (ns != nullptr) { + const int B = N / L; + const int total = Cin * B * H; + for (int idx = blockIdx.x * NT + tid; idx < total; idx += nbx * NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int B = N / L; + const int total = Cout * B * pad_len; + for (int idx = blockIdx.x * NT + tid; idx < total; idx += nbx * NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} +template +__device__ __noinline__ void transpose_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + __half *__restrict__ out, const __half *__restrict__ pad_src, + __half *__restrict__ ns, int Cin, int Cout, int Tin, + int Lin, int SCin, int jbase, int Lout, int SCout, int Hout, int pad_off, + int pad_len, int B) { + if constexpr (PDL) PDL_ENTER(); + constexpr int Ktaps = 2 * STRIDE; + constexpr int QPT = QPT_; + constexpr int VPT = QPT * STRIDE; + static_assert(QPT % 2 == 0, "half2 input loads need an even frame count"); + static_assert(VPT % 2 == 0, "half2 staging needs an even slot"); + float *ws=reinterpret_cast(native_aux()); + __half *sm=reinterpret_cast<__half*>(native_arena()); + const int co = blockIdx.y; + const int b = blockIdx.z; + const int tid = threadIdx.x, BLK = blockDim.x; + if (G != 4 || STRIDE != 7 || native_meta()->lengths[b] > 0) { + for (int i = tid; i < G * Ktaps; i += BLK) + ws[i] = __half2float(w[(size_t)(co * G + i / Ktaps) * Ktaps + (i % Ktaps)]); + __syncthreads(); + const float bs = bias[0 + co]; + const int half = Cout >> 1; + const bool snake = co < half; + const float av = snake ? alpha[co] : 0.f; + const float ri = snake ? __frcp_rn(av + 1e-9f) : 0.f; + const __half *ic = in + (size_t)(co * G) * SCin + b * Lin + jbase; + __half *oc = out + (size_t)co * SCout + b * Lout + Hout; + constexpr bool VEC16 = (VPT % 8) == 0; + const int ophase = + VEC16 ? 0 : (int)((((uintptr_t)oc) >> 1) & 7u); + __half *const sbase = sm + ophase; + const int hd = (8 - ophase) & 7; + const int lane = tid & 31; + for (int qb = 1 + blockIdx.x * BLK * QPT; qb <= Tin; + qb += gridDim.x * BLK * QPT) { + const int q = qb + QPT * tid; + const bool a0 = (q <= Tin); + float jq[QPT][G], jp0[G]; + #pragma unroll + for (int g = 0; g < G; ++g) { + #pragma unroll + for (int e2 = 0; e2 < QPT / 2; ++e2) { + const __half2 v = a0 ? *reinterpret_cast( + ic + (size_t)g * SCin + q + 2 * e2) + : __floats2half2_rn(0.f, 0.f); + jq[2 * e2][g] = __half2float(v.x); + jq[2 * e2 + 1][g] = __half2float(v.y); + } + } + #pragma unroll + for (int g = 0; g < G; ++g) { + const float up = __shfl_up_sync(0xffffffffu, jq[QPT - 1][g], 1); + jp0[g] = (lane == 0) + ? (a0 ? __half2float(ic[(size_t)g * SCin + q - 1]) : 0.f) + : up; + } + __half *const sp = sbase + tid * VPT; + __align__(16) __half vreg[VPT]; + #pragma unroll + for (int r = 0; r < STRIDE; ++r) { + float wc[G], wp[G]; + #pragma unroll + for (int g = 0; g < G; ++g) { + wc[g] = ws[g * Ktaps + r]; + wp[g] = ws[g * Ktaps + r + STRIDE]; + } + #pragma unroll + for (int e = 0; e < QPT; ++e) { + float v = bs; + #pragma unroll + for (int g = 0; g < G; ++g) { + const float cur = jq[e][g]; + const float prv = (e == 0) ? jp0[g] : jq[e - 1][g]; + v = __fmaf_rn(wc[g], cur, __fmaf_rn(wp[g], prv, v)); + } + if (snake) { + const float sn = __sinf(av * v); + v = __fmaf_rn(sn * sn, ri, v); + } else { + v = v > 0.f ? v : 0.01f * v; + } + vreg[e * STRIDE + r] = __float2half(v); + } + } + if constexpr (VEC16) { + #pragma unroll + for (int m = 0; m < VPT / 8; ++m) + *reinterpret_cast(sp + m * 8) = + *reinterpret_cast(vreg + m * 8); + } else { + #pragma unroll + for (int i = 0; i < VPT; ++i) sp[i] = vreg[i]; + } + __syncthreads(); + const int nval = min(BLK * QPT, Tin - qb + 1); + const int tot = nval * STRIDE; + __half *dst = oc + (size_t)(qb - 1) * STRIDE; + if constexpr (VEC16) { + const int tot2 = tot >> 1; + for (int i = tid; i < tot2; i += BLK) + reinterpret_cast<__half2 *>(dst)[i] = + reinterpret_cast(sbase)[i]; + if ((tot & 1) && tid == 0) dst[tot - 1] = sbase[tot - 1]; + } else { + const int head = hd < tot ? hd : tot; + for (int i = tid; i < head; i += BLK) dst[i] = sbase[i]; + const int nb4 = (tot - head) >> 3; + const int4 *s4 = reinterpret_cast(sbase + head); + int4 *d4 = reinterpret_cast(dst + head); + for (int i = tid; i < nb4; i += BLK) d4[i] = s4[i]; + for (int i = head + (nb4 << 3) + tid; i < tot; i += BLK) + dst[i] = sbase[i]; + } + __syncthreads(); + } + } + if (blockIdx.x == 0) { + if (ns != nullptr) { + for (int ci = co + tid * Cout; ci < Cin; ci += BLK * Cout) + ns[(size_t)b * Cin + ci] = + in[(size_t)ci * SCin + b * Lin + jbase + native_true_tokens(b,Tin)]; + } + if (pad_src != nullptr) { + for (int j = tid; j < pad_len; j += BLK) + out[(size_t)co * SCout + b * Lout + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } +} +template +__device__ __noinline__ void transpose_conv_early_bpack_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + __half *__restrict__ out, const __half *__restrict__ pad_src, + __half *__restrict__ ns, int Cin, int Cout, int Tin, + int Lin, int SCin, int jbase, int Lout, int SCout, int Hout, int pad_off, + int pad_len, int B) { + if constexpr (PDL) PDL_ENTER(); + static_assert(BPACK <= 32); + static_assert(LPB == 8 || LPB == 16 || LPB == 32); + constexpr int Ktaps = 2 * STRIDE; + constexpr int WARP = 32; + constexpr int NT = WARP * BPACK; + constexpr int BPB = WARP / LPB; + float *ws=reinterpret_cast(native_aux()); + __half *sm=reinterpret_cast<__half*>(native_arena()); + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int lane = tid & 31; + const int sub = lane % LPB; + const int bsub = lane / LPB; + const int co = blockIdx.y; + const int b = blockIdx.z * (BPACK * BPB) + warp * BPB + bsub; + const bool active = b < B; + for (int i = tid; i < G * Ktaps; i += NT) + ws[i] = __half2float(w[(size_t)(co * G + i / Ktaps) * Ktaps + + (i % Ktaps)]); + __syncthreads(); + const int qb = 1 + blockIdx.x * LPB; + const int q = qb + sub; + __half *sp = sm + (warp * WARP + lane) * STRIDE; + const bool live = active && q <= Tin; + { + const float bs = bias[co]; + const int half = Cout >> 1; + const bool snake = co < half; + const float av = snake ? alpha[co] : 0.f; + const float ri = snake ? __frcp_rn(av + 1e-9f) : 0.f; + const __half *ic = in + (size_t)(co * G) * SCin + b * Lin + jbase; + float jq[G], jp[G]; +#pragma unroll + for (int g = 0; g < G; ++g) + jq[g] = live ? __half2float(ic[(size_t)g * SCin + q]) : 0.f; +#pragma unroll + for (int g = 0; g < G; ++g) { + const float up = __shfl_up_sync(0xffffffffu, jq[g], 1, LPB); + jp[g] = (sub == 0) + ? (live ? __half2float(ic[(size_t)g * SCin + q - 1]) : 0.f) + : up; + } + if (live) +#pragma unroll + for (int r = 0; r < STRIDE; ++r) { + float v = bs; +#pragma unroll + for (int g = 0; g < G; ++g) + v = __fmaf_rn(ws[g * Ktaps + r], jq[g], + __fmaf_rn(ws[g * Ktaps + r + STRIDE], jp[g], v)); + if (snake) { + const float sn = __sinf(av * v); + v = __fmaf_rn(sn * sn, ri, v); + } else { + v = v > 0.f ? v : 0.01f * v; + } + sp[r] = __float2half(v); + } + } + __syncwarp(); + if (active) { + const int nval = min(LPB, Tin - (qb - 1)); + const int tot2 = (nval * STRIDE) >> 1; + const __half *src = sm + (warp * WARP + bsub * LPB) * STRIDE; + __half *dst = out + (size_t)co * SCout + b * Lout + Hout + + (size_t)(qb - 1) * STRIDE; + for (int i = sub; i < tot2; i += LPB) + reinterpret_cast<__half2 *>(dst)[i] = + reinterpret_cast(src)[i]; + if ((nval * STRIDE) & 1) + if (sub == 0) dst[nval * STRIDE - 1] = src[nval * STRIDE - 1]; + } + if (blockIdx.x == 0 && active) { + if (ns != nullptr) { + for (int ci = co + sub * Cout; ci < Cin; ci += LPB * Cout) + ns[(size_t)b * Cin + ci] = + in[(size_t)ci * SCin + b * Lin + jbase + native_true_tokens(b,Tin)]; + } + if (pad_src != nullptr) { + for (int j = sub; j < pad_len; j += LPB) + out[(size_t)co * SCout + b * Lout + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } +} +template +__device__ __noinline__ void transpose_conv_bpack2_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + __half *__restrict__ out, const __half *__restrict__ pad_src, + __half *__restrict__ ns, int Cin, int Cout, int Tin, + int Lin, int SCin, int jbase, int Lout, int SCout, int Hout, int pad_off, + int pad_len, int B) { + if constexpr (PDL) PDL_ENTER(); + static_assert(BPACK <= 32); + static_assert(LPB == 8 || LPB == 16 || LPB == 32); + constexpr int Ktaps = 2 * STRIDE; + constexpr int WARP = 32; + constexpr int NT = WARP * BPACK; + constexpr int BPB = WARP / LPB; + constexpr int QPT = 2; + constexpr int VPT = QPT * STRIDE; + float *ws=reinterpret_cast(native_aux()); + __half *sm=reinterpret_cast<__half*>(native_arena()); + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int lane = tid & 31; + const int sub = lane % LPB; + const int bsub = lane / LPB; + const int co = blockIdx.y; + const int b = blockIdx.z * (BPACK * BPB) + warp * BPB + bsub; + const bool active = b < B; + for (int i = tid; i < G * Ktaps; i += NT) + ws[i] = __half2float(w[(size_t)(co * G + i / Ktaps) * Ktaps + + (i % Ktaps)]); + __syncthreads(); + const int qb = 1 + blockIdx.x * (LPB * QPT); + const int q = qb + sub * QPT; + __half *sp = sm + (warp * WARP + lane) * VPT; + const bool a0 = active && (q <= Tin); + { + const float bs = bias[co]; + const int half = Cout >> 1; + const bool snake = co < half; + const float av = snake ? alpha[co] : 0.f; + const float ri = snake ? __frcp_rn(av + 1e-9f) : 0.f; + const __half *ic = in + (size_t)(co * G) * SCin + b * Lin + jbase; + float jq0[G], jq1[G], jp0[G]; +#pragma unroll + for (int g = 0; g < G; ++g) { + const __half2 v = a0 ? *reinterpret_cast( + ic + (size_t)g * SCin + q) + : __floats2half2_rn(0.f, 0.f); + jq0[g] = __half2float(v.x); + jq1[g] = __half2float(v.y); + } +#pragma unroll + for (int g = 0; g < G; ++g) { + const float up = __shfl_up_sync(0xffffffffu, jq1[g], 1, LPB); + jp0[g] = (sub == 0) + ? (a0 ? __half2float(ic[(size_t)g * SCin + q - 1]) + : 0.f) + : up; + } + if (a0) { + float vv[VPT]; +#pragma unroll + for (int r = 0; r < STRIDE; ++r) { + float wc[G], wp[G]; +#pragma unroll + for (int g = 0; g < G; ++g) { + wc[g] = ws[g * Ktaps + r]; + wp[g] = ws[g * Ktaps + r + STRIDE]; + } +#pragma unroll + for (int e = 0; e < QPT; ++e) { + float v = bs; +#pragma unroll + for (int g = 0; g < G; ++g) { + const float cur = (e == 0) ? jq0[g] : jq1[g]; + const float prv = (e == 0) ? jp0[g] : jq0[g]; + v = __fmaf_rn(wc[g], cur, __fmaf_rn(wp[g], prv, v)); + } + if (snake) { + const float sn = __sinf(av * v); + v = __fmaf_rn(sn * sn, ri, v); + } else { + v = v > 0.f ? v : 0.01f * v; + } + vv[e * STRIDE + r] = v; + } + } + __half2 *sp2 = reinterpret_cast<__half2 *>(sp); +#pragma unroll + for (int i = 0; i < VPT / 2; ++i) + sp2[i] = __floats2half2_rn(vv[2 * i], vv[2 * i + 1]); + } + } + __syncwarp(); + if (active) { + const int nval = min(LPB * QPT, Tin - (qb - 1)); + const int tot = nval * STRIDE; + const __half *src = sm + (warp * WARP + bsub * LPB) * VPT; + __half *dst = out + (size_t)co * SCout + b * Lout + Hout + + (size_t)(qb - 1) * STRIDE; + for (int i = sub; i < (tot >> 1); i += LPB) + reinterpret_cast<__half2 *>(dst)[i] = + reinterpret_cast(src)[i]; + if (tot & 1) + if (sub == 0) dst[tot - 1] = src[tot - 1]; + } + if (blockIdx.x == 0 && active) { + if (ns != nullptr) { + for (int ci = co + sub * Cout; ci < Cin; ci += LPB * Cout) + ns[(size_t)b * Cin + ci] = + in[(size_t)ci * SCin + b * Lin + jbase + native_true_tokens(b,Tin)]; + } + if (pad_src != nullptr) { + for (int j = sub; j < pad_len; j += LPB) + out[(size_t)co * SCout + b * Lout + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } +} +#define TCONV_BPACK(G,STRIDE,BPACK,LPB,TIN,COUT,BB,...) pdl_launch(transpose_conv_early_bpack_kernel, dim3(DIVUP(TIN, LPB), COUT, DIVUP(BB, BPACK * (32 / LPB))), dim3(32 * BPACK), (size_t)(32 * (BPACK)) * (STRIDE) * 2 + 64, s, __VA_ARGS__) +#define TCONV_BP2(G,STRIDE,LPB,TIN,COUT,BB,...) pdl_launch(transpose_conv_bpack2_kernel, dim3(DIVUP(TIN, (LPB) * 2), COUT, DIVUP(BB, 32)), dim3(32 * (LPB)), (size_t)(32 * (LPB)) * (4 * (STRIDE)) + 64, s, __VA_ARGS__) +#define TCONV_BPACK_DISPATCH1(G,STRIDE,BPACK,TIN,COUT,BB,...) do { if ((TIN) >= 32) TCONV_BP2(G, STRIDE, 16, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) >= 16) TCONV_BP2(G, STRIDE, 8, TIN, COUT, BB, __VA_ARGS__); else TCONV_BPACK(G, STRIDE, BPACK, 16, TIN, COUT, BB, __VA_ARGS__); } while (0) +#define TCONV_BPACK_DISPATCH(G,STRIDE,BPACK,TIN,COUT,BB,...) do { if ((TIN) >= 32) TCONV_BP2(G, STRIDE, 16, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) >= 16) TCONV_BP2(G, STRIDE, 8, TIN, COUT, BB, __VA_ARGS__); else TCONV_BPACK(G, STRIDE, BPACK, 16, TIN, COUT, BB, __VA_ARGS__); } while (0) +#define TCONV_BPACK_DISPATCH8(G,STRIDE,BPACK,TIN,COUT,BB,...) do { if ((TIN) >= 32) TCONV_BP2(G, STRIDE, 16, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) >= 16) TCONV_BP2(G, STRIDE, 8, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) <= 8) TCONV_BPACK(G, STRIDE, BPACK, 8, TIN, COUT, BB, __VA_ARGS__); else TCONV_BPACK(G, STRIDE, BPACK, 16, TIN, COUT, BB, __VA_ARGS__); } while (0) +__device__ __noinline__ void pack_latent_kernel(const __half *__restrict__ latent, + const __half *__restrict__ state, + __half *__restrict__ out, int Cin, int T, + int L, int SC, int H, int B) { + const int total = Cin * B * L; + for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; + idx += gridDim.x * blockDim.x) { + const int j = idx % L, r = idx / L; + const int b = r % B, ci = r / B; + out[(size_t)ci * SC + b * L + j] = + (j < H) ? state[((size_t)b * Cin + ci) * H + j] + : latent[((size_t)b * Cin + ci) * T + (j - H)]; + } +} +#define PC_VEC 8 +__device__ __noinline__ void post_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, __half *__restrict__ y, + __half *__restrict__ ns, int Cin, int T, int L, int SC, int KT, int B, + int N) { + float *pw=reinterpret_cast(native_arena()); + for (int i = threadIdx.x; i < Cin * KT; i += 256) pw[i] = __half2float(w[i]); + __syncthreads(); + const int H = KT - 1; + const float bs = bias[0]; + for (int n = (blockIdx.x * 256 + threadIdx.x) * PC_VEC; n < N; + n += gridDim.x * 256 * PC_VEC) { + float v[PC_VEC]; +#pragma unroll + for (int i = 0; i < PC_VEC; ++i) v[i] = bs; + for (int ci = 0; ci < Cin; ++ci) { + const __half *xr = in + (size_t)ci * SC + n; + const float *wr = pw + ci * KT; + const int4 lo = *reinterpret_cast(xr); + const int4 hi = *reinterpret_cast(xr + 8); + unsigned u[8]; + u[0] = lo.x; u[1] = lo.y; u[2] = lo.z; u[3] = lo.w; + u[4] = hi.x; u[5] = hi.y; u[6] = hi.z; u[7] = hi.w; + float x[PC_VEC + 6]; +#pragma unroll + for (int j = 0; j < (PC_VEC + 6) / 2; ++j) { + const float2 f = __half22float2( + *reinterpret_cast(&u[j])); + x[2 * j] = f.x; + x[2 * j + 1] = f.y; + } +#pragma unroll + for (int i = 0; i < PC_VEC; ++i) + for (int k = 0; k < 7; ++k) + if (k < KT) v[i] = __fmaf_rn(wr[k], x[i + k], v[i]); + } + const int b0 = n / L, t0 = n - (n / L) * L; +#pragma unroll + for (int i = 0; i < PC_VEC; ++i) { + int t = t0 + i, b = b0; + if (t >= L) { t -= L; b = b0 + 1; } + if (t < T && b < B) + y[(size_t)b * T + t] = + __float2half(fminf(fmaxf(v[i], -1.f), 1.f)); + } + } + if (ns != nullptr) { + const int total = B * Cin * H; + for (int idx = blockIdx.x * 256 + threadIdx.x; idx < total; + idx += gridDim.x * 256) { + const int j = idx % H, r = idx / H; + const int ci = r % Cin, b = r / Cin; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SC + b * L + native_true_tokens(b,T) + j]; + } + } +} +#define RP_JOBS 6 +struct RepackSet { + const __half *src[RP_JOBS]; + __half *dst[RP_JOBS]; + int Cout[RP_JOBS], CoutP[RP_JOBS], Cin[RP_JOBS]; + int base[RP_JOBS], nch[RP_JOBS], rs[RP_JOBS]; + int k3base; + const __half *k3src[2]; + __half *k3dst[2]; +}; +__host__ __device__ inline int native_RP_COUT(int j) {const int v[6]={768, 384, 256, 128, 64, 32};return v[j];} +__host__ __device__ inline int native_RP_CIN(int j) {const int v[6]={384, 768, 128, 256, 32, 64};return v[j];} +__host__ __device__ inline int native_RP_COUTP(int j) {const int v[6]={768, 384, 256, 128, 128, 128};return v[j];} +template +__device__ __noinline__ void repack_kernel(RepackSet js) { + const int bx = blockIdx.x; + int job, job_bx; + if (bx < 1152) { + job = 0; job_bx = 0; + } else if (bx < 2304) { + job = 1; job_bx = 1152; + } else if (bx < 2432) { + job = 2; job_bx = 2304; + } else if (bx < 2560) { + job = 3; job_bx = 2432; + } else if (bx < 2576) { + job = 4; job_bx = 2560; + } else { + job = 5; job_bx = 2576; + } + const int Cout = js.Cout[job], CoutP = js.CoutP[job], Cin = js.Cin[job]; + const int idx = (bx - job_bx) * blockDim.x + threadIdx.x; + const int co = idx / Cin, ci = idx - co * Cin; + const int c = ci / CIB, cl = ci - c * CIB; + __half *d = js.dst[job] + (size_t)c * (KT * CoutP * CIB) + + (size_t)(co >> 3) * (CIB * 8) + (cl >> 3) * 64 + (co & 7) * 8 + + (cl & 7); + const __half *sp = js.src[job] + (size_t)(co * Cin + ci) * KT; + const bool live = co < Cout; +#pragma unroll + for (int k = 0; k < KT; ++k) + d[(size_t)k * CoutP * CIB] = live ? sp[k] : __float2half(0.f); +} +template +__device__ __noinline__ void repack_rs_kernel( + const RepackSet &js, const __half *__restrict__ latent, const __half *__restrict__ latent_state, + __half *__restrict__ latent_out, int latent_C, int latent_T, int latent_L, + int latent_SC, int latent_H, int latent_B) { + PDL_TRIGGER(); + constexpr int CO_TILE = 16; + constexpr int CI_TILE = 64; + constexpr int ROW_H = CI_TILE * KT + 8; + constexpr int ROW4 = CI_TILE * KT / 8; + __half *tile=reinterpret_cast<__half*>(native_arena()); + const int bx = blockIdx.x; + { + const int idx = bx + threadIdx.x * gridDim.x; + const int total = latent_C * latent_B * latent_L; + if (idx < total) { + const int j = idx % latent_L, r = idx / latent_L; + const int b = r % latent_B, ci = r / latent_B; + latent_out[(size_t)ci * latent_SC + b * latent_L + j] = + (j < latent_H) + ? latent_state[((size_t)b * latent_C + ci) * latent_H + j] + : latent[((size_t)b * latent_C + ci) * latent_T + + (j - latent_H)]; + } + } + if (bx >= js.k3base) { + constexpr int K3 = 3; + constexpr int R3_H = CI_TILE * K3 + 8; + constexpr int R3_4 = CI_TILE * K3 / 8; + const int r = bx - js.k3base; + const int conv = r / 1152; + const int loc = r - conv * 1152; + const int Cout3 = conv ? 768 : 1536; + const int Cin3 = conv ? 1536 : 768; + const int nci3 = Cin3 / CI_TILE; + const int cot3 = loc / nci3; + if (cot3 * CO_TILE >= Cout3) return; + const int co03 = cot3 * CO_TILE; + const int ci03 = (loc - cot3 * nci3) * CI_TILE; + const int NCH3 = Cin3 >> 4; + const __half *const s3 = js.k3src[conv]; + __half *const d3 = js.k3dst[conv]; + const int t3 = threadIdx.x; +#pragma unroll + for (int n = 0; n < CO_TILE * R3_4 / 128; ++n) { + const int i = t3 + n * 128; + const int co = i / R3_4, q = i - co * R3_4; + reinterpret_cast(tile + co * R3_H)[q] = + *reinterpret_cast( + s3 + ((size_t)(co03 + co) * Cin3 + ci03) * K3 + q * 8); + } + __syncthreads(); + const int w3 = t3 >> 5, l3 = t3 & 31; + const int rA3 = l3 >> 2, q43 = l3 & 3; + const int cb3 = co03 >> 4; + __half *const db3 = d3 + ((size_t)(cb3 >> 2) * NCH3 * K3) * 1024 + + ((cb3 & 3) * 32 + l3) * 8; +#pragma unroll + for (int u = w3; u < (CI_TILE / 16) * K3; u += 4) { + const int cl = u / K3, tap = u - cl * K3; + const int cgi = (ci03 >> 4) + cl; + const __half *sa = tile + rA3 * R3_H + (cl * 16 + 2 * q43) * K3 + tap; + const __half *sb = sa + 8 * R3_H; + uint32_t o[4]; + o[0] = __half_as_ushort(sa[0]) | + ((uint32_t)__half_as_ushort(sa[K3]) << 16); + o[1] = __half_as_ushort(sb[0]) | + ((uint32_t)__half_as_ushort(sb[K3]) << 16); + o[2] = __half_as_ushort(sa[8 * K3]) | + ((uint32_t)__half_as_ushort(sa[9 * K3]) << 16); + o[3] = __half_as_ushort(sb[8 * K3]) | + ((uint32_t)__half_as_ushort(sb[9 * K3]) << 16); + *reinterpret_cast(db3 + (size_t)(cgi * K3 + tap) * 1024) = + make_int4(o[0], o[1], o[2], o[3]); + } + return; + } + int job = 0; +#pragma unroll + for (int q = 1; q < RP_JOBS; ++q) if (bx >= js.base[q]) job = q; + const int local = bx - js.base[job]; + const int Cout = js.Cout[job], Cin = js.Cin[job]; + const int nci = (Cin + CI_TILE - 1) / CI_TILE; + const int cot = local / nci; + const int co0 = cot * CO_TILE; + const int ci0 = (local - cot * nci) * CI_TILE; + const int NCH = js.nch[job]; + const int t = threadIdx.x; + const int valid4 = (min(Cin - ci0, CI_TILE) * KT) >> 3; + for (int i = t; i < CO_TILE * ROW4; i += 128) { + const int co = i / ROW4, q = i - co * ROW4; + int4 v = make_int4(0, 0, 0, 0); + if (co0 + co < Cout && q < valid4) + v = *reinterpret_cast( + js.src[job] + ((size_t)(co0 + co) * Cin + ci0) * KT + q * 8); + reinterpret_cast(tile + co * ROW_H)[q] = v; + } + __syncthreads(); + const int warp = t >> 5, lane = t & 31; + const int rA = lane >> 2, q4 = lane & 3; + const int cb = co0 >> 4; + const int mtile = cb >> 2, wsub = cb & 3; + __half *const dbase = js.dst[job] + + ((size_t)mtile * NCH * KT) * 1024 + (wsub * 32 + lane) * 8; + if (!js.rs[job]) { + const int CoutP = js.CoutP[job]; + constexpr int UN2 = 16 * (CI_TILE / 8) * KT; +#pragma unroll + for (int n = 0; n < UN2 / 128; ++n) { + const int u = t + n * 128; + const int col = u & 15, cig = (u >> 4) & (CI_TILE / 8 - 1); + const int tap = u / (16 * (CI_TILE / 8)); + const int gci = ci0 + cig * 8; + if (gci >= Cin) continue; + const int gco = co0 + col; + uint32_t o[4]; + const __half *sp = tile + col * ROW_H; +#pragma unroll + for (int e = 0; e < 4; ++e) + o[e] = __half_as_ushort(sp[(cig * 8 + 2 * e) * KT + tap]) | + ((uint32_t)__half_as_ushort( + sp[(cig * 8 + 2 * e + 1) * KT + tap]) << 16); + __half *d = js.dst[job] + + (size_t)(gci >> 4) * ((size_t)KT * CoutP * 16) + + (size_t)tap * (CoutP * 16) + (size_t)(gco >> 3) * 128 + + ((gci >> 3) & 1) * 64 + (gco & 7) * 8; + *reinterpret_cast(d) = make_int4(o[0], o[1], o[2], o[3]); + } + return; + } + constexpr int UNITS = (CI_TILE / 16) * KT; +#pragma unroll + for (int u = warp; u < UNITS; u += 4) { + const int cl = u / KT, tap = u - cl * KT; + const int cgi = (ci0 >> 4) + cl; + if (cgi >= NCH) continue; + const __half *sa = tile + rA * ROW_H + (cl * 16 + 2 * q4) * KT + tap; + const __half *sb = sa + 8 * ROW_H; + uint32_t o[4]; + o[0] = __half_as_ushort(sa[0]) | + ((uint32_t)__half_as_ushort(sa[KT]) << 16); + o[1] = __half_as_ushort(sb[0]) | + ((uint32_t)__half_as_ushort(sb[KT]) << 16); + o[2] = __half_as_ushort(sa[8 * KT]) | + ((uint32_t)__half_as_ushort(sa[9 * KT]) << 16); + o[3] = __half_as_ushort(sb[8 * KT]) | + ((uint32_t)__half_as_ushort(sb[9 * KT]) << 16); + *reinterpret_cast(dbase + (size_t)(cgi * KT + tap) * 1024) = + make_int4(o[0], o[1], o[2], o[3]); + } +} +__device__ __noinline__ void repack3_rs_kernel( + const __half *__restrict__ win, const __half *__restrict__ wsk, + __half *__restrict__ din, __half *__restrict__ dsk, + const __half *__restrict__ asrc, __half *__restrict__ adst, int SCA, + int ncol, int ntile) { + PDL_TRIGGER(); + if (blockIdx.y == 12) { + if (adst == nullptr) return; + for (int u = blockIdx.x; u < 96 * ntile; u += gridDim.x) { + const int g = u / ntile; + const int col = (u - g * ntile) * 128 + (int)threadIdx.x; + if (col >= ncol) continue; + __half v[8]; +#pragma unroll + for (int h = 0; h < 8; ++h) + v[h] = asrc[(size_t)(g * 8 + h) * SCA + col]; + *reinterpret_cast(adst + (size_t)g * ((size_t)SCA * 8) + + (size_t)col * 8) = + *reinterpret_cast(v); + } + return; + } + constexpr int KT = 3; + constexpr int CO_TILE = 16; + constexpr int CI_TILE = 64; + constexpr int ROW_H = CI_TILE * KT + 8; + constexpr int ROW4 = CI_TILE * KT / 8; + __half *tile=reinterpret_cast<__half*>(native_arena()); + const int jy = blockIdx.y; + const int conv = jy & 1, b = 5 - (jy >> 1); + const int Cout = conv ? 768 : 1536; + const int Cin = conv ? 1536 : 768; + const __half *src = (conv ? wsk : win) + (size_t)(b + 1) * Cout * Cin * KT; + __half *dst = (conv ? dsk : din) + (size_t)b * ((size_t)Cout * Cin * KT); + const int nci = Cin / CI_TILE; + const int cot = blockIdx.x / nci; + const int co0 = cot * CO_TILE; + const int ci0 = (blockIdx.x - cot * nci) * CI_TILE; + const int NCH = Cin >> 4; + const int t = threadIdx.x; +#pragma unroll + for (int n = 0; n < CO_TILE * ROW4 / 128; ++n) { + const int i = t + n * 128; + const int co = i / ROW4, q = i - co * ROW4; + reinterpret_cast(tile + co * ROW_H)[q] = + *reinterpret_cast( + src + ((size_t)(co0 + co) * Cin + ci0) * KT + q * 8); + } + __syncthreads(); + const int warp = t >> 5, lane = t & 31; + const int rA = lane >> 2, q4 = lane & 3; + const int cb = co0 >> 4; + __half *const dbase = dst + ((size_t)(cb >> 2) * NCH * KT) * 1024 + + ((cb & 3) * 32 + lane) * 8; +#pragma unroll + for (int u = warp; u < (CI_TILE / 16) * KT; u += 4) { + const int cl = u / KT, tap = u - cl * KT; + const int cgi = (ci0 >> 4) + cl; + const __half *sa = tile + rA * ROW_H + (cl * 16 + 2 * q4) * KT + tap; + const __half *sb = sa + 8 * ROW_H; + uint32_t o[4]; + o[0] = __half_as_ushort(sa[0]) | + ((uint32_t)__half_as_ushort(sa[KT]) << 16); + o[1] = __half_as_ushort(sb[0]) | + ((uint32_t)__half_as_ushort(sb[KT]) << 16); + o[2] = __half_as_ushort(sa[8 * KT]) | + ((uint32_t)__half_as_ushort(sa[9 * KT]) << 16); + o[3] = __half_as_ushort(sb[8 * KT]) | + ((uint32_t)__half_as_ushort(sb[9 * KT]) << 16); + *reinterpret_cast(dbase + (size_t)(cgi * KT + tap) * 1024) = + make_int4(o[0], o[1], o[2], o[3]); + } +} +__device__ __noinline__ void repack_activation_only(const __half *asrc,__half *adst,int SCA,int ncol,int ntile){ + if (adst == nullptr) return; + for (int u = blockIdx.x; u < 96 * ntile; u += gridDim.x) { + const int g = u / ntile; + const int col = (u - g * ntile) * 128 + (int)threadIdx.x; + if (col >= ncol) continue; + __half v[8]; +#pragma unroll + for (int h = 0; h < 8; ++h) + v[h] = asrc[(size_t)(g * 8 + h) * SCA + col]; + *reinterpret_cast(adst + (size_t)g * ((size_t)SCA * 8) + + (size_t)col * 8) = + *reinterpret_cast(v); + } + return; +} +namespace { +struct Stack { + int T, L, N, SC, H; +}; +inline __host__ __device__ int tc_blk(int Tin, int qpt = 2) { + int best = 256, bw = 1 << 30; + for (int b = 256; b >= 32; b >>= 1) { + const int per = qpt * b; + const int waste = (Tin + per - 1) / per * per - Tin; + if (waste < bw) { bw = waste; best = b; } + } + return best; +} +inline __host__ __device__ int roundup(int a, int b) { return (a + b - 1) / b * b; } +__host__ __device__ Stack mk(int B, int T, int K) { + Stack s; + s.T = T; + s.H = K - 1; + s.L = s.H + T; + s.N = B * s.L; + s.SC = roundup(s.N, 128) + 128; + return s; +} +struct Plan { + Stack s0, s1, s2, s3, s4; + __half *LB, *X0, *M0, *Y0; + __half *A, *Ac, *Bb, *M1; + __half *P, *Q, *M2; + __half *R, *S, *M3; + __half *U; + __half *W7[RP_JOBS]; + __half *W3i, *W3s; + size_t total; +}; +static constexpr int kTuneB = 32; +inline __host__ __device__ int nref(const Stack &st) { return kTuneB * st.L; } +inline __host__ __device__ bool use_rp3i(int, int F) { return (size_t)kTuneB * (4 * F + 2) > 400; } +inline __host__ __device__ bool use_rp3s(int, int F) { return (size_t)kTuneB * (4 * F + 2) > 400; } +__host__ __device__ Plan build_plan(int B, int F, __half *base) { + Plan p; + p.s0 = mk(B, 2 * F, 3); + p.s1 = mk(B, 4 * F, 3); + p.s2 = mk(B, 36 * F, 7); + p.s3 = mk(B, 252 * F, 7); + p.s4 = mk(B, 1764 * F, 7); + size_t off = 0; + auto alloc = [&](int ch, const Stack &s) -> __half * { + __half *r = base ? base + off : nullptr; + off += (size_t)ch * s.SC; + return r; + }; + p.LB = alloc(40, p.s0); + p.X0 = alloc(768, p.s0); + p.M0 = alloc(1536, p.s0); + p.Y0 = alloc(768, p.s0); + p.A = alloc(768, p.s1); + p.Ac = use_rp3i(B, F) ? alloc(768, p.s1) : nullptr; + p.Bb = alloc(768, p.s1); + p.M1 = alloc(1536, p.s1); + p.P = alloc(384, p.s2); + p.Q = alloc(384, p.s2); + p.M2 = alloc(768, p.s2); + p.R = alloc(128, p.s3); + p.S = alloc(128, p.s3); + p.M3 = alloc(256, p.s3); + p.U = alloc(32, p.s4); + for (int j = 0; j < RP_JOBS; ++j) { + p.W7[j] = base ? base + off : nullptr; + off += (size_t)native_RP_COUTP(j) * native_RP_CIN(j) * 7; + } + p.W3i = nullptr; + if (use_rp3i(B, F)) { + p.W3i = base ? base + off + (size_t)3 * 1536 * 768 : nullptr; + off += (size_t)7 * 3 * 1536 * 768; + } + p.W3s = nullptr; + if (use_rp3s(B, F)) { + p.W3s = base ? base + off + (size_t)3 * 768 * 1536 : nullptr; + off += (size_t)7 * 3 * 768 * 1536; + } + p.total = off; + return p; +} +__host__ __device__ const __half *H16(const void *p) { return static_cast(p); } +__host__ __device__ const float *F32(const void *p) { return static_cast(p); } +} +size_t codec_scratch_bytes(int B, int F) { + return build_plan(B, F, nullptr).total * sizeof(__half); +} +template +__device__ __forceinline__ void wg_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, bool prep = false) { + constexpr int WG_BM = 64 * NWG, WG_NT = 128 * NWG, WG_CLD = WG_BN + 8; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int GEMM_B = STAGES * (A_E + B_E) * 2; + constexpr int EPI_B = WG_BM * WG_CLD * 4; + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute(wg_conv_kernel, SH); + return; + } + native_invoke(wg_conv_kernel,dim3(dim3(DIVUP(Cout, WG_BM), DIVUP(st.N, WG_BN))),dim3(WG_NT),SH, + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg2_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, + bool prep = false) { + constexpr int BK = CIB * KT; + constexpr int WG_BM = 64 * MSPLIT; + constexpr int WG_NT = 128 * MSPLIT * KSPLIT * NSPLIT; + constexpr int WG_BNT = WG_BN * NSPLIT; + constexpr int NCRAW = WG_BNT / 4; + constexpr int EPN = + (WG_NT % NCRAW == 0) ? WG_BNT : (((WG_BNT + 127) / 128) * 128); + constexpr int WG_CLD = EPN + 8; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int GEMM_B = KSPLIT * (SA * A_E + SB * NSPLIT * B_E) * 2; + constexpr int EPI_B = KSPLIT * WG_BM * WG_CLD * 4; + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute( + wg2_conv_kernel, + SH); + return; + } + native_invoke_pf(wg2_conv_kernel,dim3(dim3(DIVUP(Cout, WG_BM), DIVUP(st.N, WG_BNT))),dim3(WG_NT),SH, + w, (size_t)KT * Cin * Cout * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg3_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, + bool prep = false) { + constexpr int WG_BM = 64 * NWG, WG_NT = 128 * NWG * NSG; + constexpr int WG_BNT = WG_BN * NSG, WG_CLD = WG_BNT + 8; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int GEMM_B = STAGES * (A_E + NSG * B_E) * 2; + constexpr int EPI_B = WG_BM * WG_CLD * 4; + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute(wg3_conv_kernel, SH); + return; + } + native_invoke(wg3_conv_kernel,dim3(dim3(DIVUP(Cout, WG_BM), DIVUP(st.N, WG_BNT))),dim3(WG_NT),SH, + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg4_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, int CoutP, + const Stack &st, int SCin, int do_snake, int pad_off, + int pad_len, bool prep = false) { + constexpr int NT = 128 * MSPLIT * NSPLIT; + constexpr int BM = 64 * MSPLIT, BNT = WG_BN * NSPLIT; + constexpr int NBR = BNT + KT - 1; + constexpr int NBRP = BC8 ? ((NBR + 7) & ~7) : ((NBR + 31) & ~31); + constexpr int NCRAW = BNT / 4; + constexpr int EPN = + (NT % NCRAW == 0) ? BNT : (((BNT + 127) / 128) * 128); + constexpr int B_E = CIB * NBRP; + constexpr int GEMM_B = STAGES * B_E * 2; + constexpr int EPI_B = BM * (EPN + 8) * (PREACT ? 2 : 4); + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute( + wg4_conv_kernel, + SH); + return; + } + pdl_launch_pf(wg4_conv_kernel, + dim3(DIVUP(Cout, BM), DIVUP(st.N, BNT)), dim3(NT), SH, s, + w, (size_t)KT * CoutP * Cin * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, + st.T, st.L, SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg5_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, + bool prep = false) { + constexpr int BM = 64 * MSPLIT; + constexpr int NT = 128 * MSPLIT * KSPLIT; + constexpr int NCRAW = WG_BN / 4; + constexpr int EPN = (NT % NCRAW == 0) ? WG_BN : (((WG_BN + 127) / 128) * 128); + constexpr int NBRP = BC8 ? (((WG_BN + 2) + 7) & ~7) + : (((WG_BN + 2) + 31) & ~31); + constexpr int B_E = CIB * NBRP; + constexpr int GEMM_B = KSPLIT * SB * B_E * 2; + constexpr int EPI_COLS = OC8 ? WG_BN : EPN; + constexpr int EPI_CLD_R = EPI_COLS + 8; + constexpr int EPI_CLD = (EPI_CLD_R % 16 == 0) ? EPI_CLD_R + 4 : EPI_CLD_R; + constexpr int EPI_B = KSPLIT * BM * EPI_CLD * (OC8 ? 2 : 4); + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + static_assert(SH <= 227 * 1024, "wg5 shared over budget"); + if (prep) { + prepare_attribute( + wg5_conv_kernel, + SH); + return; + } + if constexpr (PDL) { + cudaLaunchConfig_t cfg = {}; + cudaLaunchAttribute attr[1] = {}; + native_config_grid(cfg) = dim3(DIVUP(Cout, BM), DIVUP(st.N, WG_BN)); + native_config_threads(cfg) = dim3(NT); + cfg.dynamicSmemBytes = SH; + cfg.stream = s; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = 1; + cfg.attrs = attr; + cfg.numAttrs = 1; + native_cudaLaunchKernelEx_pf( + &cfg, + wg5_conv_kernel, + w, (size_t)3 * Cin * Cout * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); + } else { + native_invoke_pf(wg5_conv_kernel,dim3(dim3(DIVUP(Cout, BM), DIVUP(st.N, WG_BN))),dim3(NT),SH, + w, (size_t)3 * Cin * Cout * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, + st.L, SCin, st.SC, st.N, do_snake, pad_off, pad_len); + } +} +#ifndef D7W +#define D7W 96 +#endif +#define D7_M2_C8I(BN,LATE) wg4_launch<7, 16, 2, 1, D7W, 4, 1, false, 2, true, false, LATE>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_M2P_C8O(BN,LATE) wg4_launch<7, 16, 2, 1, D7W, 4, 1, true, 2, false, true, LATE>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_M2N2(BN) wg4_launch<7, 16, 2, 2, BN, 4, 1, false, 2>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_M1(BN) wg4_launch<7, 16, 1, 1, BN, 4, 2, false, 2>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_C8I(W) \ + wg4_launch<7, 16, 2, 1, W, 4, 1, false, 2, true, false, false>( \ + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, \ + SCin, 1, pad_off, pad_len, prep) +#define D7_C8O(W) \ + wg4_launch<7, 16, 2, 1, W, 4, 1, true, 2, false, true, false>( \ + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, \ + SCin, 1, pad_off, pad_len, prep) +template +__device__ __forceinline__ int d7_band_cost(int N, int mtiles, int resident) { + const int v = mtiles * DIVUP(N, BN); + const int sub = (DIVUP(v, resident) >= 2) ? 2 : 1; + return DIVUP(DIVUP(v, sub), resident) * sub * (128 + BN + 6); +} +__device__ __forceinline__ int d7_band(int N, int mtiles, int resident) { + int best = 64, bc = d7_band_cost<64>(N, mtiles, resident); + const int c96 = d7_band_cost<96>(N, mtiles, resident); + if (c96 < bc) { bc = c96; best = 96; } + if (d7_band_cost<112>(N, mtiles, resident) < bc) best = 112; + return best; +} +__device__ __forceinline__ void dense7(cudaStream_t s, int job, const __half *in, const __half *w, + const float *bias, const float *alpha, const __half *resid, + __half *out, const __half *pad_src, __half *ns, + const Stack &st, int SCin, int pad_off, int pad_len, + bool prep = false) { + const int Cin = native_RP_CIN(job), Cout = native_RP_COUT(job), CoutP = native_RP_COUTP(job); + const int b_ = d7_band(st.N, DIVUP(Cout, 128), native_resident()); + if ((job & 1) == 0) { + if (b_ == 64) D7_C8O(64); + else if (b_ == 96) D7_C8O(96); + else D7_C8O(112); + } else { + if (b_ == 64) D7_C8I(64); + else if (b_ == 96) D7_C8I(96); + else D7_C8I(112); + } +} +__device__ __forceinline__ int wg5_band_skip(int N, int mtiles, int resident) { + const int wmin = DIVUP(mtiles * DIVUP(N, 104), resident); + if (DIVUP(mtiles * DIVUP(N, 24), resident) == wmin) return 24; + if (DIVUP(mtiles * DIVUP(N, 40), resident) == wmin) return 40; + if (DIVUP(mtiles * DIVUP(N, 48), resident) == wmin) return 48; + if (DIVUP(mtiles * DIVUP(N, 64), resident) == wmin) return 64; + return 104; +} +__device__ __forceinline__ int wg5_band_main_long(int N, int mtiles, + int resident) { + const int wmin = DIVUP(mtiles * DIVUP(N, 104), resident); + if (DIVUP(mtiles * DIVUP(N, 64), resident) == wmin) return 64; + return 104; +} +__device__ __forceinline__ int wg5_band_main_short(int N, int mtiles, + int resident) { + const int wmin = DIVUP(mtiles * DIVUP(N, 96), resident); + if (DIVUP(mtiles * DIVUP(N, 56), resident) == wmin) return 56; + if (DIVUP(mtiles * DIVUP(N, 80), resident) == wmin) return 80; + return 96; +} +#define DENSE_LAUNCH(KT,BM,BN,NW,WM,WN) native_invoke(dense_conv_kernel,dim3(dim3(DIVUP(Cout, BM), DIVUP(st.N, BN))),dim3(NW * 32),0, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, SCin, st.SC, st.N, do_snake, pad_off, pad_len) +__device__ __forceinline__ void dense(cudaStream_t s, int KT, const __half *in, const __half *w, + const float *bias, const float *alpha, const __half *resid, + __half *out, const __half *pad_src, __half *ns, int Cin, + int Cout, const Stack &st, int SCin, int do_snake, + int pad_off, int pad_len, const __half *w3 = nullptr, + int layout_mode = 0, bool prep = false) { + if (KT == 3 && w3 != nullptr) { + if (Cin == 1536) { +#define WG5S(BN,O,R,PDL) wg5_launch<16, 1, 4, BN, 4, 3, 1, 1, true, O, R, 1536, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5SW(BN,O,R,PDL) wg5_launch<32, 1, 2, BN, 4, 3, 1, 1, true, O, R, 1536, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5S3(O,R) do { const int n_ = nref(st); if (n_ > 500 && n_ <= 640) { WG5SW(56, O, R, false); } else { const int b_ = wg5_band_skip(st.N, DIVUP(Cout, 64), native_resident()); if (b_ == 24) WG5S(24, O, R, true); else if (b_ == 40) WG5S(40, O, R, true); else if (b_ == 48) WG5S(48, O, R, true); else if (b_ == 64) WG5S(64, O, R, true); else WG5S(104, O, R, true); } } while (0) + if (layout_mode == 1) WG5S3(true, false); + else if (layout_mode == 2) WG5S3(true, true ); + else if (layout_mode == 3) WG5S3(false, true ); + else WG5S3(false, false); +#undef WG5S3 +#undef WG5SW +#undef WG5S + } else { +#define WG5I(BN,BC,MS,PDL) wg5_launch<16, MS, 2, BN, 4, 3, 1, 1, BC, true, false, 768, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5I4(BN,BC,PDL) wg5_launch<16, 1, 4, BN, 4, 3, 1, 1, BC, true, false, 768, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5I2W(BN,BC,PDL) wg5_launch<32, 1, 2, BN, 4, 3, 1, 1, BC, true, false, 768, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5I3(BC) do { const int n_ = nref(st); if (n_ > 640) { const int b_ = wg5_band_main_long(st.N, DIVUP(Cout, 128), native_resident()); if (b_ == 64) WG5I(64, BC, 2, true); else WG5I(104, BC, 2, true); } else if (n_ > 500) { WG5I2W(64, BC, false); } else { const int b_ = wg5_band_main_short(st.N, DIVUP(Cout, 64), native_resident()); if (b_ == 56) WG5I4(56, BC, true); else if (b_ == 80) WG5I4(80, BC, true); else WG5I4(96, BC, true); } } while (0) + if (layout_mode == 1) WG5I3(true); + else WG5I3(false); +#undef WG5I3 +#undef WG5I2W +#undef WG5I4 +#undef WG5I + } + } else if (KT == 3) { + const int n_ = nref(st); + if (Cin == 1536 && Cout == 768) { + if (n_ <= 704) + wg2_launch<3, 16, 1, 4, 1, 64, 2, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else + wg2_launch<3, 16, 1, 4, 1, 128, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + } else if (Cin == 768 && Cout == 1536) { + if (n_ > 1000) + wg2_launch<3, 16, 1, 2, 1, 128, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ > 640 && n_ <= 768) + wg2_launch<3, 16, 1, 2, 1, 64, 2, 2, 2>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ <= 320) + wg2_launch<3, 16, 1, 4, 1, 64, 2, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ > 384 && n_ <= 512) + wg2_launch<3, 16, 1, 4, 1, 112, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ > 512 && n_ <= 640) + wg2_launch<3, 16, 1, 4, 1, 120, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else + wg2_launch<3, 16, 1, 4, 1, 128, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + } else { + wg2_launch<3, 16, 1, 1, 1, 64, 3, 2, 2>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, + SCin, do_snake, pad_off, pad_len, prep); + } + } else if (Cout >= 128 && + ((Cin >= 384 && st.T == 288) || + (Cin < 384 && st.T == 2016))) { + wg3_launch<7, 112, 2, 2, 1, 128, 2>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, + SCin, do_snake, pad_off, pad_len); + } else { + wg3_launch<7, 112, 1, 4, 1, 128, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, + SCin, do_snake, pad_off, pad_len); + } +} +__device__ void codec_program(const CodecArgs &a, const Plan &p, const RepackSet &js, + int repack_grid, + const __half *const *state_override = nullptr) { + cudaStream_t s = nullptr; + const int B = a.B; + const void *const *W = a.inp; + const __half *const *ST = state_override ? state_override : + reinterpret_cast(a.st); + __half *const *NS = reinterpret_cast<__half *const *>(a.ns); + dense(s, 3, p.LB, H16(W[1]), F32(W[2]), nullptr, nullptr, p.X0, ST[1], NS[0], + 40, 768, p.s0, p.s0.SC, 0, 0, 2); + const __half *k3iw = H16(W[3]); + const float *k3ib = F32(W[4]); + const float *k3ia = F32(W[5]); + const __half *k3sw = H16(W[6]); + const float *k3sb = F32(W[7]); + const float *k3sa = F32(W[8]); + const size_t IW = (size_t)1536 * 768 * 3, SW = (size_t)768 * 1536 * 3; + dense(s, 3, p.X0, k3iw, k3ib, k3ia, nullptr, p.M0, ST[2], NS[1], 768, 1536, + p.s0, p.s0.SC, 1, 0, 2, p.W3i ? p.W3i - IW : nullptr, 0); + dense(s, 3, p.M0, k3sw, k3sb, k3sa, p.X0, p.Y0, ST[3], NS[2], 1536, 768, + p.s0, p.s0.SC, 1, 1, 1, p.W3s ? p.W3s - SW : nullptr, 0); + TCONV_BPACK_DISPATCH8(1, 2, 8, p.s0.T, 768, B, + p.Y0, H16(W[9]), F32(W[10]), F32(W[11]), p.A, ST[4], NS[3], 768, 768, + p.s0.T, p.s0.L, p.s0.SC, 1, p.s1.L, p.s1.SC, 2, 0, 2, B); + if (p.W3i != nullptr) + native_invoke(repack_activation_only,dim3(min(1152,96*DIVUP(p.s1.SC,128))),dim3(128),0, + p.A, p.Ac, p.s1.SC, p.s1.SC, DIVUP(p.s1.SC,128)); + __half *cur = p.A, *nxt = p.Bb; + for (int i = 0; i < 6; ++i) { + const int bi = i + 1; + const __half *ain = (i == 0 && p.Ac) ? p.Ac : cur; + dense(s, 3, ain, k3iw + bi * IW, k3ib + bi * 1536, k3ia + bi * 768, + nullptr, p.M1, ST[5 + 2 * i], NS[4 + 2 * i], 768, 1536, p.s1, + p.s1.SC, 1, 0, 2, p.W3i ? p.W3i + (bi - 1) * IW : nullptr, + (i > 0 || p.Ac) ? 1 : 0); + const __half *ps = (i < 5) ? ST[6 + 2 * i] : ST[16]; + const int po = (i < 5) ? 0 : 1, pl = 1 + (i < 5); + dense(s, 3, p.M1, k3sw + bi * SW, k3sb + bi * 768, k3sa + bi * 384, cur, + nxt, ps, NS[5 + 2 * i], 1536, 768, p.s1, p.s1.SC, 1, po, pl, + p.W3s ? p.W3s + (bi - 1) * SW : nullptr, + i == 0 ? 1 : (i < 5 ? 2 : 3)); + __half *t = cur; cur = nxt; nxt = t; + } + TCONV_BPACK_DISPATCH1(2, 9, 16, p.s1.T, 384, B, + cur, H16(W[12]), F32(W[13]), F32(W[14]), p.P, ST[17], NS[16], 768, 384, + p.s1.T, p.s1.L, p.s1.SC, 1, p.s2.L, p.s2.SC, 6, 0, 6, B); + dense7(s, 0, p.P, p.W7[0], F32(W[22]), F32(W[23]), nullptr, p.M2, ST[18], + NS[17], p.s2, p.s2.SC, 0, 6); + dense7(s, 1, p.M2, p.W7[1], F32(W[25]), F32(W[26]), p.P, p.Q, ST[19], + NS[18], p.s2, p.s2.SC, 5, 1); + TCONV_BPACK_DISPATCH(3, 7, 16, p.s2.T, 128, B, + p.Q, H16(W[15]), F32(W[16]), F32(W[17]), p.R, ST[20], NS[19], 384, 128, + p.s2.T, p.s2.L, p.s2.SC, 5, p.s3.L, p.s3.SC, 6, 0, 6, B); + dense7(s, 2, p.R, p.W7[2], F32(W[28]), F32(W[29]), nullptr, p.M3, ST[21], + NS[20], p.s3, p.s3.SC, 0, 6); + dense7(s, 3, p.M3, p.W7[3], F32(W[31]), F32(W[32]), p.R, p.S, ST[22], + NS[21], p.s3, p.s3.SC, 5, 1); + if (p.s3.T >= 1500) + pdl_launch(transpose_conv_kernel<4, 7, 8, true>, dim3(1, 32, B), + dim3(tc_blk(p.s3.T, 8)), + (size_t)tc_blk(p.s3.T, 8) * (8 * 7) * 2 + 64, s, + p.S, H16(W[18]), F32(W[19]), F32(W[20]), p.U, ST[23], + NS[22], 128, 32, p.s3.T, p.s3.L, p.s3.SC, 5, p.s4.L, + p.s4.SC, 6, 0, 6, B); + else + pdl_launch(transpose_conv_kernel<4, 7, 4, true>, dim3(1, 32, B), + dim3(tc_blk(p.s3.T, 4)), + (size_t)tc_blk(p.s3.T, 4) * (4 * 7) * 2 + 64, s, + p.S, H16(W[18]), F32(W[19]), F32(W[20]), p.U, ST[23], + NS[22], 128, 32, p.s3.T, p.s3.L, p.s3.SC, 5, p.s4.L, + p.s4.SC, 6, 0, 6, B); + { + const int ntile = DIVUP(p.s4.N, TL_BN); + const int grid = ntile < 132 ? ntile : 132; + pdl_launch_pf(tail_fused_kernel, dim3(grid), dim3(TL_NT), TL_SHARED, s, + p.U, (size_t)32 * p.s4.SC * sizeof(__half), + p.U, p.W7[4], p.W7[5], H16(W[39]), F32(W[34]), F32(W[35]), + F32(W[37]), F32(W[38]), F32(W[40]), ST[24], ST[25], + static_cast<__half *>(a.waveform), NS[23], NS[24], NS[25], B, + p.s4.T, p.s4.L, p.s4.SC, p.s4.N, ntile); + } +} diff --git a/deploy/nemotron_native/kernels/b512/native/codec/native_stage1.cu b/deploy/nemotron_native/kernels/b512/native/codec/native_stage1.cu new file mode 100644 index 0000000..485fa50 --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/native/codec/native_stage1.cu @@ -0,0 +1,359 @@ +#include "native_stage1.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#ifndef NATIVE_CODEC_PREFETCH +#define NATIVE_CODEC_PREFETCH 1 +#endif +#ifndef NATIVE_CODEC_PACK_SUBCTAS +#define NATIVE_CODEC_PACK_SUBCTAS 1 +#endif +namespace cg=cooperative_groups; +constexpr int PHYSICAL_THREADS=512,MAIN_ARENA=227840,AUX_BYTES=4096,META_OFFSET=231936,TOTAL_SHARED=232192; +constexpr int SLICE_BYTES=MAIN_ARENA+AUX_BYTES; +constexpr int MAX_SUBCTA=8; +struct NativeMeta { + dim3 grid,threads; + uint3 block[MAX_SUBCTA]; + NativeStage1Status *status;const int32_t *lengths;int padded_frames; + int submask,tshift,tmask,slice; + int resident; + int identity_slots; +}; +static_assert(sizeof(NativeMeta)<=TOTAL_SHARED-META_OFFSET); +__device__ __forceinline__ char *native_base(){extern __shared__ __align__(16) char arena[];return arena;} +__device__ __forceinline__ NativeMeta *native_meta(){return reinterpret_cast(native_base()+META_OFFSET);} +__device__ __forceinline__ int native_subid(){const NativeMeta *m=native_meta();return (int)((threadIdx.x>>m->tshift)&(unsigned)m->submask);} +__device__ __forceinline__ char *native_arena(){const NativeMeta *m=native_meta();return native_base()+native_subid()*m->slice;} +__device__ __forceinline__ char *native_aux(){const NativeMeta *m=native_meta();return native_base()+native_subid()*m->slice+(m->slice-AUX_BYTES);} +__device__ __forceinline__ uint3 native_tid(){uint3 r;r.x=threadIdx.x&(unsigned)native_meta()->tmask;r.y=0u;r.z=0u;return r;} +__device__ __forceinline__ dim3 native_block(){const uint3 b=native_meta()->block[native_subid()];return dim3(b.x,b.y,b.z);} +__device__ __forceinline__ dim3 native_grid(){return native_meta()->grid;} +__device__ __forceinline__ dim3 native_threads(){return native_meta()->threads;} +__device__ __forceinline__ int native_resident(){return native_meta()->resident;} +__device__ __forceinline__ void native_local_sync(){ + const NativeMeta *m=native_meta();unsigned n=m->threads.x,id=15u-(unsigned)native_subid(); + asm volatile("barrier.sync %0, %1;"::"r"(id),"r"(n):"memory"); +} +__device__ __forceinline__ dim3 &native_config_grid(cudaLaunchConfig_t &c){return c.gridDim;} +__device__ __forceinline__ dim3 &native_config_threads(cudaLaunchConfig_t &c){return c.blockDim;} +__device__ __forceinline__ int native_true_tokens(int row,int padded_tokens){ + const NativeMeta *m=native_meta(); + return (padded_tokens/m->padded_frames)*m->lengths[row]; +} +__device__ __forceinline__ void native_warm(const char *p,size_t n,int k,int nk,int me){ + if(p==nullptr||n<16||nk<=0)return; + const size_t a=(size_t)p&15u;p-=a;n=(n+a)&~(size_t)15; + const uint4 *q=reinterpret_cast(p); + const size_t cnt=n>>4,step=(size_t)nk*PHYSICAL_THREADS; + uint32_t acc=0u; + for(size_t i=(size_t)k*PHYSICAL_THREADS+(size_t)me;i __device__ __forceinline__ void native_invoke_pf(F fn,dim3 g,dim3 t,size_t shared,const void *pf,size_t pfb,A&&... args){ + const int me=threadIdx.x,owner=blockIdx.x,owners=gridDim.x; + const cg::grid_group all=cg::this_grid();NativeMeta *meta=native_meta(); + if(t.x>512||t.y!=1||t.z!=1||t.x%32||shared>MAIN_ARENA){ + if(me==0&&owner==0)meta->status->code=10; + all.sync();return; + } + int sub=1,tshift=9; + const unsigned tx=t.x; + const int vtotal=(int)(g.x*g.y*g.z); + if(NATIVE_CODEC_PACK_SUBCTAS && (tx&(tx-1))==0u){ + int s=0;while((1u<MAX_SUBCTA)cand=MAX_SUBCTA; + const int want=(vtotal+owners-1)/owners; + while(cand>1&&cand>want)cand>>=1; + while(cand>1&&shared>(size_t)(SLICE_BYTES/cand-AUX_BYTES))cand>>=1; + sub=cand; + } + const int slice=(SLICE_BYTES/sub)&~127; + if(me==0){meta->grid=g;meta->threads=t;meta->submask=sub-1;meta->tshift=tshift; + meta->tmask=(int)((1u<slice=slice;} + const int total=vtotal,grp=me>>tshift; + const bool holder=me=busy&&pf!=nullptr)native_warm((const char*)pf,pfb,owner-busy,owners-busy,me); + for(int base=owner*sub;baseblock[grp]=(g.y==1u&&g.z==1u)?make_uint3((unsigned)i,0u,0u) + :make_uint3(i%g.x,(i/g.x)%g.y,i/(g.x*g.y)); + __syncthreads(); + if(holder&&i(args)...); + asm volatile("cp.async.wait_all;":::"memory"); + if(t.x>=128)asm volatile("wgmma.wait_group.sync.aligned 0;":::"memory"); + } + __syncthreads(); + } + all.sync(); +} +template __device__ __forceinline__ void native_invoke(F fn,dim3 g,dim3 t,size_t shared,A&&... args){ + native_invoke_pf(fn,g,t,shared,nullptr,(size_t)0,static_cast(args)...); +} +template __device__ __forceinline__ void native_cudaLaunchKernelEx(cudaLaunchConfig_t *c,F fn,A... args){ + native_invoke(fn,c->gridDim,c->blockDim,c->dynamicSmemBytes,args...); +} +template __device__ __forceinline__ void native_cudaLaunchKernelEx_pf(cudaLaunchConfig_t *c,F fn,const void *pf,size_t pfb,A... args){ + native_invoke_pf(fn,c->gridDim,c->blockDim,c->dynamicSmemBytes,pf,pfb,args...); +} +#define blockIdx native_block() +#define gridDim native_grid() +#define blockDim native_threads() +#define threadIdx native_tid() +#define __syncthreads() native_local_sync() +#include "legacy_codec.cuh" +#undef __syncthreads +#undef threadIdx +#undef blockDim +#undef gridDim +#undef blockIdx +namespace { +struct FastArgs { + CodecArgs codec;Plan plan;RepackSet repack;int repack_grid;int input_frames;int blocks; + const int64_t *codes;const int32_t *lengths,*slots; + const __half *state[26];__half *next[26];float *pcm;int32_t *samples; + NativeStage1Status *status;int nslots; + int channels[26],histories[26]; +}; +static_assert(sizeof(FastArgs)<=4096,"precomputed kernel argument must fit the conservative parameter budget"); +struct Handle {FastArgs args{};Plan plans[17]{};RepackSet repacks[17]{};int repack_grids[17]{};void *scratch=nullptr;size_t bytes=0;void *immutable=nullptr;size_t immutable_bytes=0;int setup_launches=0;std::vector packed;int max_frames=0,blocks=0,device=-1;}; +thread_local std::string error; +bool check(cudaError_t e,const char *where){if(e==cudaSuccess)return true;error=std::string(where)+": "+cudaGetErrorString(e);return false;} +__device__ void reject(NativeStage1Status *s,int code,int row,int frame,int detail){ + if(atomicCAS(&s->code,0,code)==0){s->row=row;s->frame=frame;s->detail=detail;} +} +__global__ __launch_bounds__(512,1) void stage1_fast_persistent(const __grid_constant__ FastArgs a){ + { + const auto all=cg::this_grid(); + const size_t me=(size_t)blockIdx.x*512+threadIdx.x,step=(size_t)gridDim.x*512; + if(threadIdx.x==0){native_meta()->status=a.status;native_meta()->lengths=a.lengths;native_meta()->padded_frames=a.codec.F;native_meta()->resident=a.blocks;native_meta()->identity_slots=(a.nslots==a.codec.B);} + if(me==0)*a.status=NativeStage1Status{0,-1,-1,0}; + all.sync();const int B=a.codec.B,F=a.codec.F,inputF=a.input_frames; + if(threadIdx.xidentity_slots,0); + __syncthreads(); + for(int row=me;rowinputF)reject(a.status,8,row,-1,n); + if(slot<0||slot>=a.nslots)reject(a.status,2,row,-1,slot); + for(int j=0;j=0&&n<=inputF)?n*1764:0; + } + for(size_t i=me;i<(size_t)B*inputF;i+=step){ + int row=i/inputF,t=i%inputF;bool any=false; + for(int c=0;c<16;++c){int64_t v=a.codes[i*16+c];if(v<0||v>=1024)reject(a.status,4,row,t,c);any|=v!=0;} + if(t=a.lengths[row]&&any)reject(a.status,6,row,t,0); + } + if(a.nslots!=B) + for(int p=0;p<26;++p){ + const size_t vectors=(size_t)a.nslots*a.channels[p]*a.histories[p]/8; + const uint4 *src=reinterpret_cast(a.state[p]); + uint4 *dst=reinterpret_cast(a.next[p]); + for(size_t i=me;icode)return; + if(!native_meta()->identity_slots) for(int p=1;p<26;++p){ + const size_t stride=(size_t)a.channels[p]*a.histories[p]; + __half *dest=(__half*)a.codec.st[p]; + for(int row=blockIdx.x;row(a.state[p]+(size_t)slot*stride); + uint4 *dst=reinterpret_cast(dest+(size_t)row*stride); + for(size_t offset=threadIdx.x;offset>(2*(c%5)))&3)-2)*.5f; + } + v=__float2half_rn(x); + } + packed[(size_t)c*SC+(size_t)row*L+j]=v; + } + all.sync(); + } + codec_program(a.codec,a.plan,a.repack,a.repack_grid, + native_meta()->identity_slots?a.state:nullptr); + if(a.status->code)return; + const int B=a.codec.B,F=a.codec.F,inputF=a.input_frames; + for(int p=0;p<26;++p){ + const size_t stride=(size_t)a.channels[p]*a.histories[p];const __half *src=(const __half*)a.codec.ns[p]; + for(int row=blockIdx.x;row(src+(size_t)row*stride); + const uint4 *prior=reinterpret_cast(a.state[p]+(size_t)slot*stride); + uint4 *dest=reinterpret_cast(a.next[p]+(size_t)slot*stride); + for(size_t offset=threadIdx.x;offset(a.codec.waveform); + float2 *pcm=reinterpret_cast(a.pcm); + const int row_pairs=inputF*882; + for(int row=blockIdx.x;rowstatus=a.status;native_meta()->resident=a.blocks;} + if(blockIdx.x==0&&threadIdx.x==0)*a.status=NativeStage1Status{0,-1,-1,0}; + all.sync(); + native_invoke(repack_rs_kernel<7>,dim3(a.repack_grid),dim3(128),MAIN_ARENA, + a.repack,H16(a.codec.inp[0]),H16(a.codec.st[0]),a.plan.LB, + 40,a.plan.s0.T,a.plan.s0.L,a.plan.s0.SC,2,a.codec.B); + native_invoke(repack3_rs_kernel,dim3(1152,12),dim3(128),MAIN_ARENA, + H16(a.codec.inp[3]),H16(a.codec.inp[6]),a.plan.W3i,a.plan.W3s, + a.plan.A,a.plan.Ac,a.plan.s1.SC,a.plan.s1.SC,DIVUP(a.plan.s1.SC,128)); +} +bool stack(Handle *h,int slot,const NativeStage1Weights *w,int first,int field,size_t elements){ + const int ids[7]={first,first+3,first+5,first+7,first+9,first+11,first+13}; + const size_t stride=elements*(field==0?sizeof(__half):sizeof(float));void *ptr=nullptr; + if(!check(cudaMalloc(&ptr,7*stride),"packed-weight allocation"))return false; + h->packed.push_back(ptr);h->args.codec.inp[slot]=ptr; + for(int i=0;i<7;++i){const void *src=field==0?w[ids[i]].weight:field==1?(const void*)w[ids[i]].bias:(const void*)w[ids[i]].alpha; + if(!check(cudaMemcpy((char*)ptr+i*stride,src,stride,cudaMemcpyDeviceToDevice),"immutable weight stack"))return false;} + return true; +} +void release(Handle *h){if(h->immutable)cudaFree(h->immutable);if(h->scratch)cudaFree(h->scratch);for(void *p:h->packed)cudaFree(p);delete h;} +} +extern "C" void *native_stage1_create(int b,int maxf,int slots,int blocks,const NativeStage1Weights *w){ + error.clear();if((b!=1&&b!=16&&b!=32&&b!=64&&b!=128&&b!=256&&b!=512)||maxf<1||maxf>16||slots4096||blocks<0||!w){error="invalid create arguments";return nullptr;} + for(int p=0;p<26;++p)if(!w[p].weight||!w[p].bias||(p!=0&&p!=25&&!w[p].alpha)){error="missing weight pointer";return nullptr;} + Handle *h=new(std::nothrow)Handle;if(!h){error="host allocation";return nullptr;} + h->max_frames=maxf;h->args.nslots=slots;h->args.codec.B=b; + const int cs[26]={40,768,1536,768,768,1536,768,1536,768,1536,768,1536,768,1536,768,1536,768,384,768,384,128,256,128,32,64,32}; + const int hs[26]={2,2,2,1,2,2,2,2,2,2,2,2,2,2,2,2,1,6,6,1,6,6,1,6,6,6}; + cudaDeviceProp prop{};int occupancy=0,coop=0; + if(!check(cudaGetDevice(&h->device),"device")||!check(cudaGetDeviceProperties(&prop,h->device),"properties")|| + !check(cudaDeviceGetAttribute(&coop,cudaDevAttrCooperativeLaunch,h->device),"cooperative support")|| + !check(cudaFuncSetAttribute(stage1_fast_persistent,cudaFuncAttributeMaxDynamicSharedMemorySize,TOTAL_SHARED),"shared-memory optin")|| + !check(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&occupancy,stage1_fast_persistent,512,TOTAL_SHARED),"occupancy")) {release(h);return nullptr;} + if(prop.major!=9||prop.minor!=0||!coop||occupancy<1){error="requires resident SM90 cooperative kernel";release(h);return nullptr;} + h->blocks=blocks?blocks:occupancy*prop.multiProcessorCount; + if(h->blocks>occupancy*prop.multiProcessorCount){error="grid exceeds cooperative occupancy";release(h);return nullptr;} + h->args.blocks=h->blocks; + const int capacityF=maxf<=3?3:(maxf<=5?5:std::max(8,maxf)); + size_t decoder=codec_scratch_bytes(b,capacityF),halfs=(size_t)b*40*2*capacityF+(size_t)b*1764*capacityF; + for(int p=0;p<26;++p){h->args.channels[p]=cs[p];h->args.histories[p]=hs[p];halfs+=2*(size_t)b*cs[p]*hs[p];} + h->bytes=decoder+halfs*sizeof(__half); + if(!check(cudaMalloc(&h->scratch,h->bytes),"scratch allocation")){release(h);return nullptr;} + if(!check(cudaMemset(h->scratch,0,h->bytes),"initial scratch padding")){release(h);return nullptr;} + h->args.codec.scratch=h->scratch;__half *q=(__half*)((char*)h->scratch+decoder); + h->args.codec.inp[0]=q;q+=(size_t)b*40*2*capacityF;h->args.codec.waveform=q;q+=(size_t)b*1764*capacityF; + for(int p=0;p<26;++p){h->args.codec.st[p]=q;q+=(size_t)b*cs[p]*hs[p];h->args.codec.ns[p]=q;q+=(size_t)b*cs[p]*hs[p];} + const void **in=h->args.codec.inp; + in[1]=w[0].weight;in[2]=w[0].bias; + if(!stack(h,3,w,1,0,(size_t)1536*768*3)||!stack(h,4,w,1,1,1536)||!stack(h,5,w,1,2,768)|| + !stack(h,6,w,2,0,(size_t)768*1536*3)||!stack(h,7,w,2,1,768)||!stack(h,8,w,2,2,384)){release(h);return nullptr;} + const int upops[4]={3,16,19,22};for(int i=0;i<4;++i){in[9+3*i]=w[upops[i]].weight;in[10+3*i]=w[upops[i]].bias;in[11+3*i]=w[upops[i]].alpha;} + for(int i=0;i<3;++i){int p=17+3*i;in[21+6*i]=w[p].weight;in[22+6*i]=w[p].bias;in[23+6*i]=w[p].alpha; + in[24+6*i]=w[p+1].weight;in[25+6*i]=w[p+1].bias;in[26+6*i]=w[p+1].alpha;} + in[39]=w[25].weight;in[40]=w[25].bias; + __half *immutable_w7[6]{};__half *immutable_w3i=nullptr,*immutable_w3s=nullptr; + if(capacityF>=3){ + size_t wh=0;for(int j=0;j<6;++j)wh+=(size_t)native_RP_COUTP(j)*native_RP_CIN(j)*7; + const size_t IW=(size_t)3*1536*768;wh+=14*IW; + h->immutable_bytes=wh*sizeof(__half); + if(!check(cudaMalloc(&h->immutable,h->immutable_bytes+sizeof(NativeStage1Status)),"immutable weight allocation")){release(h);return nullptr;} + if(!check(cudaMemset(h->immutable,0,h->immutable_bytes+sizeof(NativeStage1Status)),"immutable padding initialization")){release(h);return nullptr;} + __half *q=(__half*)h->immutable; + for(int j=0;j<6;++j){immutable_w7[j]=q;q+=(size_t)native_RP_COUTP(j)*native_RP_CIN(j)*7;} + immutable_w3i=q+IW;q+=7*IW;immutable_w3s=q+IW;q+=7*IW; + if((char*)q!=(char*)h->immutable+h->immutable_bytes){error="immutable layout size mismatch";release(h);return nullptr;} + } + for(int f:{3,5,8}) if(f<=capacityF){ + Plan p=build_plan(b,f,static_cast<__half*>(h->scratch)); + for(int j=0;j<6;++j)p.W7[j]=immutable_w7[j];p.W3i=immutable_w3i;p.W3s=immutable_w3s; + const void *const *W=h->args.codec.inp; + RepackSet js{}; + const void *SRC[RP_JOBS] = {W[21], W[24], W[27], W[30], W[33], W[36]}; + for (int j = 0; j < RP_JOBS; ++j) { + js.src[j] = H16(SRC[j]); + js.dst[j] = p.W7[j]; + js.Cout[j] = native_RP_COUT(j); + js.CoutP[j] = native_RP_COUTP(j); + js.Cin[j] = native_RP_CIN(j); + } + int acc = 0; + for (int j = 0; j < RP_JOBS; ++j) { + js.base[j] = acc; + js.nch[j] = native_RP_CIN(j) / 16; + js.rs[j] = (j < 4) ? 1 : 0; + acc += (native_RP_COUTP(j) / 16) * DIVUP(native_RP_CIN(j), 64); + } + js.k3base = acc; + js.k3src[0] = H16(W[3]); + js.k3src[1] = H16(W[6]); + js.k3dst[0] = p.W3i ? p.W3i - (size_t)3 * 1536 * 768 : nullptr; + js.k3dst[1] = p.W3s ? p.W3s - (size_t)3 * 768 * 1536 : nullptr; + const int k3n = p.W3i ? 2304 : 0; + h->plans[f]=p;h->repacks[f]=js;h->repack_grids[f]=acc+k3n; + } + if(capacityF>=3){ + FastArgs a=h->args;a.codec.F=3;a.plan=h->plans[3];a.repack=h->repacks[3];a.repack_grid=h->repack_grids[3]; + a.status=(NativeStage1Status*)((char*)h->immutable+h->immutable_bytes); + int setup_occupancy=0; + if(!check(cudaFuncSetAttribute(stage1_prepare_immutable,cudaFuncAttributeMaxDynamicSharedMemorySize,TOTAL_SHARED),"setup shared-memory optin")|| + !check(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&setup_occupancy,stage1_prepare_immutable,512,TOTAL_SHARED),"setup occupancy")|| + setup_occupancy<1||h->blocks>setup_occupancy*prop.multiProcessorCount){ + if(error.empty())error="setup grid exceeds cooperative occupancy";release(h);return nullptr;} + void *params[]={&a}; + if(!check(cudaLaunchCooperativeKernel((void*)stage1_prepare_immutable,dim3(h->blocks),dim3(512),params,TOTAL_SHARED,nullptr),"immutable setup launch")|| + !check(cudaStreamSynchronize(nullptr),"immutable setup completion")){release(h);return nullptr;} + NativeStage1Status status{}; + if(!check(cudaMemcpy(&status,a.status,sizeof(status),cudaMemcpyDeviceToHost),"immutable setup status")){release(h);return nullptr;} + if(status.code){error="immutable setup device status "+std::to_string(status.code);release(h);return nullptr;} + h->setup_launches=1; + } + return h; +} +extern "C" int native_stage1_run(void *handle,const int64_t *codes,const int32_t *lengths,const int32_t *slots, + const void *const *states,void *const *next,float *pcm,int32_t *samples,NativeStage1Status *status,int f,void *stream){ + error.clear();Handle *h=(Handle*)handle; + if(!h||!codes||!lengths||!slots||!states||!next||!pcm||!samples||!status){error="null run argument";return 1;} + if(f<1||f>8||f>h->max_frames){error="true-state variant supports public F1..8";return 1;} + const int paddedF=f<=3?3:(f<=5?5:8); + int device=-1;if(!check(cudaGetDevice(&device),"device")||device!=h->device){error="device mismatch";return 1;} + FastArgs a=h->args;a.input_frames=f;a.plan=h->plans[paddedF];a.repack=h->repacks[paddedF];a.repack_grid=h->repack_grids[paddedF];a.codec.F=paddedF;a.codes=codes;a.lengths=lengths;a.slots=slots;a.pcm=pcm;a.samples=samples;a.status=status; + for(int p=0;p<26;++p){if(!states[p]||!next[p]||states[p]==next[p]){error="invalid/aliased state";return 1;}a.state[p]=(const __half*)states[p];a.next[p]=(__half*)next[p];} + void *params[]={&a};if(!check(cudaLaunchCooperativeKernel((void*)stage1_fast_persistent,dim3(h->blocks),dim3(512),params,TOTAL_SHARED,(cudaStream_t)stream),"fast Stage1 launch"))return 1; + return check(cudaPeekAtLastError(),"immediate launch status")?0:1; +} +extern "C" int native_stage1_destroy(void *p){error.clear();if(!p)return 0;release((Handle*)p);return check(cudaPeekAtLastError(),"destroy status")?0:1;} +extern "C" const char *native_stage1_last_error(){return error.c_str();} +extern "C" size_t native_stage1_scratch_bytes(void *p){return p?((Handle*)p)->bytes:0;} +extern "C" int native_stage1_resident_blocks(void *p){return p?((Handle*)p)->blocks:0;} +extern "C" int native_stage1_setup_kernel_launches(void *p){return p?((Handle*)p)->setup_launches:0;} +extern "C" size_t native_stage1_immutable_weight_bytes(void *p){return p?((Handle*)p)->immutable_bytes:0;} diff --git a/deploy/nemotron_native/kernels/b512/native/codec/native_stage1.h b/deploy/nemotron_native/kernels/b512/native/codec/native_stage1.h new file mode 100644 index 0000000..a95b5f0 --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/native/codec/native_stage1.h @@ -0,0 +1,26 @@ +#pragma once +#include +#include +struct NativeStage1Weights { const void *weight; const float *bias; const float *alpha; }; +struct NativeStage1Status { int32_t code, row, frame, detail; }; +#ifdef __cplusplus +extern "C" { +#endif +void *native_stage1_create(int batch, int max_frames, int state_slots, + int resident_blocks, const NativeStage1Weights *weights26); +int native_stage1_run(void *handle, const int64_t *codes, const int32_t *lengths, + const int32_t *slots, const void *const *states26, + void *const *next_states26, float *pcm, int32_t *samples, + NativeStage1Status *status, int frames, void *stream); +int native_stage1_run_profile(void *handle, const int64_t *codes, const int32_t *lengths, + const int32_t *slots, const void *const *states26, + void *const *next_states26, float *pcm, int32_t *samples, + NativeStage1Status *status, int frames, void *stream, + uint64_t *phase_timestamps30); +int native_stage1_destroy(void *handle); +const char *native_stage1_last_error(); +size_t native_stage1_scratch_bytes(void *handle); +int native_stage1_resident_blocks(void *handle); +#ifdef __cplusplus +} +#endif diff --git a/deploy/nemotron_native/kernels/b512/native/codec/repack3.inc b/deploy/nemotron_native/kernels/b512/native/codec/repack3.inc new file mode 100644 index 0000000..2f61c93 --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/native/codec/repack3.inc @@ -0,0 +1,66 @@ +__device__ __noinline__ void repack3_kernel( + const __half *__restrict__ win, const __half *__restrict__ wsk, + __half *__restrict__ din, __half *__restrict__ dsk) { + int4 (*sm)[49]=reinterpret_cast(native_arena()); + const int job = blockIdx.y; + const int Cout = job ? 768 : 1536; + const int Cin = job ? 1536 : 768; + const __half *src = job ? wsk : win; + __half *dst = job ? dsk : din; + const int GPB = Cin >> 3; + const int Kg = Cin * 3; + const int GT = GPB >> 4; + const int RT = Cout >> 5; + const int t = threadIdx.x; + for (int tile = blockIdx.x; tile < 6 * RT * GT; tile += gridDim.x) { + const int gblk = tile % GT; + const int rest = tile / GT; + const int corow = rest % RT; + const int b = rest / RT; + const int co0 = corow * 32; + const int g0 = gblk * 16; + const __half *sbase = src + (size_t)(b + 1) * Cout * Kg + + (size_t)co0 * Kg + g0 * 24; +#pragma unroll + for (int j = 0; j < 6; ++j) { + const int idx = t + j * 256; + const int r = idx / 48, c = idx - (idx / 48) * 48; + sm[r][c] = *reinterpret_cast(sbase + (size_t)r * Kg + + c * 8); + } + __syncthreads(); + __half *dbase = dst + (size_t)b * (3 * (size_t)Cout * Cin); +#pragma unroll + for (int p = 0; p < 2; ++p) { + const int u = t + p * 256; + const int col = (u & 7) + 8 * (u >> 7); + const int g = g0 + ((u >> 3) & 15); + const int co = co0 + col; + const int c = g >> 1, v = g & 1; + const int4 a = sm[col][((u >> 3) & 15) * 3]; + const int4 e = sm[col][((u >> 3) & 15) * 3 + 1]; + const int4 f = sm[col][((u >> 3) & 15) * 3 + 2]; + uint32_t h[12]; + h[0] = a.x; h[1] = a.y; h[2] = a.z; h[3] = a.w; + h[4] = e.x; h[5] = e.y; h[6] = e.z; h[7] = e.w; + h[8] = f.x; h[9] = f.y; h[10] = f.z; h[11] = f.w; + __half *d = dbase + (size_t)c * (48 * Cout) + + (size_t)(co >> 3) * 128 + v * 64 + (co & 7) * 8; +#pragma unroll + for (int tp = 0; tp < 3; ++tp) { + uint32_t o[4]; +#pragma unroll + for (int jj = 0; jj < 4; ++jj) { + const uint32_t g0w = h[3 * jj], g1w = h[3 * jj + 1], + g2w = h[3 * jj + 2]; + o[jj] = (tp == 0) ? __byte_perm(g0w, g1w, 0x7610) + : (tp == 1) ? __byte_perm(g0w, g2w, 0x5432) + : __byte_perm(g1w, g2w, 0x7610); + } + *reinterpret_cast(d + (size_t)tp * (16 * Cout)) = + make_int4(o[0], o[1], o[2], o[3]); + } + } + __syncthreads(); + } +} diff --git a/deploy/nemotron_native/kernels/b512/native/codec/tail_kernel.inc b/deploy/nemotron_native/kernels/b512/native/codec/tail_kernel.inc new file mode 100644 index 0000000..054ee0f --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/native/codec/tail_kernel.inc @@ -0,0 +1,502 @@ +#define TL_NT 512 +#define TL_MW 512 +#define TL_MLD TL_MW +#define TL_VW 504 +#define TL_UW 520 +#define TL_BN 496 +#define TL_VLD 520 +#define TL_LBO_U (TL_UW * 8) +#define TL_LBO_M (TL_MLD * 8) +#define TL_LBO_V (TL_VLD * 8) +#define TL_SBO_A 128 +#define TL_PU_E (4 * TL_UW * 8) +#define TL_PM_E (8 * TL_MLD * 8) +#define TL_PV_E (32 * TL_VLD) +#define TL_AI_E (14 * 1024) +#define TL_AS_E (32 * 512 + 512) +#define TL_SHARED ((2 * TL_PU_E + TL_PM_E + TL_PV_E + TL_AI_E + TL_AS_E) * 2) +__device__ __forceinline__ void tl_stmatrix4(void *dst, const uint32_t r[4], + bool trans) { + const uint32_t a = static_cast(__cvta_generic_to_shared(dst)); + if (trans) + asm volatile( + "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1,%2,%3,%4};\n" ::"r"( + a), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3])); + else + asm volatile( + "stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1,%2,%3,%4};\n" ::"r"( + a), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3])); +} +__device__ __forceinline__ void tl_ldmatrix4t(uint32_t r[4], const void *src) { + const uint32_t a = static_cast(__cvta_generic_to_shared(src)); + asm volatile( + "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(a)); +} +__device__ __forceinline__ void tl_stmatrix1t(void *dst, uint32_t r) { + const uint32_t a = static_cast(__cvta_generic_to_shared(dst)); + asm volatile( + "stmatrix.sync.aligned.m8n8.x1.trans.shared.b16 [%0], {%1};\n" :: + "r"(a), "r"(r)); +} +template +__device__ __forceinline__ void tl_skip_band( + const __half *PU, __half *PM, const __half *AS, + __half *PV, int c0, int warp, int lane, int er0, const float *sh_ask, + const float *sh_ari_v, float b0v, float b1v) { + constexpr int NB8 = BW / 8; + float acc[BW / 2]; +#pragma unroll + for (int i = 0; i < BW / 2; ++i) acc[i] = ((i & 2) == 0) ? b0v : b1v; + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int c = 0; c < 4; ++c) { + const uint64_t da0 = + wg_desc(AS + c * 8 * 512, 128u, TL_SBO_A * 2u); + const uint64_t db0 = + wg_desc(PM + (2 * c) * TL_LBO_M + (c0 + 2) * 8, + TL_LBO_M * 2u, 128u); +#pragma unroll + for (int p = 0; p < 4; ++p) { + const uint64_t da = da0 + 64 * (2 * p); + const uint64_t db = db0 + 2 * p; + if constexpr (BW == 120) { + WGMMA_M64N120K16_T(acc, da, db, 0, 0); + } else { + WGMMA_M64N128K16_T(acc, da, db, 0, 0); + } + } + } + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + const int m_addr = lane >> 3, row = lane & 7; + const int cg = (warp & 1) * 2; + if (er0 >= 32) { +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = __floats2half2_rn( + acc[jj * 4 + h * 2 + 0], acc[jj * 4 + h * 2 + 1]); + r[m] = *reinterpret_cast(&hv); + } else { + r[m] = 0; + } + } + const int jj = j + (m_addr >> 1), h = m_addr & 1; + tl_stmatrix4(PV + (cg + h) * TL_LBO_V + (c0 + jj * 8 + row) * 8, + r, true); + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + if (er0 < 32) { +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; + const int jja = j + (m_addr >> 1), ha = m_addr & 1; + tl_ldmatrix4t(r, PV + (cg + ha) * TL_LBO_V + + (c0 + jja * 8 + row + 1) * 8); +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = *reinterpret_cast<__half2 *>(&r[m]); + acc[jj * 4 + h * 2 + 0] += __half2float(hv.x); + acc[jj * 4 + h * 2 + 1] += __half2float(hv.y); + } + } + } +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; + const int jja = j + (m_addr >> 1), ha = m_addr & 1; + tl_ldmatrix4t(r, PU + (cg + ha) * TL_LBO_U + + (c0 + jja * 8 + row + 16) * 8); +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = *reinterpret_cast<__half2 *>(&r[m]); + acc[jj * 4 + h * 2 + 0] += __half2float(hv.x); + acc[jj * 4 + h * 2 + 1] += __half2float(hv.y); + } + } + } + } + if (er0 >= 32) return; + if (er0 < 16) { + const float a0 = sh_ask[er0], r0 = sh_ari_v[er0]; + const float a1 = sh_ask[er0 + 8], r1 = sh_ari_v[er0 + 8]; +#pragma unroll + for (int i = 0; i < BW / 2; ++i) { + const int h = (i & 2) >> 1; + const float v = acc[i]; + const float sn = __sinf((h ? a1 : a0) * v); + acc[i] = __fmaf_rn(sn * sn, h ? r1 : r0, v); + } + } else { +#pragma unroll + for (int i = 0; i < BW / 2; ++i) { + const float v = acc[i]; + acc[i] = v > 0.f ? v : 0.01f * v; + } + } +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = __floats2half2_rn( + acc[jj * 4 + h * 2 + 0], acc[jj * 4 + h * 2 + 1]); + r[m] = *reinterpret_cast(&hv); + } else { + r[m] = 0; + } + } + const int jj = j + (m_addr >> 1), h = m_addr & 1; + tl_stmatrix4(PM + (cg + h) * TL_LBO_V + + (c0 + jj * 8 + row) * 8, + r, true); + } +} +__device__ __noinline__ void tail_fused_kernel( + const __half *__restrict__ U, const __half *__restrict__ win, + const __half *__restrict__ wsk, const __half *__restrict__ wpo, + const float *__restrict__ bin, const float *__restrict__ ain, + const float *__restrict__ bsk, const float *__restrict__ ask, + const float *__restrict__ bpo, const __half *__restrict__ st_m4, + const __half *__restrict__ st_v, __half *__restrict__ wave, + __half *__restrict__ ns_u, __half *__restrict__ ns_m4, + __half *__restrict__ ns_v, int B, int T4, int L4, int SC4, int N4, + int NTILE) { + char *tl_arena=reinterpret_cast(native_arena()); + __half *const PU = reinterpret_cast<__half *>(tl_arena); + __half *const PM = PU + 2 * TL_PU_E; + __half *const PV = PM + TL_PM_E; + __half *const AI = PV + TL_PV_E; + __half *const AS = AI + TL_AI_E; + __half *const PB = AS + 32 * 512; + float *sh_ain=(float*)native_aux(), *sh_ari=sh_ain+32, *sh_ask=sh_ari+32, *sh_ari_v=sh_ask+16; + const int tid = threadIdx.x; + const int wg = tid >> 7; + const int warp = (tid & 127) >> 5; + const int lane = tid & 31; + const int er0 = warp * 16 + (lane >> 2); + const int cb2 = (lane & 3) * 2; + for (int i = tid; i < 32; i += TL_NT) { + const float a = ain[i]; + sh_ain[i] = a; + sh_ari[i] = __frcp_rn(a + 1e-9f); + } + for (int i = tid; i < 16; i += TL_NT) { + const float a = ask[i]; + sh_ask[i] = a; + sh_ari_v[i] = __frcp_rn(a + 1e-9f); + } + for (int i = tid; i < TL_AI_E / 8; i += TL_NT) { + const int blk = i / 128, off = (i - blk * 128) * 8; + const int c = blk / 7, k = blk - c * 7; + *reinterpret_cast(AI + blk * 1024 + off) = + *reinterpret_cast(win + (size_t)c * 14336 + + (size_t)k * 2048 + off); + } + for (int i = tid; i < 32 * 64; i += TL_NT) { + const int blk = i / 64, off = (i - blk * 64) * 8; + const int c = blk >> 3, k = blk & 7; + *reinterpret_cast(AS + blk * 512 + off) = + (k < 7) ? *reinterpret_cast( + wsk + (size_t)c * 14336 + (size_t)k * 2048 + off) + : make_int4(0, 0, 0, 0); + } + for (int i = tid; i < 4 * (TL_VLD - TL_VW) * 8; i += TL_NT) { + const int g = i / ((TL_VLD - TL_VW) * 8); + PV[g * TL_LBO_V + TL_VW * 8 + (i - g * ((TL_VLD - TL_VW) * 8))] = + __float2half(0.f); + } + for (int i = tid; i < 512; i += TL_NT) PB[i] = __float2half(0.f); + for (int i = tid; i < 2 * 16 * 8; i += TL_NT) { + const int c = i / (16 * 8); + const int r = i - c * 16 * 8; + const int ki = r / 8, k = r & 7; + PB[c * 128 + (ki >> 3) * 64 + (ki & 7) * 8 + k] = + (k < 7) ? wpo[(c * 16 + ki) * 7 + k] : __float2half(0.f); + } + const float bias_i0 = bin[er0], bias_i1 = bin[er0 + 8]; + const float bias_s0 = (er0 < 32) ? bsk[er0] : 0.f; + const float bias_s1 = (er0 + 8 < 32) ? bsk[er0 + 8] : 0.f; + const float bpo0 = bpo[0]; + constexpr int UCOL = TL_UW / 32; + constexpr int UUNIT = 4 * UCOL; + constexpr int UWARP = TL_NT / 32; + constexpr int UITER = (UUNIT + UWARP - 1) / UWARP; + const int uwarp = tid >> 5; + const int uq = lane & 3; + int ugc[UITER], uoc[UITER]; +#pragma unroll + for (int t = 0; t < UITER; ++t) { + const int bi = uwarp + t * UWARP; + const int gc = bi / UCOL; + ugc[t] = gc; + uoc[t] = bi - gc * UCOL; + } + const __half *const urow = + U + (size_t)(lane >> 2) * SC4 + 8 * uq; + __syncthreads(); + PDL_ENTER(); +#define TL_STAGE_U(PUD,UB) { _Pragma("unroll") for (int t2 = 0; t2 < UITER; t2 += 2) { uint32_t u[2][4]; _Pragma("unroll") for (int d = 0; d < 2; ++d) { const int t = t2 + d; const int n = (UB) + uoc[t] * 32; const int gn = n + 8 * uq; const bool ok = gn >= 0 && gn < N4; int4 v = make_int4(0, 0, 0, 0); if (ok) v = *reinterpret_cast( urow + (size_t)ugc[t] * (8 * (size_t)SC4) + n); u[d][0] = v.x; u[d][1] = v.y; u[d][2] = v.z; u[d][3] = v.w; } _Pragma("unroll") for (int d = 0; d < 2; ++d) { const int t = t2 + d; uint32_t *q = u[d]; uint32_t x, y; x = (uq & 1) ? q[0] : q[1]; y = (uq & 1) ? q[2] : q[3]; x = __shfl_xor_sync(0xffffffffu, x, 1); y = __shfl_xor_sync(0xffffffffu, y, 1); if (uq & 1) { q[0] = x; q[2] = y; } else { q[1] = x; q[3] = y; } x = (uq & 2) ? q[0] : q[2]; y = (uq & 2) ? q[1] : q[3]; x = __shfl_xor_sync(0xffffffffu, x, 2); y = __shfl_xor_sync(0xffffffffu, y, 2); if (uq & 2) { q[0] = x; q[1] = y; } else { q[2] = x; q[3] = y; } tl_stmatrix4((PUD) + ugc[t] * TL_LBO_U + (uoc[t] * 32 + lane) * 8, q, true); } } if (uwarp < 4) { const int n = (UB) + 512; const int gn = n + 2 * (lane & 3); uint32_t q = 0; if (gn >= 0 && gn + 1 < N4) q = *reinterpret_cast( U + (size_t)(uwarp * 8 + (lane >> 2)) * SC4 + gn); tl_stmatrix1t((PUD) + uwarp * TL_LBO_U + (512 + (lane & 7)) * 8, q); } } + TL_STAGE_U(PU, (int)blockIdx.x * TL_BN - 16) + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + int ubuf = 0; + for (int tile = blockIdx.x; tile < NTILE; tile += gridDim.x, ubuf ^= 1) { + const int g0 = tile * TL_BN; + bool active_tile = false; + for (int row = max(0, g0 - 16) / L4; + row <= min(N4 - 1, g0 - 16 + TL_UW - 1) / L4; ++row) { + if (native_meta()->lengths[row] > 0) { + active_tile = true; + break; + } + } + __half *const PUc = PU + ubuf * TL_PU_E; + __half *const PUn = PU + (ubuf ^ 1) * TL_PU_E; + { + const int b0 = g0 / L4; + for (int s = 0; s < 2; ++s) { + const int sb = b0 + s; + if (sb >= B) break; + const int cv = (int)((long long)sb * L4 + native_true_tokens(sb,T4) - g0); + if (cv < 0 || cv >= TL_BN) continue; + for (int i = tid; i < 32 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + ns_u[((size_t)sb * 32 + co) * 6 + h] = + PUc[(co >> 3) * TL_LBO_U + + (cv + h + 16) * 8 + (co & 7)]; + } + } + } + if (active_tile) { + float acc[64]; +#pragma unroll + for (int i = 0; i < 64; ++i) + acc[i] = ((i & 2) == 0) ? bias_i0 : bias_i1; + const int c0 = wg * 128; + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int c = 0; c < 2; ++c) { + const uint64_t da0 = + wg_desc(AI + c * 7 * 1024, 128u, TL_SBO_A * 2u); + const uint64_t db0 = + wg_desc(PUc + (2 * c) * TL_LBO_U + (c0 + 2) * 8, + TL_LBO_U * 2u, 128u); +#pragma unroll + for (int k = 0; k < 7; ++k) { + const uint64_t da = da0 + 128 * k; + const uint64_t db = db0 + k; + WGMMA_M64N128K16_T(acc, da, db, 0, 0); + } + } + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); + if (tile + (int)gridDim.x < NTILE) + TL_STAGE_U(PUn, (tile + (int)gridDim.x) * TL_BN - 16) + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + if (er0 < 32) { + const float a0 = sh_ain[er0], r0 = sh_ari[er0]; + const float a1 = sh_ain[er0 + 8], r1 = sh_ari[er0 + 8]; +#pragma unroll + for (int i = 0; i < 64; ++i) { + const int h = (i & 2) >> 1; + const float v = acc[i]; + const float sn = __sinf((h ? a1 : a0) * v); + acc[i] = __fmaf_rn(sn * sn, h ? r1 : r0, v); + } + } else { +#pragma unroll + for (int i = 0; i < 64; ++i) { + const float v = acc[i]; + acc[i] = v > 0.f ? v : 0.01f * v; + } + } +#pragma unroll + for (int j = 0; j < 16; j += 2) { + uint32_t r[4]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + const __half2 hv = __floats2half2_rn( + acc[jj * 4 + h * 2 + 0], acc[jj * 4 + h * 2 + 1]); + r[m] = *reinterpret_cast(&hv); + } + const int m_addr = lane >> 3, row = lane & 7; + const int jj = j + (m_addr >> 1), h = m_addr & 1; + tl_stmatrix4(PM + (warp * 2 + h) * TL_LBO_M + + (c0 + jj * 8 + row) * 8, + r, true); + } + } else if (tile + (int)gridDim.x < NTILE) { + TL_STAGE_U(PUn, (tile + (int)gridDim.x) * TL_BN - 16) + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + { + const int base = g0 - 8; + const int b0 = (base >= 0) ? base / L4 : 0; + const int cm0 = (int)((long long)b0 * L4 - base); + const int cm1 = cm0 + L4; + const bool hit0 = (b0 < B) && (cm0 + 6 > 0) && (cm0 < TL_MW); + const bool hit1 = (b0 + 1 < B) && (cm1 + 6 > 0) && (cm1 < TL_MW); + if (hit0 || hit1) { + for (int s = 0; s < 2; ++s) { + if (!(s ? hit1 : hit0)) continue; + const int sb = b0 + s, cm = s ? cm1 : cm0; + for (int i = tid; i < 64 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + const int c = cm + h; + if (c >= 0 && c < TL_MW) + PM[(co >> 3) * TL_LBO_M + c * 8 + (co & 7)] = + st_m4[((size_t)sb * 64 + co) * 6 + h]; + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + } + } + { + const int b0 = g0 / L4; + for (int s = 0; s < 2; ++s) { + const int sb = b0 + s; + if (sb >= B) break; + const int cv = (int)((long long)sb * L4 + native_true_tokens(sb,T4) - g0); + if (cv < 0 || cv >= TL_BN) continue; + for (int i = tid; i < 64 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + ns_m4[((size_t)sb * 64 + co) * 6 + h] = + PM[(co >> 3) * TL_LBO_M + + (cv + h + 8) * 8 + (co & 7)]; + } + } + } + if (active_tile) { + const int c0 = wg * 128; + if (wg == 3) + tl_skip_band<120>(PUc, PM, AS, PV, c0, warp, lane, er0, sh_ask, + sh_ari_v, bias_s0, bias_s1); + else + tl_skip_band<128>(PUc, PM, AS, PV, c0, warp, lane, er0, sh_ask, + sh_ari_v, bias_s0, bias_s1); + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + { + const int b0 = g0 / L4; + const int cv0 = (int)((long long)b0 * L4 - g0), cv1 = cv0 + L4; + const bool hv0 = (b0 < B) && (cv0 + 6 > 0) && (cv0 < TL_VW); + const bool hv1 = (b0 + 1 < B) && (cv1 + 6 > 0) && (cv1 < TL_VW); + if (hv0 || hv1) { + for (int s = 0; s < 2; ++s) { + if (!(s ? hv1 : hv0)) continue; + const int sb = b0 + s, cv = s ? cv1 : cv0; + for (int i = tid; i < 32 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + const int c = cv + h; + if (c >= 0 && c < TL_VW) + PM[(co >> 3) * TL_LBO_V + c * 8 + (co & 7)] = + st_v[((size_t)sb * 32 + co) * 6 + h]; + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + } + } + { + const int b0 = g0 / L4; + for (int s = 0; s < 2; ++s) { + const int sb = b0 + s; + if (sb >= B) break; + const int cv = (int)((long long)sb * L4 + native_true_tokens(sb,T4) - g0); + if (cv < 0 || cv >= TL_BN) continue; + for (int i = tid; i < 32 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + ns_v[((size_t)sb * 32 + co) * 6 + h] = + PM[(co >> 3) * TL_LBO_V + (cv + h) * 8 + (co & 7)]; + } + } + } + __syncthreads(); + constexpr int PC_LD = 10; + if (active_tile) { + float *const PC = reinterpret_cast(PV); +#pragma unroll + for (int round = 0; round < 2; ++round) { + float acc[4]; + const int p0 = round * 256 + wg * 64; + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); + const uint64_t da0 = + wg_desc(PM + p0 * 8, TL_LBO_V * 2u, 128u); + const uint64_t db0 = wg_desc(PB, 128u, 128u); + const uint64_t da1 = + wg_desc(PM + 2 * TL_LBO_V + p0 * 8, + TL_LBO_V * 2u, 128u); + const uint64_t db1 = wg_desc(PB + 128, 128u, 128u); + WGMMA_M64N8K16_ZERO_T(acc, da0, db0, 0, 1); + WGMMA_M64N8K16_T(acc, da1, db1, 0, 1); + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + const int row0 = p0 + warp * 16 + (lane >> 2); + const int col = (lane & 3) * 2; + *reinterpret_cast(PC + row0 * PC_LD + col) = + make_float2(acc[0], acc[1]); + *reinterpret_cast(PC + (row0 + 8) * PC_LD + col) = + make_float2(acc[2], acc[3]); + } + } + __syncthreads(); + if (active_tile) { + const float *const PC = reinterpret_cast(PV); + const int p0 = tid; + const bool live = p0 < TL_BN; + const unsigned live_mask = __ballot_sync(0xffffffffu, live); + if (live) { + const float s0 = __fadd_rn(PC[(p0 + 0) * PC_LD + 0], + PC[(p0 + 1) * PC_LD + 1]); + const float s1 = __fadd_rn(PC[(p0 + 2) * PC_LD + 2], + PC[(p0 + 3) * PC_LD + 3]); + const float s2 = __fadd_rn(PC[(p0 + 4) * PC_LD + 4], + PC[(p0 + 5) * PC_LD + 5]); + const float s3 = PC[(p0 + 6) * PC_LD + 6]; + float y0 = __fadd_rn(__fadd_rn(s0, s1), + __fadd_rn(s2, s3)); + y0 = fminf(fmaxf(y0 + bpo0, -1.f), 1.f); + const float y1 = __shfl_down_sync(live_mask, y0, 1); + if ((p0 & 1) == 0) { + const int b0 = g0 / L4; + const int rr = g0 - b0 * L4; + int t0 = rr + p0, bo0 = b0; + int t1 = t0 + 1, bo1 = bo0; + if (t0 >= L4) { t0 -= L4; ++bo0; } + if (t1 >= L4) { t1 -= L4; ++bo1; } + if (bo0 == bo1 && bo0 < B && t1 < T4) { + *reinterpret_cast<__half2 *>( + wave + (size_t)bo0 * T4 + t0) = + __floats2half2_rn(y0, y1); + } else { + if (bo0 < B && t0 < T4) + wave[(size_t)bo0 * T4 + t0] = __float2half(y0); + if (bo1 < B && t1 < T4) + wave[(size_t)bo1 * T4 + t1] = __float2half(y1); + } + } + } + } + } +#undef TL_STAGE_U +} diff --git a/deploy/nemotron_native/kernels/b512/native/codec/wg2_kernel.inc b/deploy/nemotron_native/kernels/b512/native/codec/wg2_kernel.inc new file mode 100644 index 0000000..4ffa7ad --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/native/codec/wg2_kernel.inc @@ -0,0 +1,280 @@ +__device__ __forceinline__ void cpa16(void *dsm, const void *src, bool pred) { + unsigned a = static_cast(__cvta_generic_to_shared(dsm)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(a), + "l"(src), "r"(pred ? 16 : 0) + : "memory"); +} +template +__device__ __forceinline__ void cpa16o(unsigned a, const void *src, bool pred) { + asm volatile("cp.async.cg.shared.global [%0+%3], [%1], 16, %2;\n" ::"r"(a), + "l"(src), "r"(pred ? 16 : 0), "n"(OFF) + : "memory"); +} +template +__device__ __noinline__ +void wg2_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int BK = CIB * KT; + constexpr int WG_BM = 64 * MSPLIT; + constexpr int WG_BNT = WG_BN * NSPLIT; + constexpr int NTS = 128 * MSPLIT * NSPLIT; + constexpr int WG_NT = NTS * KSPLIT; + constexpr int NCRAW = WG_BNT / 4; + constexpr int EPN = + (WG_NT % NCRAW == 0) ? WG_BNT : (((WG_BNT + 127) / 128) * 128); + constexpr int WG_CLD = EPN + 8; + constexpr bool NB8_P2 = ((WG_BN / 8) & (WG_BN / 8 - 1)) == 0; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int ST_E = SA * A_E + SB * NSPLIT * B_E; + constexpr int SBO_AE = KG * WG_SLOT; + constexpr int LBO_BE = (WG_BN / 8) * WG_SLOT; + constexpr int NB8 = WG_BN / 8; + constexpr bool ALDG = (WG_BN <= 64); + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *smem = reinterpret_cast<__half *>(arena); + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int kgrp = tid / NTS; + const int stid = tid - kgrp * NTS; + const int wgl = stid >> 7; + const int mg = wgl % MSPLIT; + const int ngid = wgl / MSPLIT; + const int warp = (stid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * WG_BM; + const int nb = blockIdx.y * WG_BNT; + const int n0 = nb + ngid * WG_BN; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + const int NCH = (Kg + BK - 1) / BK; + if (tid < WG_BM) { + const int co = m0 + tid; + if (do_snake && co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + for (int i = tid; i < EPN; i += WG_NT) { + const int n = nb + i; + colok[i] = (i < WG_BNT && n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mg * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (kgrp == 0 && eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (kgrp == 0 && eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + constexpr int KGX = ALDG ? (KG <= 8 ? 8 : (KG <= 16 ? 16 : 32)) : KG; + constexpr int ATOT = WG_BM * KGX; + constexpr int AREG = (ATOT + NTS - 1) / NTS; + constexpr int BTOT = NSPLIT * CIB * NB8; + constexpr int BREG = (BTOT + NTS - 1) / NTS; + const __half *aptr[AREG]; + int aoff[AREG]; + int akmax[AREG]; +#pragma unroll + for (int r = 0; r < AREG; ++r) { + const int idx = stid + r * NTS; + const int m = idx / KGX, v = idx - (idx / KGX) * KGX; + const int co = m0 + m; + const bool live = (ATOT % NTS == 0 || idx < ATOT) && (co < Cout) && + v < KG; + aoff[r] = (m >> 3) * SBO_AE + (v < KG ? v : 0) * WG_SLOT + (m & 7) * 8; + akmax[r] = live ? (Kg - v * 8) : 0; + aptr[r] = w + (size_t)(live ? co : 0) * Kg + v * 8; + } + const __half *bptr[BREG]; + int bcimax[BREG], bnc[BREG], bcil[BREG], bgrp[BREG]; +#pragma unroll + for (int r = 0; r < BREG; ++r) { + const int idx = stid + r * NTS; + const int gg = idx / (CIB * NB8); + const int rr = idx - gg * (CIB * NB8); + const int cil = rr / NB8, nc = rr - (rr / NB8) * NB8; + bcil[r] = cil; + bnc[r] = nc; + bgrp[r] = gg; + const bool live = (BTOT % NTS == 0 || idx < BTOT); + bcimax[r] = live ? (Cin - cil) : 0; + bptr[r] = in + (size_t)cil * SCin + nb + gg * WG_BN + nc * 8; + } + int4 br0[BREG], br1[BREG]; + int4 ar[ALDG ? AREG : 1]; + uint32_t edge[BREG]; + __half *const kaB = smem + kgrp * ST_E; + __half *const kbB = kaB + SA * A_E; + const unsigned barid = 1u + kgrp + (unsigned)(KSPLIT * native_subid()); +#define WG2_BAR() asm volatile("bar.sync %0, %1;\n" ::"r"(barid), "r"(NTS)) +#define WG2_LOADB(c) { const int _ci0 = (c) * CIB; _Pragma("unroll") for (int r = 0; r < BREG; ++r) { br0[r] = make_int4(0, 0, 0, 0); if constexpr (KT == 3) edge[r] = 0; else br1[r] = make_int4(0, 0, 0, 0); if (_ci0 < bcimax[r]) { const int4 *src = reinterpret_cast( bptr[r] + (size_t)_ci0 * SCin); br0[r] = src[0]; if constexpr (KT == 3) { if (!NB8_P2 || bnc[r] == NB8 - 1) edge[r] = src[1].x; } else { br1[r] = src[1]; } } } } +#define WG2_LOADA(c) if constexpr (ALDG) { const int _kg0 = (c) * BK; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { ar[r] = make_int4(0, 0, 0, 0); if (_kg0 < akmax[r]) ar[r] = *reinterpret_cast(aptr[r] + _kg0); } } +#define WG2_PUTA(c,AD) { _Pragma("unroll") for (int r = 0; r < AREG; ++r) { if constexpr (ALDG) { if (akmax[r] > 0) *reinterpret_cast((AD) + aoff[r]) = ar[r]; } else { const int _kg0 = (c) * BK; if (akmax[r] > 0) cpa16((AD) + aoff[r], aptr[r] + _kg0, _kg0 < akmax[r]); } } if constexpr (!ALDG) asm volatile("cp.async.commit_group;\n" ::: "memory"); } +#define WG2_STSB(BD) { _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int cil = bcil[r], nc = bnc[r]; unsigned int u[8]; u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; if constexpr (KT == 3) { if constexpr (NB8_P2) { u[4] = __shfl_down_sync(0xffffffffu, u[0], 1, NB8); if (nc == NB8 - 1) u[4] = edge[r]; } else { u[4] = edge[r]; } u[5] = u[6] = u[7] = 0; } else { u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; } if (BTOT % NTS == 0 || (stid + r * NTS) < BTOT) { _Pragma("unroll") for (int k = 0; k < KT; ++k) { const int kh = k >> 1; int4 o; if (k & 1) { o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); } else { o.x = u[kh + 0]; o.y = u[kh + 1]; o.z = u[kh + 2]; o.w = u[kh + 3]; } const int row = cil * KT + k; *reinterpret_cast((BD) + bgrp[r] * B_E + nc * WG_SLOT + (row >> 3) * LBO_BE + (row & 7) * 8) = o; } } } } +#define WG2_MMA(AD,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t da0 = wg_desc((AD) + mg * 8 * SBO_AE, WG_SLOT * 2, SBO_AE * 2); const uint64_t db0 = wg_desc((BD) + ngid * B_E, LBO_BE * 2, WG_SLOT * 2); _Pragma("unroll") for (int kk = 0; kk < BK; kk += 16) { const uint64_t da = da0 + (kk >> 3) * (WG_SLOT / 8); const uint64_t db = db0 + (kk >> 3) * (LBO_BE / 8); if constexpr (WG_BN == 64) { WGMMA_M64N64K16(acc, da, db); } else if constexpr (WG_BN == 256) { WGMMA_M64N256K16(acc, da, db); } else if constexpr (WG_BN == 104) { WGMMA_M64N104K16(acc, da, db); } else if constexpr (WG_BN == 112) { WGMMA_M64N112K16(acc, da, db); } else if constexpr (WG_BN == 120) { WGMMA_M64N120K16(acc, da, db); } else if constexpr (WG_BN == 224) { WGMMA_M64N224K16(acc, da, db); } else { static_assert(WG_BN == 128, "no wgmma macro for WG_BN"); WGMMA_M64N128K16(acc, da, db); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + constexpr int SMAX = SA > SB ? SA : SB; + WG2_LOADA(kgrp) + WG2_LOADB(kgrp) +#pragma unroll 1 + for (int s = 0; s < SMAX; ++s) { + const int c = s * KSPLIT + kgrp; + if (s < SA) { + if (c < NCH) WG2_PUTA(c, kaB + s * A_E) + if (c + KSPLIT < NCH) WG2_LOADA(c + KSPLIT) + } + if (s < SB) { + if (c < NCH) WG2_STSB(kbB + s * NSPLIT * B_E) + if (c + KSPLIT < NCH) WG2_LOADB(c + KSPLIT) + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + const int NIT = (NCH - kgrp + KSPLIT - 1) / KSPLIT; + for (int i = 0; i < NIT; ++i) { + if constexpr (!ALDG) + asm volatile("cp.async.wait_group %0;\n" ::"n"(SA - 2) : "memory"); + WG2_BAR(); + WG2_MMA(kaB + (i % SA) * A_E, kbB + (i % SB) * NSPLIT * B_E) + asm volatile("wgmma.wait_group.sync.aligned 1;\n" ::: "memory"); + if constexpr (NTS > 128) WG2_BAR(); + if (i >= 1) { + const int ca = (i + SA - 1) * KSPLIT + kgrp; + if (ca < NCH) WG2_PUTA(ca, kaB + ((i - 1) % SA) * A_E) + if (ALDG && ca + KSPLIT < NCH) WG2_LOADA(ca + KSPLIT) + const int cb = (i + SB - 1) * KSPLIT + kgrp; + if (cb < NCH) { + WG2_STSB(kbB + ((i - 1) % SB) * NSPLIT * B_E) + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if (cb + KSPLIT < NCH) WG2_LOADB(cb + KSPLIT) + } else if constexpr (ALDG) { + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + } + } + } + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); +#undef WG2_LOADB +#undef WG2_LOADA +#undef WG2_PUTA +#undef WG2_STSB +#undef WG2_MMA +#undef WG2_BAR + __syncthreads(); + float *Ck = Cs + kgrp * WG_BM * WG_CLD; + const int cb = (lane & 3) * 2; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Ck + (er0 + h * 8) * WG_CLD + + ngid * WG_BN + j * 8 + cb) = v; + } + } + __syncthreads(); + constexpr int NC = EPN / 4; + constexpr int MSTEP = WG_NT / NC; + constexpr int MITER = WG_BM / MSTEP; + constexpr int UNR = (MITER % 4 == 0) ? 4 : ((MITER % 2 == 0) ? 2 : 1); + static_assert(MITER % UNR == 0, "row batch must divide the row count"); + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + const int n = nb + lc; +#pragma unroll 1 + for (int t0 = 0; t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + v[u] = *reinterpret_cast(Cs + m * WG_CLD + lc); + if (resid != nullptr) { + const int co = m0 + m; + const int cc = co < Cout ? co : 0; + const __half2 *rp = reinterpret_cast( + resid + (size_t)cc * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + const int co = m0 + m; + if (co >= Cout) continue; + float4 vv = v[u]; +#pragma unroll + for (int g = 1; g < KSPLIT; ++g) { + float4 p = *reinterpret_cast( + Cs + g * WG_BM * WG_CLD + m * WG_CLD + lc); + vv.x += p.x; vv.y += p.y; vv.z += p.z; vv.w += p.w; + } + if (resid != nullptr) { + vv.x += __half2float(r0[u].x); vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (do_snake) { + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + } + __half *op = out + (size_t)co * SCout + n + H; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B * H; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int total = Cout * B * pad_len; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} diff --git a/deploy/nemotron_native/kernels/b512/native/codec/wg3_kernel.inc b/deploy/nemotron_native/kernels/b512/native/codec/wg3_kernel.inc new file mode 100644 index 0000000..fa55b32 --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/native/codec/wg3_kernel.inc @@ -0,0 +1,177 @@ +template +__device__ __noinline__ void wg3_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int WG_BNT = WG_BN * NSG; + constexpr int WG_CLD = WG_BNT + 8; + constexpr int WG_NT = 128 * NWG * NSG; + constexpr int WG_BM = 64 * NWG; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int SBO_AE = KG * WG_SLOT; + constexpr int LBO_BE = (WG_BN / 8) * WG_SLOT; + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + constexpr int ST_E = A_E + NSG * B_E; + __half *As0 = reinterpret_cast<__half *>(arena); + __half *Bs0 = As0 + A_E; + __half *As1 = (STAGES == 2) ? As0 + ST_E : As0; + __half *Bs1 = (STAGES == 2) ? Bs0 + ST_E : Bs0; + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int wgid = tid >> 7; + const int mgid = wgid % NWG, ngid = wgid / NWG; + const int warp = (tid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * WG_BM; + constexpr int WGROWS = 8; + const int nb = blockIdx.y * WG_BNT; + const int n0 = nb + ngid * WG_BN; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + if (tid < WG_BM) { + const int co = m0 + tid; + if (do_snake && co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + if (tid < WG_BNT) { + const int n = nb + tid; + colok[tid] = (n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mgid * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + constexpr int ATOT = WG_BM * KG; + constexpr int AREG = (ATOT + WG_NT - 1) / WG_NT; + constexpr int CIB = BK / KT; + constexpr int BTOT = NSG * CIB * (WG_BN / 8); + constexpr int BREG = (BTOT + WG_NT - 1) / WG_NT; + int4 ar[AREG]; + int4 br0[BREG], br1[BREG]; + uint32_t edge[BREG]; +#define WG_LOADG(kg0) { const int _ci0 = (kg0) / KT; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; ar[r] = make_int4(0, 0, 0, 0); if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; const int co = m0 + m, kg = (kg0) + v * 8; if (co < Cout && kg < Kg) ar[r] = *reinterpret_cast( w + (size_t)co * Kg + kg); } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; br0[r] = make_int4(0, 0, 0, 0); if constexpr (KT == 3) edge[r] = 0; else br1[r] = make_int4(0, 0, 0, 0); if (BTOT % WG_NT == 0 || idx < BTOT) { const int _g = idx / (CIB * (WG_BN / 8)); const int _r = idx - _g * (CIB * (WG_BN / 8)); const int ci = _ci0 + _r / (WG_BN / 8); const int nc = _r % (WG_BN / 8); if (ci < Cin) { const int4 *src = reinterpret_cast( in + (size_t)ci * SCin + nb + _g * WG_BN + nc * 8); br0[r] = src[0]; if constexpr (KT == 3) { if (nc == WG_BN / 8 - 1) edge[r] = src[1].x; } else { br1[r] = src[1]; } } } } } + WG_LOADG(0) +#define WG_STS(AD,BD) { _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; *reinterpret_cast( (AD) + (m >> 3) * SBO_AE + v * WG_SLOT + (m & 7) * 8) = ar[r]; } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; if (BTOT % WG_NT == 0 || idx < BTOT) { const int _g = idx / (CIB * (WG_BN / 8)); const int _r = idx - _g * (CIB * (WG_BN / 8)); const int cil = _r / (WG_BN / 8), nc = _r % (WG_BN / 8); unsigned int u[8]; u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; if constexpr (KT == 3) { u[4] = __shfl_down_sync(0xffffffffu, u[0], 1, WG_BN / 8); if (nc == WG_BN / 8 - 1) u[4] = edge[r]; u[5] = u[6] = u[7] = 0; } else { u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; } _Pragma("unroll") for (int k = 0; k < KT; ++k) { const int kh = k >> 1; int4 o; if (k & 1) { o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); } else { o.x = u[kh + 0]; o.y = u[kh + 1]; o.z = u[kh + 2]; o.w = u[kh + 3]; } const int row = cil * KT + k; *reinterpret_cast((BD) + _g * B_E + nc * WG_SLOT + (row >> 3) * LBO_BE + (row & 7) * 8) = o; } } } if constexpr (WG_BN == 64) { asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); } } +#define WG_MMA(AD,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t da0 = wg_desc((AD) + mgid * WGROWS * SBO_AE, WG_SLOT * 2, SBO_AE * 2); const uint64_t db0 = wg_desc((BD) + ngid * B_E, LBO_BE * 2, WG_SLOT * 2); _Pragma("unroll") for (int kk = 0; kk < BK; kk += 16) { const uint64_t da = da0 + (kk >> 3) * (WG_SLOT / 8); const uint64_t db = db0 + (kk >> 3) * (LBO_BE / 8); if constexpr (WG_BN == 64) { WGMMA_M64N64K16(acc, da, db); } else { WGMMA_M64N128K16(acc, da, db); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + WG_STS(As0, Bs0) + if (BK < Kg) WG_LOADG(BK) + __syncthreads(); + int cur = 0; + for (int kg0 = 0; kg0 < Kg; kg0 += BK) { + if (cur == 0) { WG_MMA(As0, Bs0) } else { WG_MMA(As1, Bs1) } + if (kg0 + BK < Kg) { + if (STAGES == 2) + asm volatile("wgmma.wait_group.sync.aligned 1;\n" ::: "memory"); + else + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + __syncthreads(); + if (cur == 0) { WG_STS(As1, Bs1) } else { WG_STS(As0, Bs0) } + if (kg0 + 2 * BK < Kg) WG_LOADG(kg0 + 2 * BK) + __syncthreads(); + } else { + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + } + if (STAGES == 2) cur ^= 1; + } +#undef WG_LOADG +#undef WG_STS +#undef WG_MMA + __syncthreads(); + const int r0 = er0; + const int cb = (lane & 3) * 2; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Cs + (r0 + h * 8) * WG_CLD + + ngid * WG_BN + j * 8 + cb) = v; + } + } + __syncthreads(); + constexpr int NC = WG_BNT / 4; + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + for (int m = lr; m < WG_BM; m += WG_NT / NC) { + const int co = m0 + m; + if (co >= Cout) continue; + const int n = nb + lc; + float4 v = *reinterpret_cast(Cs + m * WG_CLD + lc); + __half *op = out + (size_t)co * SCout + n + H; + if (resid != nullptr) { + const __half2 *rp = reinterpret_cast( + resid + (size_t)co * SCout + n + H); + __half2 r0h = rp[0], r1h = rp[1]; + v.x += __half2float(r0h.x); v.y += __half2float(r0h.y); + v.z += __half2float(r1h.x); v.w += __half2float(r1h.y); + } + if (do_snake) { + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + } + if (ok4) { + __half2 o0 = __floats2half2_rn(v.x, v.y); + __half2 o1 = __floats2half2_rn(v.z, v.w); + reinterpret_cast<__half2 *>(op)[0] = o0; + reinterpret_cast<__half2 *>(op)[1] = o1; + } else { + const float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B * H; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int total = Cout * B * pad_len; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} diff --git a/deploy/nemotron_native/kernels/b512/native/codec/wg4_step_widths.inc b/deploy/nemotron_native/kernels/b512/native/codec/wg4_step_widths.inc new file mode 100644 index 0000000..1287a6f --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/native/codec/wg4_step_widths.inc @@ -0,0 +1,468 @@ +template +__device__ __noinline__ void wg4_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int CoutP, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + if constexpr (PDL && !LATE_PDL) PDL_ENTER(); + constexpr int NT = 128 * MSPLIT * NSPLIT; + constexpr int BM = 64 * MSPLIT; + constexpr int BNT = WG_BN * NSPLIT; + constexpr int NCRAW = BNT / 4; + constexpr int EPN = (NT % NCRAW == 0) ? BNT : (((BNT + 127) / 128) * 128); + constexpr int WG_CLD = EPN + 8; + constexpr int CG = CIB / 8; + constexpr int NBR = BNT + KT - 1; + constexpr int NBRP = BC8 ? ((NBR + 7) & ~7) : ((NBR + 31) & ~31); + constexpr int UCOL = NBRP / 32; + constexpr int NUNIT = CG * UCOL; + constexpr int NWARP = NT / 32; + constexpr int BITER = (NUNIT + NWARP - 1) / NWARP; + constexpr int B_E = CIB * NBRP; + constexpr int ST_E = B_E; + constexpr int LBO_B = NBRP * 8; + constexpr int ABUF = ABUF_; + static_assert(!OC8 || PREACT, + "K=7 C8 output is only used for residual-free pair heads"); + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *smem = reinterpret_cast<__half *>(arena); + __half *Hs = smem; + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int wgid = tid >> 7; + const int mg = wgid % MSPLIT, ng = wgid / MSPLIT; + const int warp = (tid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * BM; + const int nb = blockIdx.y * BNT; + const int H = KT - 1; + bool active_tile = false; + for (int row = nb / L; row <= min(N - 1, nb + BNT - 1) / L; ++row) { + if (native_meta()->lengths[row] > 0) { + active_tile = true; + break; + } + } + if (active_tile) { + const int half = Cout >> 1; + const int NCH = (Cin + CIB - 1) / CIB; + for (int i = tid; i < BM; i += NT) { + const int co = m0 + i; + if (co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[i] = a; + row_ri[i] = __frcp_rn(a + 1e-9f); + } + } + for (int i = tid; i < EPN; i += NT) { + const int n = nb + i; + colok[i] = (i < BNT && n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mg * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + const int mtile = blockIdx.x * MSPLIT + mg; + const __half *const apt = w + ((size_t)mtile * NCH * KT) * 1024 + + (size_t)(tid & 127) * 8; + uint32_t areg[ABUF][KT][4]; + const int wi = tid >> 5; + int bdst[BITER], boff[BITER]; + bool blive[BITER]; + const __half *const bbase = in + (size_t)(lane >> 2) * SCin + nb + + 8 * (lane & 3); +#pragma unroll + for (int t = 0; t < BITER; ++t) { + const int u = wi + t * NWARP; + const int gc = (u < NUNIT) ? u / UCOL : 0; + const int cu = (u < NUNIT) ? u - gc * UCOL : 0; + blive[t] = u < NUNIT; + bdst[t] = gc * LBO_B + (cu * 32 + lane) * 8; + boff[t] = gc * 8 * SCin + cu * 32; + } + const size_t SC8i = (size_t)SCin * 8; + constexpr int BQ = B_E / 8; + constexpr int BCREG = (BQ + NT - 1) / NT; + int cdst[BCREG]; + const __half *csrc[BCREG]; + bool clive[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) { + const int q = tid + r * NT; + const int gc = q / NBRP, rem = q - gc * NBRP; + clive[r] = ((BQ % NT == 0) || q < BQ) && gc < CG; + cdst[r] = clive[r] ? gc * LBO_B + rem * 8 : 0; + csrc[r] = in + (clive[r] ? (size_t)gc * SC8i + + (size_t)nb * 8 + rem * 8 : 0); + } + uint32_t br[BITER][4]; + constexpr int CPA_STAGES = 4; + constexpr int BSTAGES = BC8 ? CPA_STAGES : STAGES; + static_assert(STAGES >= CPA_STAGES, "C8 B-panel pipeline allocation"); + static_assert(ABUF == 2 && BSTAGES == 4, "wg4 unrolled pipeline shape"); + uint64_t dslot[BSTAGES]; +#pragma unroll + for (int s = 0; s < BSTAGES; ++s) + dslot[s] = wg_desc(smem + s * ST_E + ng * WG_BN * 8, LBO_B * 2u, 128u); + uint32_t csm0[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) + csm0[r] = static_cast(__cvta_generic_to_shared(smem + cdst[r])); + uint32_t bsm0[BITER]; +#pragma unroll + for (int t = 0; t < BITER; ++t) + bsm0[t] = static_cast(__cvta_generic_to_shared(smem + bdst[t])); + const size_t CSTEP = (size_t)CG * SC8i; + const __half *crol[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) crol[r] = csrc[r] + CPA_STAGES * CSTEP; + const __half *arol = apt + (size_t)ABUF * (KT * 1024); +#define WG4_CPB(c,BD) { const size_t _o = (size_t)(c) * CG * SC8i; _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16((BD) + cdst[r], csrc[r] + _o, true); } +#define WG4_CPB_R(SLOT) { _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16o<(SLOT) * ST_E * 2>(csm0[r], crol[r], true); _Pragma("unroll") for (int r = 0; r < BCREG; ++r) crol[r] += CSTEP; } +#define WG4_LOADB(c) { const size_t _o = (size_t)(c) * CIB * SCin; _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { const int4 _v = *reinterpret_cast( bbase + _o + boff[t]); br[t][0] = _v.x; br[t][1] = _v.y; br[t][2] = _v.z; br[t][3] = _v.w; } } } +#define WG4_XPOSE(u) { const int _q = lane & 3; uint32_t _x, _y; _x = (_q & 1) ? u[0] : u[1]; _y = (_q & 1) ? u[2] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 1); _y = __shfl_xor_sync(0xffffffffu, _y, 1); if (_q & 1) { u[0] = _x; u[2] = _y; } else { u[1] = _x; u[3] = _y; } _x = (_q & 2) ? u[0] : u[2]; _y = (_q & 2) ? u[1] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 2); _y = __shfl_xor_sync(0xffffffffu, _y, 2); if (_q & 2) { u[0] = _x; u[1] = _y; } else { u[2] = _x; u[3] = _y; } } +#define WG4_STOREB(BD) { _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { uint32_t _u[4]; _u[0] = br[t][0]; _u[1] = br[t][1]; _u[2] = br[t][2]; _u[3] = br[t][3]; WG4_XPOSE(_u) const uint32_t _a = static_cast( __cvta_generic_to_shared((BD) + bdst[t])); asm volatile( "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 " "[%0], {%1,%2,%3,%4};\n" ::"r"(_a), "r"(_u[0]), "r"(_u[1]), "r"(_u[2]), "r"(_u[3])); } } } +#define WG4_STOREB_R(SLOT) { _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { uint32_t _u[4]; _u[0] = br[t][0]; _u[1] = br[t][1]; _u[2] = br[t][2]; _u[3] = br[t][3]; WG4_XPOSE(_u) asm volatile( "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 " "[%0+%5], {%1,%2,%3,%4};\n" ::"r"(bsm0[t]), "r"(_u[0]), "r"(_u[1]), "r"(_u[2]), "r"(_u[3]), "n"((SLOT) * ST_E * 2)); } } } +#define WG4_LOADA_R(BUF) { _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { const int4 _v = *reinterpret_cast(arol + kk * 1024); areg[BUF][kk][0] = _v.x; areg[BUF][kk][1] = _v.y; areg[BUF][kk][2] = _v.z; areg[BUF][kk][3] = _v.w; } arol += KT * 1024; } +#define WG4_MMA_R(BUF,SLOT) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { WGMMA_RS_ANY_T(WG_BN, acc, areg[BUF][kk], dslot[SLOT] + kk, 0); } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } +#define WG4_LOADA(BUF,c) { const __half *_p = apt + (size_t)(c) * (KT * 1024); _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { const int4 _v = *reinterpret_cast(_p + kk * 1024); areg[BUF][kk][0] = _v.x; areg[BUF][kk][1] = _v.y; areg[BUF][kk][2] = _v.z; areg[BUF][kk][3] = _v.w; } } +#define WG4_KEEPA(BUF) { _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) asm volatile("" : "+r"(areg[BUF][kk][0]), "+r"(areg[BUF][kk][1]), "+r"(areg[BUF][kk][2]), "+r"(areg[BUF][kk][3])); } +#define WG4_MMA(BUF,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t db0 = wg_desc((BD) + ng * WG_BN * 8, LBO_B * 2u, 128u); _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { WGMMA_RS_ANY_T(WG_BN, acc, areg[BUF][kk], db0 + kk, 0); } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + if constexpr (PDL && LATE_PDL) PDL_ENTER(); + constexpr int WGD = 1; + if constexpr (BC8) { +#pragma unroll 1 + for (int st = 0; st < CPA_STAGES; ++st) { + if (st < NCH) WG4_CPB(st, smem + st * ST_E) + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } + WG4_LOADA(0, 0) + if (ABUF > 1 && 1 < NCH) WG4_LOADA(1, 1) + if (ABUF > 2 && 2 < NCH) WG4_LOADA(2, 2) + } else { + WG4_LOADB(0) + WG4_LOADA(0, 0) + if (ABUF > 1 && 1 < NCH) WG4_LOADA(1, 1) + if (ABUF > 2 && 2 < NCH) WG4_LOADA(2, 2) + WG4_STOREB(smem) + if (1 < NCH) WG4_LOADB(1) + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + } +#define WG4_BODY(BUF,i,SLOT) { if constexpr (BC8) { asm volatile("cp.async.wait_group %0;\n" : : "n"(CPA_STAGES - 3) : "memory"); __syncthreads(); if ((i) >= 2) { if ((i) + 2 < NCH) { WG4_CPB_R(((SLOT) + 2) % CPA_STAGES) } asm volatile("cp.async.commit_group;\n" ::: "memory"); } } else { if (i) __syncthreads(); } WG4_MMA_R(BUF, SLOT) asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(WGD) : "memory"); if constexpr (NSPLIT > 1) __syncthreads(); if constexpr (!BC8) { const int cb = (i) + 1; if (cb < NCH) { WG4_STOREB_R(((SLOT) + 1) % BSTAGES) asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); if (cb + 1 < NCH) WG4_LOADB(cb + 1) } } if ((i) >= 1) { const int nb2 = ((BUF) + ABUF - 1) % ABUF; if ((i) + ABUF - 1 < NCH) { WG4_KEEPA(nb2) WG4_LOADA_R(nb2) } } } + for (int i = 0; i < NCH; i += 4) { + WG4_BODY(0, i, 0) + if (i + 1 < NCH) WG4_BODY(1, i + 1, 1) + if (i + 2 < NCH) WG4_BODY(0, i + 2, 2) + if (i + 3 < NCH) WG4_BODY(1, i + 3, 3) + } + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + if constexpr (BC8) + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); +#undef WG4_CPB +#undef WG4_CPB_R +#undef WG4_STOREB_R +#undef WG4_LOADA_R +#undef WG4_MMA_R +#undef WG4_LOADB +#undef WG4_XPOSE +#undef WG4_STOREB +#undef WG4_LOADA +#undef WG4_KEEPA +#undef WG4_MMA +#undef WG4_BODY +#define TL4_SW(c,r) ((c) * BM + (((r) & ~7) ^ (((c) & 7) << 3)) + ((r) & 7)) + __syncthreads(); + const int cb2 = (lane & 3) * 2; + const bool snake = PREACT && (m0 + er0) < half; + const float sn_a0 = snake ? row_alpha[er0] : 0.f; + const float sn_r0 = snake ? row_ri[er0] : 0.f; + const float sn_a1 = snake ? row_alpha[er0 + 8] : 0.f; + const float sn_r1 = snake ? row_ri[er0 + 8] : 0.f; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + const int row = er0 + h * 8; + if constexpr (PREACT) { + if (snake) { + const float a = h ? sn_a1 : sn_a0; + const float ri = h ? sn_r1 : sn_r0; + const float s0 = __sinf(a * v.x); + const float s1 = __sinf(a * v.y); + v.x = __fmaf_rn(s0 * s0, ri, v.x); + v.y = __fmaf_rn(s1 * s1, ri, v.y); + } else { + v.x = v.x > 0.f ? v.x : 0.01f * v.x; + v.y = v.y > 0.f ? v.y : 0.01f * v.y; + } + if constexpr (OC8) { + const int c0 = ng * WG_BN + j * 8 + cb2; + Hs[TL4_SW(c0, row)] = __float2half(v.x); + Hs[TL4_SW(c0 + 1, row)] = __float2half(v.y); + } else { + *reinterpret_cast<__half2 *>( + Hs + row * WG_CLD + ng * WG_BN + j * 8 + cb2) = + __floats2half2_rn(v.x, v.y); + } + } else { + *reinterpret_cast( + Cs + row * WG_CLD + ng * WG_BN + j * 8 + cb2) = v; + } + } + } + __syncthreads(); + if constexpr (PREACT) { + if constexpr (OC8) { + constexpr int GRP = BM / 8; + constexpr int TOTG = GRP * BNT; + constexpr int GITER = (TOTG + NT - 1) / NT; + const size_t SC8o = (size_t)SCout * 8; +#pragma unroll + for (int t = 0; t < GITER; ++t) { + const int idx = tid + t * NT; + if (TOTG % NT != 0 && idx >= TOTG) continue; + const int g = idx / BNT; + const int j = idx - g * BNT; + if (!colok[j]) continue; + const int m0g = g * 8; + if (m0 + m0g >= Cout) continue; + *reinterpret_cast( + out + (size_t)((m0 + m0g) >> 3) * SC8o + + (size_t)(nb + j + H) * 8) = + *reinterpret_cast(Hs + TL4_SW(j, m0g)); + } + } else { + static_assert((BNT & 3) == 0, "four-column preactivated epilogue"); + constexpr int NC = BNT / 4; + constexpr int TOT = BM * NC; + for (int q = tid; q < TOT; q += NT) { + const int m = q / NC; + const int lc = (q - m * NC) * 4; + const int co = m0 + m; + const bool any4 = (colok[lc] | colok[lc + 1] | + colok[lc + 2] | colok[lc + 3]) != 0; + if (co >= Cout || !any4) continue; + const __half2 *sp = reinterpret_cast( + Hs + m * WG_CLD + lc); + const __half2 v0 = sp[0], v1 = sp[1]; + __half *op = out + (size_t)co * SCout + nb + lc + H; + const bool ok4 = (colok[lc] & colok[lc + 1] & + colok[lc + 2] & colok[lc + 3]) != 0; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = v0; + reinterpret_cast<__half2 *>(op)[1] = v1; + } else { + if (colok[lc + 0]) op[0] = v0.x; + if (colok[lc + 1]) op[1] = v0.y; + if (colok[lc + 2]) op[2] = v1.x; + if (colok[lc + 3]) op[3] = v1.y; + } + } + } + } else { + constexpr bool COMPACT_EPI = + (WG_BN == 136 || WG_BN == 160 || WG_BN == 208); + if constexpr (COMPACT_EPI) { + static_assert((BNT & 3) == 0, "4-column epilogue groups"); + constexpr int NC = BNT / 4; + constexpr int TOT = BM * NC; + constexpr int MITER = (TOT + NT - 1) / NT; + constexpr int UNR = (MITER >= 8) ? 8 : ((MITER >= 4) ? 4 : + ((MITER >= 2) ? 2 : 1)); +#pragma unroll 1 + for (int t0 = 0; t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int q = tid + (t0 + u) * NT; + if (q >= TOT) continue; + const int m = q / NC; + const int lc = (q - m * NC) * 4; + const int co = m0 + m; + const bool any4 = (colok[lc] | colok[lc + 1] | + colok[lc + 2] | colok[lc + 3]) != 0; + if (co >= Cout || !any4) continue; + v[u] = *reinterpret_cast(Cs + m * WG_CLD + lc); + if constexpr (!PREACT) { + const int n = nb + lc; + const __half2 *rp = reinterpret_cast( + resid + (size_t)co * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int q = tid + (t0 + u) * NT; + if (q >= TOT) continue; + const int m = q / NC; + const int lc = (q - m * NC) * 4; + const int co = m0 + m; + const bool any4 = (colok[lc] | colok[lc + 1] | + colok[lc + 2] | colok[lc + 3]) != 0; + if (co >= Cout || !any4) continue; + const int n = nb + lc; + float4 vv = v[u]; + if constexpr (!PREACT) { + vv.x += __half2float(r0[u].x); vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (co < half) { + const float a = row_alpha[m], ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + __half *op = out + (size_t)co * SCout + n + H; + const bool ok4 = (colok[lc] & colok[lc + 1] & + colok[lc + 2] & colok[lc + 3]) != 0; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + } else { + constexpr int NC = EPN / 4; + constexpr int MSTEP = NT / NC; + constexpr int MITER = BM / MSTEP; + constexpr int UNR = (MITER % 8 == 0) ? 8 + : ((MITER % 4 == 0) ? 4 + : ((MITER % 2 == 0) ? 2 : 1)); + static_assert(MITER % UNR == 0, "row batch must divide row count"); + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && + colok[lc + 2] && colok[lc + 3]; + const bool live4 = colok[lc] || colok[lc + 1] || + colok[lc + 2] || colok[lc + 3]; + const int n = nb + lc; +#pragma unroll 1 + for (int t0 = 0; live4 && t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + v[u] = *reinterpret_cast( + Cs + m * WG_CLD + lc); + if constexpr (!PREACT) { + const int co = m0 + m; + const int cc = co < Cout ? co : 0; + const __half2 *rp = reinterpret_cast( + resid + (size_t)cc * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + const int co = m0 + m; + if (co >= Cout) continue; + float4 vv = v[u]; + if constexpr (!PREACT) { + vv.x += __half2float(r0[u].x); + vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); + vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (co < half) { + const float a = row_alpha[m], ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + __half *op = out + (size_t)co * SCout + n + H; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + } + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B; + for (int idx = bxy * NT + tid; idx < total; idx += nbx * NT) { + const int b = idx / Cin, ci = idx - b * Cin; + const int x = b * L + native_true_tokens(b,T); + __half *const dst = ns + ((size_t)b * Cin + ci) * H; + if constexpr (BC8) { + const __half *const src = + in + (size_t)(ci >> 3) * ((size_t)SCin * 8) + + (size_t)x * 8 + (ci & 7); +#pragma unroll + for (int j = 0; j < H; ++j) dst[j] = src[j * 8]; + } else { + const __half *const src = in + (size_t)ci * SCin + x; +#pragma unroll + for (int j = 0; j < H; j += 2) + *reinterpret_cast(dst + j) = + *reinterpret_cast(src + j); + } + } + } + if (pad_src != nullptr) { + const int total = Cout * B; + for (int idx = bxy * NT + tid; idx < total; idx += nbx * NT) { + const int b = idx / Cout, co = idx - b * Cout; + const int x = b * L + pad_off; + const __half *const src = + pad_src + ((size_t)b * Cout + co) * pad_len; + if constexpr (OC8) { + __half *const dst = + out + (size_t)(co >> 3) * ((size_t)SCout * 8) + + (size_t)x * 8 + (co & 7); + for (int j = 0; j < pad_len; ++j) dst[j * 8] = src[j]; + } else { + __half *const dst = out + (size_t)co * SCout + x; + for (int j = 0; j < pad_len; ++j) dst[j] = src[j]; + } + } + } + } +} diff --git a/deploy/nemotron_native/kernels/b512/native/codec/wg5_step_widths.inc b/deploy/nemotron_native/kernels/b512/native/codec/wg5_step_widths.inc new file mode 100644 index 0000000..fe920a3 --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/native/codec/wg5_step_widths.inc @@ -0,0 +1,528 @@ +template +__device__ __noinline__ +void wg5_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int KT = 3; + constexpr int BM = 64 * MSPLIT; + constexpr int BNT = WG_BN; + constexpr int NTS = 128 * MSPLIT; + constexpr int WG_NT = NTS * KSPLIT; + constexpr int NCRAW = BNT / 4; + constexpr int EPN = + (WG_NT % NCRAW == 0) ? BNT : (((BNT + 127) / 128) * 128); + constexpr int WG_CLD_R = (OC8 ? WG_BN : EPN) + 8; + constexpr int WG_CLD = (WG_CLD_R % 16 == 0) ? WG_CLD_R + 4 : WG_CLD_R; + constexpr bool C8_XPOSE_MAILBOX = + OC8 && CIB == 16 && MSPLIT == 1 && KSPLIT == 4 && + ((CIN == 768 && WG_BN == 96) || + (CIN == 1536 && WG_BN == 104)); + constexpr int CG = CIB / 8; + constexpr int CSUB = CIB / 16; + constexpr int AK = KT * CSUB; + constexpr int NBRP = BC8 ? (((BNT + KT - 1) + 7) & ~7) + : (((BNT + KT - 1) + 31) & ~31); + constexpr int UCOL = NBRP / 32; + constexpr int NUNIT = CG * UCOL; + constexpr int NWARP = NTS / 32; + constexpr int BITER = (NUNIT + NWARP - 1) / NWARP; + constexpr int B_E = CIB * NBRP; + constexpr int ST_E = SB * B_E; + constexpr int LBO_B = NBRP * 8; + constexpr int ABUF = 3; + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *smem = reinterpret_cast<__half *>(arena); + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int kgrp = tid / NTS; + const int stid = tid - kgrp * NTS; + const int mg = stid >> 7; + const int warp = (stid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * BM; + const int nb = blockIdx.y * BNT; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + constexpr int NCH = (CIN + CIB - 1) / CIB; + static_assert(NCH % KSPLIT == 0, "wg5 K-split must divide the chunk count"); + bool active_tile = false; + for (int row = nb / L; row <= min(N - 1, nb + BNT - 1) / L; ++row) { + if (native_meta()->lengths[row] > 0) { + active_tile = true; + break; + } + } + if (active_tile) { + if (tid < BM) { + const int co = m0 + tid; + if (co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + for (int i = tid; i < EPN; i += WG_NT) { + const int n = nb + i; + colok[i] = (i < BNT && n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mg * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (kgrp == 0 && eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (kgrp == 0 && eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + const int mtile = blockIdx.x * MSPLIT + mg; + const __half *const apt = w + ((size_t)mtile * NCH * AK) * 1024 + + (size_t)(stid & 127) * 8; + uint32_t areg[ABUF][AK][4]; + const int wi = stid >> 5; + int bdst[BITER], boff[BITER]; + bool blive[BITER]; + const __half *const bbase = in + (size_t)(lane >> 2) * SCin + nb + + 8 * (lane & 3); +#pragma unroll + for (int t = 0; t < BITER; ++t) { + const int u = wi + t * NWARP; + const int gc = (u < NUNIT) ? u / UCOL : 0; + const int cu = (u < NUNIT) ? u - gc * UCOL : 0; + blive[t] = u < NUNIT; + bdst[t] = gc * LBO_B + (cu * 32 + lane) * 8; + boff[t] = gc * 8 * SCin + cu * 32; + } + const size_t SC8i = (size_t)SCin * 8; + constexpr int BQ = B_E / 8; + constexpr int BCREG = (BQ + NTS - 1) / NTS; + int cdst[BCREG]; + const __half *csrc[BCREG]; + bool clive[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) { + const int q = stid + r * NTS; + const int gc = q / NBRP, rem = q - gc * NBRP; + clive[r] = ((BQ % NTS == 0) || q < BQ) && gc < CG; + cdst[r] = clive[r] ? gc * LBO_B + rem * 8 : 0; + csrc[r] = in + (clive[r] ? (size_t)gc * SC8i + + (size_t)nb * 8 + rem * 8 : 0); + } + uint32_t br[BITER][4]; + __half *const kbB = smem + kgrp * ST_E; + const unsigned barid = 1u + kgrp + (unsigned)(KSPLIT * native_subid()); + static_assert((B_E & 7) == 0, "descriptor slot alignment"); + const uint64_t dbase = wg_desc(kbB, LBO_B * 2u, 128u); + uint32_t csm0[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) + csm0[r] = static_cast( + __cvta_generic_to_shared(kbB + cdst[r])); + const size_t CSTEP = (size_t)KSPLIT * CG * SC8i; + constexpr int ASTEP = KSPLIT * AK * 1024; + const __half *crol[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) + crol[r] = csrc[r] + (size_t)(ABUF * KSPLIT + kgrp) * CG * SC8i; + const __half *arol = apt + (size_t)(ABUF * KSPLIT + kgrp) * (AK * 1024); +#define WG5_BAR() asm volatile("bar.sync %0, %1;\n" ::"r"(barid), "r"(NTS)) +#define WG5_CPB(c,BD) { const size_t _o = (size_t)(c) * CG * SC8i; _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16((BD) + cdst[r], csrc[r] + _o, true); } +#define WG5_LOADA(BUF,c) { const __half *_p = apt + (size_t)(c) * (AK * 1024); _Pragma("unroll") for (int kk = 0; kk < AK; ++kk) { const int4 _v = *reinterpret_cast(_p + kk * 1024); areg[BUF][kk][0] = _v.x; areg[BUF][kk][1] = _v.y; areg[BUF][kk][2] = _v.z; areg[BUF][kk][3] = _v.w; } } +#define WG5_KEEPA(BUF) { _Pragma("unroll") for (int kk = 0; kk < AK; ++kk) asm volatile("" : "+r"(areg[BUF][kk][0]), "+r"(areg[BUF][kk][1]), "+r"(areg[BUF][kk][2]), "+r"(areg[BUF][kk][3])); } +#define WG5_LOADB(c) { const size_t _o = (size_t)(c) * CIB * SCin; _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { const int4 _v = *reinterpret_cast( bbase + _o + boff[t]); br[t][0] = _v.x; br[t][1] = _v.y; br[t][2] = _v.z; br[t][3] = _v.w; } } } +#define WG5_BXPOSE(u) { const int _q = lane & 3; uint32_t _x, _y; _x = (_q & 1) ? u[0] : u[1]; _y = (_q & 1) ? u[2] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 1); _y = __shfl_xor_sync(0xffffffffu, _y, 1); if (_q & 1) { u[0] = _x; u[2] = _y; } else { u[1] = _x; u[3] = _y; } _x = (_q & 2) ? u[0] : u[2]; _y = (_q & 2) ? u[1] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 2); _y = __shfl_xor_sync(0xffffffffu, _y, 2); if (_q & 2) { u[0] = _x; u[1] = _y; } else { u[2] = _x; u[3] = _y; } } +#define WG5_STOREB(BD) { _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { uint32_t _u[4]; _u[0] = br[t][0]; _u[1] = br[t][1]; _u[2] = br[t][2]; _u[3] = br[t][3]; WG5_BXPOSE(_u) const uint32_t _a = static_cast( __cvta_generic_to_shared((BD) + bdst[t])); asm volatile( "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 " "[%0], {%1,%2,%3,%4};\n" ::"r"(_a), "r"(_u[0]), "r"(_u[1]), "r"(_u[2]), "r"(_u[3])); } } } +#define WG5_MMA(BUF) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); _Pragma("unroll") for (int h = 0; h < CSUB; ++h) { _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { WGMMA_RS_ANY_T(WG_BN, acc, areg[BUF][h * KT + kk], dbase + (BUF) * (B_E / 8) + kk + h * (LBO_B / 4), 0); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } +#define WG5_CPB_R(SLOT) { _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16o<(SLOT) * B_E * 2>(csm0[r], crol[r], true); } +#define WG5_LOADA_R(BUF) { _Pragma("unroll") for (int kk = 0; kk < AK; ++kk) { const int4 _v = *reinterpret_cast(arol + kk * 1024); areg[BUF][kk][0] = _v.x; areg[BUF][kk][1] = _v.y; areg[BUF][kk][2] = _v.z; areg[BUF][kk][3] = _v.w; } } +#define WG5_ROLL() { arol += ASTEP; if constexpr (BC8) { _Pragma("unroll") for (int r = 0; r < BCREG; ++r) crol[r] += CSTEP; } } + if constexpr (PDL) { + (void)0; + if (tid == 0) + (void)0; + } + static_assert(SB == 3 && WGD == 1 && ABUF == 3, "wg5 pipeline shape"); + if constexpr (BC8) { +#pragma unroll 1 + for (int s = 0; s < SB; ++s) { + const int c = s * KSPLIT + kgrp; + if (c < NCH) WG5_CPB(c, kbB + s * B_E) + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } + } else { + WG5_LOADB(kgrp) +#pragma unroll 1 + for (int s = 0; s < SB; ++s) { + const int c = s * KSPLIT + kgrp; + if (c < NCH) WG5_STOREB(kbB + s * B_E) + if (c + KSPLIT < NCH) WG5_LOADB(c + KSPLIT) + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + } + WG5_LOADA(0, kgrp) + if (KSPLIT + kgrp < NCH) WG5_LOADA(1, KSPLIT + kgrp) + if (2 * KSPLIT + kgrp < NCH) WG5_LOADA(2, 2 * KSPLIT + kgrp) + constexpr int NIT = NCH / KSPLIT; + static_assert(NIT % ABUF == 0, "wg5 mainloop unroll shape"); +#define WG5_BODY(BUF,i) { if constexpr (BC8) asm volatile("cp.async.wait_group %0;\n" : : "n"(SB - 2) : "memory"); WG5_BAR(); WG5_MMA(BUF) asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(WGD) : "memory"); if constexpr (NTS > 128) WG5_BAR(); if ((i) >= 1) { const int nb2 = ((BUF) + 2) % ABUF; if (((i) + 2) * KSPLIT < NCH) { if constexpr (BC8) { WG5_CPB_R(nb2) } else { WG5_STOREB(kbB + nb2 * B_E) asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); if (((i) + 3) * KSPLIT < NCH) WG5_LOADB(((i) + 3) * KSPLIT + kgrp) } WG5_KEEPA(nb2) WG5_LOADA_R(nb2) } WG5_ROLL() if constexpr (BC8) asm volatile("cp.async.commit_group;\n" ::: "memory"); } } +#pragma unroll + for (int i = 0; i < NIT; i += 3) { + WG5_BODY(0, i) + WG5_BODY(1, i + 1) + WG5_BODY(2, i + 2) + } + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + if constexpr (BC8) + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); +#undef WG5_BXPOSE +#undef WG5_CPB +#undef WG5_CPB_R +#undef WG5_LOADA_R +#undef WG5_ROLL +#undef WG5_LOADB +#undef WG5_STOREB +#undef WG5_MMA +#undef WG5_LOADA +#undef WG5_KEEPA +#undef WG5_BODY +#undef WG5_BAR + __syncthreads(); + const int cb2 = (lane & 3) * 2; + if constexpr (C8_XPOSE_MAILBOX) { + __half *Ck = smem + kgrp * BM * WG_CLD; + static_assert(WG_BN % 32 == 0 || WG_BN % 32 == 8, + "x4 mailbox plus optional x1 eight-column tail"); +#pragma unroll + for (int jb = 0; jb < WG_BN / 32; ++jb) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + const __half2 p0 = __floats2half2_rn( + acc[(jb * 4 + 0) * 4 + h * 2], + acc[(jb * 4 + 0) * 4 + h * 2 + 1]); + const __half2 p1 = __floats2half2_rn( + acc[(jb * 4 + 1) * 4 + h * 2], + acc[(jb * 4 + 1) * 4 + h * 2 + 1]); + const __half2 p2 = __floats2half2_rn( + acc[(jb * 4 + 2) * 4 + h * 2], + acc[(jb * 4 + 2) * 4 + h * 2 + 1]); + const __half2 p3 = __floats2half2_rn( + acc[(jb * 4 + 3) * 4 + h * 2], + acc[(jb * 4 + 3) * 4 + h * 2 + 1]); + const uint32_t r0 = + *reinterpret_cast(&p0); + const uint32_t r1 = + *reinterpret_cast(&p1); + const uint32_t r2 = + *reinterpret_cast(&p2); + const uint32_t r3 = + *reinterpret_cast(&p3); + const int row_group = mg * 8 + warp * 2 + h; + const int col = jb * 32 + (lane >> 3) * 8 + (lane & 7); + const uint32_t a = static_cast( + __cvta_generic_to_shared( + Ck + ((row_group * WG_CLD + col) * 8))); + asm volatile( + "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 " + "[%0], {%1,%2,%3,%4};\n" + : + : "r"(a), "r"(r0), "r"(r1), "r"(r2), "r"(r3) + : "memory"); + } + } + if constexpr ((WG_BN % 32) == 8) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + constexpr int jt = WG_BN / 8 - 1; + const __half2 p = __floats2half2_rn( + acc[jt * 4 + h * 2], + acc[jt * 4 + h * 2 + 1]); + const uint32_t r = + *reinterpret_cast(&p); + const int row_group = mg * 8 + warp * 2 + h; + const int col = (WG_BN & ~31) + (lane & 7); + const uint32_t a = static_cast( + __cvta_generic_to_shared( + Ck + ((row_group * WG_CLD + col) * 8))); + asm volatile( + "stmatrix.sync.aligned.m8n8.x1.trans.shared.b16 " + "[%0], {%1};\n" + : + : "r"(a), "r"(r) + : "memory"); + } + } + } else if constexpr (OC8) { + __half *Ck = smem + kgrp * BM * WG_CLD; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + *reinterpret_cast<__half2 *>(Ck + (er0 + h * 8) * WG_CLD + + j * 8 + cb2) = + __floats2half2_rn(acc[j * 4 + h * 2], + acc[j * 4 + h * 2 + 1]); + } + } + } else { + float *Ck = Cs + kgrp * BM * WG_CLD; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], + acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Ck + (er0 + h * 8) * WG_CLD + + j * 8 + cb2) = v; + } + } + } + __syncthreads(); + if constexpr (OC8) { + constexpr int GRP = BM / 8; + constexpr int TOTG = GRP * WG_BN; + constexpr int GITER = (TOTG + WG_NT - 1) / WG_NT; + const size_t SC8o = (size_t)SCout * 8; +#pragma unroll + for (int t = 0; t < GITER; ++t) { + const int idx = tid + t * WG_NT; + if (TOTG % WG_NT != 0 && idx >= TOTG) continue; + const int j = idx - (idx / WG_BN) * WG_BN, g = idx / WG_BN; + if (!colok[j]) continue; + const int m0g = g * 8; + if (m0 + m0g >= Cout) continue; + const int n = nb + j; + __half rr[8]; + if constexpr (CIN == 1536) { + if constexpr (RC8) { + *reinterpret_cast(rr) = + *reinterpret_cast( + resid + (size_t)((m0 + m0g) >> 3) * SC8o + + (size_t)(n + H) * 8); + } else { +#pragma unroll + for (int h = 0; h < 8; ++h) + rr[h] = resid[(size_t)(m0 + m0g + h) * SCout + n + H]; + } + } + uint32_t pk[4]; + const bool snake = (m0 + m0g) < half; + if constexpr (C8_XPOSE_MAILBOX) { + const __half *const Ccol = smem + (g * WG_CLD + j) * 8; +#pragma unroll + for (int h = 0; h < 8; h += 4) { + const int2 z0 = *reinterpret_cast(Ccol + h); + const __half2 z00 = + *reinterpret_cast(&z0.x); + const __half2 z01 = + *reinterpret_cast(&z0.y); + float a0 = __half2float(z00.x); + float a1 = __half2float(z00.y); + float a2 = __half2float(z01.x); + float a3 = __half2float(z01.y); +#pragma unroll + for (int kg = 1; kg < KSPLIT; ++kg) { + const int2 z = *reinterpret_cast( + Ccol + kg * BM * WG_CLD + h); + const __half2 z10 = + *reinterpret_cast(&z.x); + const __half2 z11 = + *reinterpret_cast(&z.y); + a0 += __half2float(z10.x); + a1 += __half2float(z10.y); + a2 += __half2float(z11.x); + a3 += __half2float(z11.y); + } + if constexpr (CIN == 1536) { + a0 += __half2float(rr[h + 0]); + a1 += __half2float(rr[h + 1]); + a2 += __half2float(rr[h + 2]); + a3 += __half2float(rr[h + 3]); + } + const int m = m0g + h; + if (snake) { + const float sn0 = __sinf(row_alpha[m + 0] * a0); + const float sn1 = __sinf(row_alpha[m + 1] * a1); + const float sn2 = __sinf(row_alpha[m + 2] * a2); + const float sn3 = __sinf(row_alpha[m + 3] * a3); + a0 = __fmaf_rn(sn0 * sn0, row_ri[m + 0], a0); + a1 = __fmaf_rn(sn1 * sn1, row_ri[m + 1], a1); + a2 = __fmaf_rn(sn2 * sn2, row_ri[m + 2], a2); + a3 = __fmaf_rn(sn3 * sn3, row_ri[m + 3], a3); + } else { + a0 = a0 > 0.f ? a0 : 0.01f * a0; + a1 = a1 > 0.f ? a1 : 0.01f * a1; + a2 = a2 > 0.f ? a2 : 0.01f * a2; + a3 = a3 > 0.f ? a3 : 0.01f * a3; + } + const __half2 p0 = __floats2half2_rn(a0, a1); + const __half2 p1 = __floats2half2_rn(a2, a3); + pk[h >> 1] = + *reinterpret_cast(&p0); + pk[(h >> 1) + 1] = + *reinterpret_cast(&p1); + } + } else { +#pragma unroll + for (int h = 0; h < 8; h += 2) { + float f[2]; +#pragma unroll + for (int e = 0; e < 2; ++e) { + const int m = m0g + h + e; + float a = __half2float(smem[m * WG_CLD + j]); +#pragma unroll + for (int kg = 1; kg < KSPLIT; ++kg) + a += __half2float( + smem[kg * BM * WG_CLD + m * WG_CLD + j]); + if constexpr (CIN == 1536) + a += __half2float(rr[h + e]); + if (snake) { + const float sn = __sinf(row_alpha[m] * a); + a = __fmaf_rn(sn * sn, row_ri[m], a); + } else { + a = a > 0.f ? a : 0.01f * a; + } + f[e] = a; + } + const __half2 p = __floats2half2_rn(f[0], f[1]); + pk[h >> 1] = *reinterpret_cast(&p); + } + } + *reinterpret_cast( + out + (size_t)((m0 + m0g) >> 3) * SC8o + + (size_t)(n + H) * 8) = make_int4(pk[0], pk[1], pk[2], pk[3]); + } + } else { + constexpr int NC = EPN / 4; + constexpr int MSTEP = WG_NT / NC; + constexpr int MITER = BM / MSTEP; + constexpr int UNR = (MITER % 4 == 0) ? 4 : ((MITER % 2 == 0) ? 2 : 1); + static_assert(MITER % UNR == 0, "row batch must divide the row count"); + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + const bool live4 = colok[lc] || colok[lc + 1] || colok[lc + 2] || + colok[lc + 3]; + const int n = nb + lc; +#pragma unroll 1 + for (int t0 = 0; live4 && t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + v[u] = *reinterpret_cast(Cs + m * WG_CLD + lc); + if constexpr (CIN == 1536) { + const int co = m0 + m; + const int cc = co < Cout ? co : 0; + if constexpr (RC8) { + const __half *rp = + resid + (size_t)(cc >> 3) * ((size_t)SCout * 8) + + (size_t)(n + H) * 8 + (cc & 7); + r0[u].x = rp[0]; r0[u].y = rp[8]; + r1[u].x = rp[16]; r1[u].y = rp[24]; + } else { + const __half2 *rp = reinterpret_cast( + resid + (size_t)cc * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + const int co = m0 + m; + if (co >= Cout) continue; + float4 vv = v[u]; +#pragma unroll + for (int g = 1; g < KSPLIT; ++g) { + float4 p = *reinterpret_cast( + Cs + g * BM * WG_CLD + m * WG_CLD + lc); + vv.x += p.x; vv.y += p.y; vv.z += p.z; vv.w += p.w; + } + if constexpr (CIN == 1536) { + vv.x += __half2float(r0[u].x); vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + __half *op = out + (size_t)co * SCout + n + H; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + static_assert(CIN == 768 || CIN == 1536); + constexpr int COUT = (CIN == 768) ? 1536 : 768; + if (ns != nullptr) { + for (int b = bxy; b < B; b += nbx) { + const int x = b * L + native_true_tokens(b,T); + for (int ci = tid; ci < CIN; ci += WG_NT) { + __half *const dst = ns + ((size_t)b * CIN + ci) * H; + if constexpr (BC8) { + const __half *const src = + in + (size_t)(ci >> 3) * ((size_t)SCin * 8) + + (size_t)x * 8 + (ci & 7); +#pragma unroll + for (int j = 0; j < H; ++j) dst[j] = src[j * 8]; + } else { + const __half *const src = in + (size_t)ci * SCin + x; +#pragma unroll + for (int j = 0; j < H; j += 2) + *reinterpret_cast(dst + j) = + *reinterpret_cast(src + j); + } + } + } + } + if (pad_src != nullptr) { + for (int b = bxy; b < B; b += nbx) { + const int x = b * L + pad_off; + for (int co = tid; co < COUT; co += WG_NT) { + const __half *const src = + pad_src + ((size_t)b * COUT + co) * pad_len; + if constexpr (OC8) { + __half *const dst = + out + (size_t)(co >> 3) * ((size_t)SCout * 8) + + (size_t)x * 8 + (co & 7); + for (int j = 0; j < pad_len; ++j) dst[j * 8] = src[j]; + } else { + __half *const dst = out + (size_t)co * SCout + x; + for (int j = 0; j < pad_len; ++j) dst[j] = src[j]; + } + } + } + } + } +} diff --git a/deploy/nemotron_native/kernels/b512/native/codec/wg_kernel.inc b/deploy/nemotron_native/kernels/b512/native/codec/wg_kernel.inc new file mode 100644 index 0000000..2d5c766 --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/native/codec/wg_kernel.inc @@ -0,0 +1,182 @@ +__device__ __forceinline__ uint64_t wg_desc(const void *p, uint32_t lbo_bytes, + uint32_t sbo_bytes) { + uint32_t a = static_cast(__cvta_generic_to_shared(p)); + uint64_t d = static_cast((a & 0x3FFFFu) >> 4); + d |= static_cast((lbo_bytes & 0x3FFFFu) >> 4) << 16; + d |= static_cast((sbo_bytes & 0x3FFFFu) >> 4) << 32; + return d; +} +#include "wgmma_macro.inc" +constexpr int WG_SLOT = 72; +template +__device__ __noinline__ void wg_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int WG_CLD = WG_BN + 8; + constexpr int WG_NT = 128 * NWG; + constexpr int WG_BM = 64 * NWG; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int SBO_AE = KG * WG_SLOT; + constexpr int LBO_BE = (WG_BN / 8) * WG_SLOT; + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *As0 = reinterpret_cast<__half *>(arena); + __half *Bs0 = As0 + A_E; + __half *As1 = (STAGES == 2) ? Bs0 + B_E : As0; + __half *Bs1 = (STAGES == 2) ? Bs0 + B_E + A_E : Bs0; + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int wgid = tid >> 7; + const int warp = (tid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * WG_BM; + constexpr int WGROWS = 8; + const int n0 = blockIdx.y * WG_BN; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + if (tid < WG_BM) { + const int co = m0 + tid; + if (do_snake && co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + if (tid < WG_BN) { + const int n = n0 + tid; + colok[tid] = (n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = wgid * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + constexpr int ATOT = WG_BM * KG; + constexpr int AREG = (ATOT + WG_NT - 1) / WG_NT; + constexpr int CIB = BK / KT; + constexpr int BTOT = CIB * (WG_BN / 8); + constexpr int BREG = (BTOT + WG_NT - 1) / WG_NT; + int4 ar[AREG]; + int4 br0[BREG], br1[BREG]; + uint32_t edge[BREG]; +#define WG_LOADG(kg0) { const int _ci0 = (kg0) / KT; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; ar[r] = make_int4(0, 0, 0, 0); if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; const int co = m0 + m, kg = (kg0) + v * 8; if (co < Cout && kg < Kg) ar[r] = *reinterpret_cast( w + (size_t)co * Kg + kg); } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; br0[r] = make_int4(0, 0, 0, 0); if constexpr (KT == 3) edge[r] = 0; else br1[r] = make_int4(0, 0, 0, 0); if (BTOT % WG_NT == 0 || idx < BTOT) { const int ci = _ci0 + idx / (WG_BN / 8); const int nc = idx % (WG_BN / 8); if (ci < Cin) { const int4 *src = reinterpret_cast( in + (size_t)ci * SCin + n0 + nc * 8); br0[r] = src[0]; if constexpr (KT == 3) { if (nc == WG_BN / 8 - 1) edge[r] = src[1].x; } else { br1[r] = src[1]; } } } } } + WG_LOADG(0) +#define WG_STS(AD,BD) { _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; *reinterpret_cast( (AD) + (m >> 3) * SBO_AE + v * WG_SLOT + (m & 7) * 8) = ar[r]; } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; if (BTOT % WG_NT == 0 || idx < BTOT) { const int cil = idx / (WG_BN / 8), nc = idx % (WG_BN / 8); unsigned int u[8]; u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; if constexpr (KT == 3) { u[4] = __shfl_down_sync(0xffffffffu, u[0], 1, WG_BN / 8); if (nc == WG_BN / 8 - 1) u[4] = edge[r]; u[5] = u[6] = u[7] = 0; } else { u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; } _Pragma("unroll") for (int k = 0; k < KT; ++k) { const int kh = k >> 1; int4 o; if (k & 1) { o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); } else { o.x = u[kh + 0]; o.y = u[kh + 1]; o.z = u[kh + 2]; o.w = u[kh + 3]; } const int row = cil * KT + k; *reinterpret_cast((BD) + nc * WG_SLOT + (row >> 3) * LBO_BE + (row & 7) * 8) = o; } } } if constexpr (WG_BN == 64) { asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); } } +#define WG_MMA(AD,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t da0 = wg_desc((AD) + wgid * WGROWS * SBO_AE, WG_SLOT * 2, SBO_AE * 2); const uint64_t db0 = wg_desc((BD), LBO_BE * 2, WG_SLOT * 2); _Pragma("unroll") for (int kk = 0; kk < BK; kk += 16) { const uint64_t da = da0 + (kk >> 3) * (WG_SLOT / 8); const uint64_t db = db0 + (kk >> 3) * (LBO_BE / 8); if constexpr (WG_BN == 64) { WGMMA_M64N64K16(acc, da, db); } else { WGMMA_M64N128K16(acc, da, db); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + WG_STS(As0, Bs0) + if (BK < Kg) WG_LOADG(BK) + __syncthreads(); + int cur = 0; + for (int kg0 = 0; kg0 < Kg; kg0 += BK) { + if (cur == 0) { WG_MMA(As0, Bs0) } else { WG_MMA(As1, Bs1) } + if (kg0 + BK < Kg) { + if (STAGES == 2) + asm volatile("wgmma.wait_group.sync.aligned 1;\n" ::: "memory"); + else + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + __syncthreads(); + if (cur == 0) { WG_STS(As1, Bs1) } else { WG_STS(As0, Bs0) } + if (kg0 + 2 * BK < Kg) WG_LOADG(kg0 + 2 * BK) + __syncthreads(); + } else { + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + } + cur ^= 1; + } +#undef WG_LOADG +#undef WG_STS +#undef WG_MMA + __syncthreads(); + const int r0 = er0; + const int cb = (lane & 3) * 2; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Cs + (r0 + h * 8) * WG_CLD + j * 8 + cb) = v; + } + } + __syncthreads(); + constexpr int NC = WG_BN / 4; + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + for (int m = lr; m < WG_BM; m += WG_NT / NC) { + const int co = m0 + m; + if (co >= Cout) continue; + const int n = n0 + lc; + float4 v = *reinterpret_cast(Cs + m * WG_CLD + lc); + __half *op = out + (size_t)co * SCout + n + H; + if (resid != nullptr) { + const __half2 *rp = reinterpret_cast( + resid + (size_t)co * SCout + n + H); + __half2 r0h = rp[0], r1h = rp[1]; + v.x += __half2float(r0h.x); v.y += __half2float(r0h.y); + v.z += __half2float(r1h.x); v.w += __half2float(r1h.y); + } + if (do_snake) { + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + } + if (ok4) { + __half2 o0 = __floats2half2_rn(v.x, v.y); + __half2 o1 = __floats2half2_rn(v.z, v.w); + reinterpret_cast<__half2 *>(op)[0] = o0; + reinterpret_cast<__half2 *>(op)[1] = o1; + } else { + const float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B * H; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int total = Cout * B * pad_len; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} diff --git a/deploy/nemotron_native/kernels/b512/native/codec/wgmma_macro.inc b/deploy/nemotron_native/kernels/b512/native/codec/wgmma_macro.inc new file mode 100644 index 0000000..db624e9 --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/native/codec/wgmma_macro.inc @@ -0,0 +1,102 @@ +#define WGMMA_M64N8K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " "{" "%0,%1,%2,%3}" ", %4, %5, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N8K16(D,DA,DB) WGMMA_M64N8K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N8K16_ZERO_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " "{" "%0,%1,%2,%3}" ", %4, %5, p, 1, 1, " #TA ", " #TB ";\n}\n" : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) : "l"(DA), "l"(DB), "r"(0)) +#define WGMMA_M64N16K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %10, 0;\n" "wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7}" ", %8, %9, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N16K16(D,DA,DB) WGMMA_M64N16K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N24K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %14, 0;\n" "wgmma.mma_async.sync.aligned.m64n24k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11}" ", %12, %13, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N24K16(D,DA,DB) WGMMA_M64N24K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N32K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %18, 0;\n" "wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15}" ", %16, %17, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N32K16(D,DA,DB) WGMMA_M64N32K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N40K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %22, 0;\n" "wgmma.mma_async.sync.aligned.m64n40k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19}" ", %20, %21, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N40K16(D,DA,DB) WGMMA_M64N40K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N48K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %26, 0;\n" "wgmma.mma_async.sync.aligned.m64n48k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N48K16(D,DA,DB) WGMMA_M64N48K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N56K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %30, 0;\n" "wgmma.mma_async.sync.aligned.m64n56k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N56K16(D,DA,DB) WGMMA_M64N56K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N64K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %34, 0;\n" "wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31}" ", %32, %33, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N64K16(D,DA,DB) WGMMA_M64N64K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N72K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %38, 0;\n" "wgmma.mma_async.sync.aligned.m64n72k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35}" ", %36, %37, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N72K16(D,DA,DB) WGMMA_M64N72K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N80K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %42, 0;\n" "wgmma.mma_async.sync.aligned.m64n80k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39}" ", %40, %41, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N80K16(D,DA,DB) WGMMA_M64N80K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N88K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %46, 0;\n" "wgmma.mma_async.sync.aligned.m64n88k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43}" ", %44, %45, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N88K16(D,DA,DB) WGMMA_M64N88K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N96K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %50, 0;\n" "wgmma.mma_async.sync.aligned.m64n96k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47}" ", %48, %49, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N96K16(D,DA,DB) WGMMA_M64N96K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N104K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %54, 0;\n" "wgmma.mma_async.sync.aligned.m64n104k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51}" ", %52, %53, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N104K16(D,DA,DB) WGMMA_M64N104K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N112K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %58, 0;\n" "wgmma.mma_async.sync.aligned.m64n112k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55}" ", %56, %57, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N112K16(D,DA,DB) WGMMA_M64N112K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N120K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %62, 0;\n" "wgmma.mma_async.sync.aligned.m64n120k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59}" ", %60, %61, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N120K16(D,DA,DB) WGMMA_M64N120K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N128K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %66, 0;\n" "wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63}" ", %64, %65, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N128K16(D,DA,DB) WGMMA_M64N128K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N136K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %70, 0;\n" "wgmma.mma_async.sync.aligned.m64n136k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67}" ", %68, %69, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N136K16(D,DA,DB) WGMMA_M64N136K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N144K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %74, 0;\n" "wgmma.mma_async.sync.aligned.m64n144k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71}" ", %72, %73, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N144K16(D,DA,DB) WGMMA_M64N144K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N152K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %78, 0;\n" "wgmma.mma_async.sync.aligned.m64n152k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75}" ", %76, %77, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N152K16(D,DA,DB) WGMMA_M64N152K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N160K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %82, 0;\n" "wgmma.mma_async.sync.aligned.m64n160k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79}" ", %80, %81, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N160K16(D,DA,DB) WGMMA_M64N160K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N168K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %86, 0;\n" "wgmma.mma_async.sync.aligned.m64n168k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83}" ", %84, %85, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N168K16(D,DA,DB) WGMMA_M64N168K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N176K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %90, 0;\n" "wgmma.mma_async.sync.aligned.m64n176k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87}" ", %88, %89, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N176K16(D,DA,DB) WGMMA_M64N176K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N184K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %94, 0;\n" "wgmma.mma_async.sync.aligned.m64n184k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91}" ", %92, %93, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N184K16(D,DA,DB) WGMMA_M64N184K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N192K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %98, 0;\n" "wgmma.mma_async.sync.aligned.m64n192k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95}" ", %96, %97, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N192K16(D,DA,DB) WGMMA_M64N192K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N200K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %102, 0;\n" "wgmma.mma_async.sync.aligned.m64n200k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99}" ", %100, %101, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N200K16(D,DA,DB) WGMMA_M64N200K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N208K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %106, 0;\n" "wgmma.mma_async.sync.aligned.m64n208k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103}" ", %104, %105, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N208K16(D,DA,DB) WGMMA_M64N208K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N216K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %110, 0;\n" "wgmma.mma_async.sync.aligned.m64n216k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107}" ", %108, %109, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N216K16(D,DA,DB) WGMMA_M64N216K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N224K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %114, 0;\n" "wgmma.mma_async.sync.aligned.m64n224k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111}" ", %112, %113, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N224K16(D,DA,DB) WGMMA_M64N224K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N232K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %118, 0;\n" "wgmma.mma_async.sync.aligned.m64n232k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115}" ", %116, %117, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N232K16(D,DA,DB) WGMMA_M64N232K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N240K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %122, 0;\n" "wgmma.mma_async.sync.aligned.m64n240k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119}" ", %120, %121, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N240K16(D,DA,DB) WGMMA_M64N240K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N248K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %126, 0;\n" "wgmma.mma_async.sync.aligned.m64n248k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123}" ", %124, %125, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N248K16(D,DA,DB) WGMMA_M64N248K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N256K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %130, 0;\n" "wgmma.mma_async.sync.aligned.m64n256k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123,%124,%125,%126,%127}" ", %128, %129, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]), "+f"(D[124]), "+f"(D[125]), "+f"(D[126]), "+f"(D[127]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N256K16(D,DA,DB) WGMMA_M64N256K16_T(D, DA, DB, 0, 1) +#define WGMMA_ANY_T(NW,D,DA,DB,TA,TB) do { if constexpr (NW == 8) { WGMMA_M64N8K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 16) { WGMMA_M64N16K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 24) { WGMMA_M64N24K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 32) { WGMMA_M64N32K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 40) { WGMMA_M64N40K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 48) { WGMMA_M64N48K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 56) { WGMMA_M64N56K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 64) { WGMMA_M64N64K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 72) { WGMMA_M64N72K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 80) { WGMMA_M64N80K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 88) { WGMMA_M64N88K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 96) { WGMMA_M64N96K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 104) { WGMMA_M64N104K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 112) { WGMMA_M64N112K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 120) { WGMMA_M64N120K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 128) { WGMMA_M64N128K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 136) { WGMMA_M64N136K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 144) { WGMMA_M64N144K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 152) { WGMMA_M64N152K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 160) { WGMMA_M64N160K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 168) { WGMMA_M64N168K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 176) { WGMMA_M64N176K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 184) { WGMMA_M64N184K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 192) { WGMMA_M64N192K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 200) { WGMMA_M64N200K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 208) { WGMMA_M64N208K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 216) { WGMMA_M64N216K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 224) { WGMMA_M64N224K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 232) { WGMMA_M64N232K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 240) { WGMMA_M64N240K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 248) { WGMMA_M64N248K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 256) { WGMMA_M64N256K16_T(D, DA, DB, TA, TB); } } while (0) +#define WGMMA_ANY(NW,D,DA,DB) WGMMA_ANY_T(NW, D, DA, DB, 0, 1) +#define WGMMA_RS_M64N96K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %53, 0;\n" "wgmma.mma_async.sync.aligned.m64n96k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47}" ", {%48,%49,%50,%51}, %52, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N104K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %57, 0;\n" "wgmma.mma_async.sync.aligned.m64n104k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51}" ", {%52,%53,%54,%55}, %56, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N112K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %61, 0;\n" "wgmma.mma_async.sync.aligned.m64n112k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55}" ", {%56,%57,%58,%59}, %60, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N120K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %65, 0;\n" "wgmma.mma_async.sync.aligned.m64n120k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59}" ", {%60,%61,%62,%63}, %64, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N128K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %69, 0;\n" "wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63}" ", {%64,%65,%66,%67}, %68, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_ANY_T(NW,D,A,DB,TB) do { if constexpr (NW == 96) { WGMMA_RS_M64N96K16_T(D, A, DB, TB); } else if constexpr (NW == 104) { WGMMA_RS_M64N104K16_T(D, A, DB, TB); } else if constexpr (NW == 112) { WGMMA_RS_M64N112K16_T(D, A, DB, TB); } else if constexpr (NW == 120) { WGMMA_RS_M64N120K16_T(D, A, DB, TB); } else if constexpr (NW == 128) { WGMMA_RS_M64N128K16_T(D, A, DB, TB); } } while (0) +#define WGMMA_RS_M64N8K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %9, 0;\n" "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " "{" "%0,%1,%2,%3}" ", {%4,%5,%6,%7}, %8, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N16K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %13, 0;\n" "wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7}" ", {%8,%9,%10,%11}, %12, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N24K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %17, 0;\n" "wgmma.mma_async.sync.aligned.m64n24k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11}" ", {%12,%13,%14,%15}, %16, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N32K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %21, 0;\n" "wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15}" ", {%16,%17,%18,%19}, %20, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N40K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %25, 0;\n" "wgmma.mma_async.sync.aligned.m64n40k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N48K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %29, 0;\n" "wgmma.mma_async.sync.aligned.m64n48k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N56K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %33, 0;\n" "wgmma.mma_async.sync.aligned.m64n56k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N64K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %37, 0;\n" "wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31}" ", {%32,%33,%34,%35}, %36, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N72K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %41, 0;\n" "wgmma.mma_async.sync.aligned.m64n72k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35}" ", {%36,%37,%38,%39}, %40, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N80K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %45, 0;\n" "wgmma.mma_async.sync.aligned.m64n80k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39}" ", {%40,%41,%42,%43}, %44, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N88K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %49, 0;\n" "wgmma.mma_async.sync.aligned.m64n88k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43}" ", {%44,%45,%46,%47}, %48, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N136K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %73, 0;\n" "wgmma.mma_async.sync.aligned.m64n136k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67}" ", {%68,%69,%70,%71}, %72, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N144K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %77, 0;\n" "wgmma.mma_async.sync.aligned.m64n144k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71}" ", {%72,%73,%74,%75}, %76, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N152K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %81, 0;\n" "wgmma.mma_async.sync.aligned.m64n152k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75}" ", {%76,%77,%78,%79}, %80, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N160K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %85, 0;\n" "wgmma.mma_async.sync.aligned.m64n160k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79}" ", {%80,%81,%82,%83}, %84, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N168K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %89, 0;\n" "wgmma.mma_async.sync.aligned.m64n168k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83}" ", {%84,%85,%86,%87}, %88, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N176K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %93, 0;\n" "wgmma.mma_async.sync.aligned.m64n176k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87}" ", {%88,%89,%90,%91}, %92, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N184K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %97, 0;\n" "wgmma.mma_async.sync.aligned.m64n184k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91}" ", {%92,%93,%94,%95}, %96, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N192K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %101, 0;\n" "wgmma.mma_async.sync.aligned.m64n192k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95}" ", {%96,%97,%98,%99}, %100, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N200K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %105, 0;\n" "wgmma.mma_async.sync.aligned.m64n200k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99}" ", {%100,%101,%102,%103}, %104, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N208K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %109, 0;\n" "wgmma.mma_async.sync.aligned.m64n208k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103}" ", {%104,%105,%106,%107}, %108, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N216K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %113, 0;\n" "wgmma.mma_async.sync.aligned.m64n216k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107}" ", {%108,%109,%110,%111}, %112, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N224K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %117, 0;\n" "wgmma.mma_async.sync.aligned.m64n224k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111}" ", {%112,%113,%114,%115}, %116, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N232K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %121, 0;\n" "wgmma.mma_async.sync.aligned.m64n232k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115}" ", {%116,%117,%118,%119}, %120, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N240K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %125, 0;\n" "wgmma.mma_async.sync.aligned.m64n240k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119}" ", {%120,%121,%122,%123}, %124, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N248K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %129, 0;\n" "wgmma.mma_async.sync.aligned.m64n248k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123}" ", {%124,%125,%126,%127}, %128, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N256K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %133, 0;\n" "wgmma.mma_async.sync.aligned.m64n256k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123,%124,%125,%126,%127}" ", {%128,%129,%130,%131}, %132, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]), "+f"(D[124]), "+f"(D[125]), "+f"(D[126]), "+f"(D[127]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#undef WGMMA_RS_ANY_T +#define WGMMA_RS_ANY_T(NW,D,A,DB,TB) do { if constexpr (NW == 8) { WGMMA_RS_M64N8K16_T(D, A, DB, TB); } else if constexpr (NW == 16) { WGMMA_RS_M64N16K16_T(D, A, DB, TB); } else if constexpr (NW == 24) { WGMMA_RS_M64N24K16_T(D, A, DB, TB); } else if constexpr (NW == 32) { WGMMA_RS_M64N32K16_T(D, A, DB, TB); } else if constexpr (NW == 40) { WGMMA_RS_M64N40K16_T(D, A, DB, TB); } else if constexpr (NW == 48) { WGMMA_RS_M64N48K16_T(D, A, DB, TB); } else if constexpr (NW == 56) { WGMMA_RS_M64N56K16_T(D, A, DB, TB); } else if constexpr (NW == 64) { WGMMA_RS_M64N64K16_T(D, A, DB, TB); } else if constexpr (NW == 72) { WGMMA_RS_M64N72K16_T(D, A, DB, TB); } else if constexpr (NW == 80) { WGMMA_RS_M64N80K16_T(D, A, DB, TB); } else if constexpr (NW == 88) { WGMMA_RS_M64N88K16_T(D, A, DB, TB); } else if constexpr (NW == 96) { WGMMA_RS_M64N96K16_T(D, A, DB, TB); } else if constexpr (NW == 104) { WGMMA_RS_M64N104K16_T(D, A, DB, TB); } else if constexpr (NW == 112) { WGMMA_RS_M64N112K16_T(D, A, DB, TB); } else if constexpr (NW == 120) { WGMMA_RS_M64N120K16_T(D, A, DB, TB); } else if constexpr (NW == 128) { WGMMA_RS_M64N128K16_T(D, A, DB, TB); } else if constexpr (NW == 136) { WGMMA_RS_M64N136K16_T(D, A, DB, TB); } else if constexpr (NW == 144) { WGMMA_RS_M64N144K16_T(D, A, DB, TB); } else if constexpr (NW == 152) { WGMMA_RS_M64N152K16_T(D, A, DB, TB); } else if constexpr (NW == 160) { WGMMA_RS_M64N160K16_T(D, A, DB, TB); } else if constexpr (NW == 168) { WGMMA_RS_M64N168K16_T(D, A, DB, TB); } else if constexpr (NW == 176) { WGMMA_RS_M64N176K16_T(D, A, DB, TB); } else if constexpr (NW == 184) { WGMMA_RS_M64N184K16_T(D, A, DB, TB); } else if constexpr (NW == 192) { WGMMA_RS_M64N192K16_T(D, A, DB, TB); } else if constexpr (NW == 200) { WGMMA_RS_M64N200K16_T(D, A, DB, TB); } else if constexpr (NW == 208) { WGMMA_RS_M64N208K16_T(D, A, DB, TB); } else if constexpr (NW == 216) { WGMMA_RS_M64N216K16_T(D, A, DB, TB); } else if constexpr (NW == 224) { WGMMA_RS_M64N224K16_T(D, A, DB, TB); } else if constexpr (NW == 232) { WGMMA_RS_M64N232K16_T(D, A, DB, TB); } else if constexpr (NW == 240) { WGMMA_RS_M64N240K16_T(D, A, DB, TB); } else if constexpr (NW == 248) { WGMMA_RS_M64N248K16_T(D, A, DB, TB); } else if constexpr (NW == 256) { WGMMA_RS_M64N256K16_T(D, A, DB, TB); } } while (0) diff --git a/deploy/nemotron_native/kernels/b512/native/prefill/dedup_native_prefill_impl.inc b/deploy/nemotron_native/kernels/b512/native/prefill/dedup_native_prefill_impl.inc new file mode 100644 index 0000000..e7e1308 --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/native/prefill/dedup_native_prefill_impl.inc @@ -0,0 +1,249 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 16 || BATCH == 32 || BATCH == 64 || BATCH == 128 || BATCH == 256 || BATCH == 512)) + throw std::runtime_error("native prefill requires matching batch 1/16/32/64/128"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (padded > INT_MAX / (BB_EXPERTS * BB_MOE_INTER)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + float *r9_workspace = nullptr; + mk2_check(owned.alloc(&r9_workspace,size_t(BATCH)*ssm_elements*sizeof(float))); + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + r9_state_t *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(r9_state_t))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(r9_state_t))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = r9_workspace; + st->mb_host[i].r9_private = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + mk2_check(owned.alloc(&st->mb_host[i].r8_canonical_ssm, ssm_elements * sizeof(float))); + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_dtda, 2 * MB_HEADS); + { + const size_t estride = padded; + const size_t grows = size_t(BB_EXPERTS) * estride; + mk2_check(owned.alloc(&hb.pf_moe_gath, grows * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_gath, 0, grows * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_hc, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_hc, 0, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_gidx, grows * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_gidx, 0, grows * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ecnt, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ecnt, 0, BB_EXPERTS * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ntiles, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ntiles, 0, BB_EXPERTS * sizeof(int))); + hb.pf_moe_estride = int(estride); + hb.pf_moe_gath_v = act_gl{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_v = act_gl{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + hb.pf_moe_gath_vn = act_gl_n{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_vn = act_gl_n{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + scratch_bytes += grows * (D_MODEL + BB_MOE_INTER) * sizeof(bf16); + } + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(owned.alloc(&hb.native_pf_cohort_check, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(hb.native_pf_cohort_check, 0, 2 * sizeof(unsigned int))); +#if PF_ACTIVE_TOKEN_EXPERTS + mk2_check(owned.alloc(&hb.native_pf_expert_masks, (padded/BATCH_PAD)*sizeof(unsigned int))); +#endif + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} +extern "C" int mk2_native_prefill_effective_sequences(MK2NativePrefill *p, int *out) { + try { + if (!p || !out) throw std::runtime_error("null effective-sequence query"); + auto *st = reinterpret_cast(p->base); + mk2_check(cudaMemcpy(out, st->host_bb.native_pf_cohort_check + 1, + sizeof(int), cudaMemcpyDeviceToHost)); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" int mk2_r8_initial_state(MK2NativePrefill *p, void **initial, void **ready) { + if (!p || !initial || !ready) return 1; + auto *st = reinterpret_cast(p->base); + for (int i = 0; i < 14; ++i) initial[i] = st->mb_host[i].r8_canonical_ssm; + *ready = st->host_bb.native_pf_cohort_check + 1; + return 0; +} +extern "C" int mk2_r9_prefill_bind(MK2NativePrefill*p,const void*const*scales,const void*const*maps){if(!p)return 1;auto*st=(MK2Stage0State*)p->base;for(int i=0;i<14;++i){st->mb_host[i].r9_scales=(const float*)scales[i];st->mb_host[i].r9_map=(const int*)maps[i];}return (int)cudaMemcpy((void*)st->host_bb.mamba,st->mb_host.data(),14*sizeof(MambaLayer),cudaMemcpyHostToDevice);} diff --git a/deploy/nemotron_native/kernels/b512/native/prefill/dedup_predictor_mk2.cuh b/deploy/nemotron_native/kernels/b512/native/prefill/dedup_predictor_mk2.cuh new file mode 100644 index 0000000..1d8a7c2 --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/native/prefill/dedup_predictor_mk2.cuh @@ -0,0 +1,4594 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef PF_ACTIVE_TOKEN_EXPERTS +#define PF_ACTIVE_TOKEN_EXPERTS 0 +#endif +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +#ifndef MK2_PREFILL_ONLY +#define MK2_PREFILL_ONLY 0 +#endif +#ifndef MK2_BATCH_TILE +#define MK2_BATCH_TILE (((MK2_BATCH + 15) / 16) * 16) +#endif +constexpr int BATCH_PAD = MK2_BATCH_TILE; +static_assert(BATCH >= 1 && (MK2_PREFILL_ONLY || BATCH <= BATCH_PAD)); +static_assert(BATCH_PAD == 16 || BATCH_PAD == 32 || BATCH_PAD == 64 || BATCH_PAD == 128, + "native prefill tile must be 16, 32, 64 or 128"); +static_assert(MK2_PREFILL_ONLY || BATCH_PAD == ((BATCH + 15) / 16) * 16, + "a wider token tile requires a prefill-only library"); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_PW_MULTI +#define PF_SSM_PW_MULTI 4 +#endif +#ifndef PF_SSM_PW_SINGLE +#define PF_SSM_PW_SINGLE 2 +#endif +static_assert(MB_HEAD_DIM % PF_SSM_PW_MULTI == 0 && + MB_HEAD_DIM % PF_SSM_PW_SINGLE == 0, + "head_dim must split evenly into scan channel groups"); +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef MK2_SMEM_SLACK +#define MK2_SMEM_SLACK 1024 +#endif +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +#ifndef MK2_THREADS +#define MK2_THREADS 256 +#endif +constexpr int THREADS = MK2_THREADS; +constexpr int NORM_THREADS = 256; +constexpr int NORM_WARPS = NORM_THREADS / 32; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +#ifndef MK2_GEMM_CONSUMERS +#define MK2_GEMM_CONSUMERS 1 +#endif +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 +#define MK2_HEAD_STAGE_BYTES 0 +#else +#define MK2_HEAD_STAGE_BYTES sizeof(out_tile_f) +#endif +constexpr int NCONS = MK2_GEMM_CONSUMERS; +static_assert(NCONS >= 1 && NCONS <= 3, "consumers occupy warpgroups 0, 2 and 3"); +static_assert(THREADS >= (NCONS + 1) * 128, "one producer plus NCONS consumers"); +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +constexpr int PAD_N = 32; +using act_tile_n = st_bf; +using out_tile_n = st_bf; +using act_gl_n = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_MOE_ACC_C, EPI_ACC_F32 }; +#include +constexpr int R9_STATE_MODE=1; +using r9_state_t=bf16; +__device__ __forceinline__ int r9_scale_index(int head,int channel,int n) {return (((head/8)*4+(head%8)/2)*4+channel/6)*128+n;} +__device__ __forceinline__ float r9_load(r9_state_t v,float scale) { +return __bfloat162float(v); +} +__device__ __forceinline__ r9_state_t r9_store(float x,float scale) { +return __float2bfloat16(x); +} +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; + float *r8_canonical_ssm = nullptr; + r9_state_t *r9_private = nullptr; + const float *r9_scales = nullptr; + const int *r9_map = nullptr; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + unsigned int *native_pf_expert_masks = nullptr; + float *pf_dtda = nullptr; + bf16 *pf_moe_gath = nullptr; + bf16 *pf_moe_hc = nullptr; + int *pf_moe_gidx = nullptr; + int *pf_moe_ecnt = nullptr; + int *pf_moe_ntiles = nullptr; + int pf_moe_estride = 0; + act_gl pf_moe_gath_v; + act_gl pf_moe_hc_v; + act_gl_n pf_moe_gath_vn; + act_gl_n pf_moe_hc_vn; + unsigned int *native_pf_cohort_check = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + volatile unsigned int *generation = state + 1; + const unsigned int seen = *generation; + if (atomicAdd(state, 1u) == gridDim.x - 1) { + atomicExch(state, 0u); + atomicAdd(state + 1, 1u); + } else { + while (*generation == seen) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const bf16 *in = src + (size_t)row * D_MODEL; + float v[PER_THREAD]; + float sum = 0.f, sq = 0.f; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (threadIdx.x < NORM_THREADS) { + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + v[j] = __bfloat162float(in[threadIdx.x + j * NORM_THREADS]); + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + if (lane == 0) { scratch[warp] = sum; scratch[NORM_WARPS + warp] = sq; } + } + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f, q = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) { s += scratch[w]; q += scratch[NORM_WARPS + w]; } + scratch[2 * NORM_WARPS] = s / D_MODEL; + scratch[2 * NORM_WARPS + 1] = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + } + __syncthreads(); + const float mean = scratch[2 * NORM_WARPS], inv = scratch[2 * NORM_WARPS + 1]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + if (threadIdx.x >= NORM_THREADS) break; + const int i = threadIdx.x + j * NORM_THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * weight[i]); + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, + int idx, unsigned mask); +__device__ __forceinline__ void norm_sub_sync(int sub) { + if (THREADS == NORM_THREADS) { __syncthreads(); return; } + asm volatile("bar.sync %0, %1;" :: "r"(6 + sub), "n"(NORM_THREADS) : "memory"); +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, const float *fold_src = nullptr, + const float *fold_routes = nullptr) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + constexpr int NSUB = THREADS / NORM_THREADS; + if (threadIdx.x >= NSUB * NORM_THREADS) return; + const int sub = threadIdx.x / NORM_THREADS; + const int tid = threadIdx.x - sub * NORM_THREADS; + float *const sc = scratch + sub * (NORM_WARPS + 1); + for (int row = blockIdx.x * NSUB + sub; row < n_rows; + row += gridDim.x * NSUB) { + const size_t base = (size_t)row * D_MODEL; + unsigned fold_mask = (1u << BB_EXPERTS) - 1u; + if (fold_src && fold_routes) { + fold_mask = 0u; + const float *rw = fold_routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) + if (rw[e] != 0.f) fold_mask |= 1u << e; + } + float v[PER_THREAD]; + float sq = 0.f; + const int warp = tid >> 5, lane = tid & 31; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + bf16 hv; + if (fold_src) { + hv = bb_moe_fold_deterministic(fold_src, (int)(base + i), fold_mask); + h[base + i] = hv; + } else { + hv = h[base + i]; + } + v[j] = __bfloat162float(hv) + __bfloat162float(r[base + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + if (lane == 0) sc[warp] = sq; + norm_sub_sync(sub); + if (tid == 0) { + float s = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) s += sc[w]; + sc[NORM_WARPS] = rsqrtf(s / D_MODEL + RMS_EPS); + } + norm_sub_sync(sub); + const float inv = sc[NORM_WARPS]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + out[base + i] = __float2bfloat16( + __bfloat162float(weight[i]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + } + norm_sub_sync(sub); + } +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + const int c = idx - row * MB_CONV_DIM; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + const float raw = __bfloat162float( + zxbcdt[(size_t)row * MB_IN_PROJ + MB_INNER + c]); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + xbc_act[(size_t)row * MB_CONV_DIM + c] = __float2bfloat16(acc); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + const int pairs = BATCH * MB_HEADS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_HEADS * split); + const int head = (pair / split) % MB_HEADS; + const int first_channel = (pair % split) * channels; + const int group = head / MB_HEADS_PER_GROUP; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + L.dt_bias[head]); + const float dA = __expf(L.A[head] * dt); + const float Dh = L.D[head]; + float *st = L.ssm_state + (size_t)slot_map[row] * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + for (int p = first_channel; p < first_channel + channels; ++p) { + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + const size_t at = (size_t)p * MB_STATE + n; + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + static_assert(MB_CONV_K == 4, "packed conv taps assume a 4-tap kernel"); + constexpr int CPT = 4; + static_assert(MB_CONV_DIM % CPT == 0, "channel vector width"); + constexpr int CGROUPS = MB_CONV_DIM / CPT; + const int total = n_tokens * CGROUPS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / CGROUPS; + const int c = (idx - t * CGROUPS) * CPT; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + float wv[CPT][MB_CONV_K]; + { + const float4 *wsrc = reinterpret_cast( + L.conv_w + (size_t)c * MB_CONV_K); + #pragma unroll + for (int j = 0; j < CPT / 2; ++j) { + const float4 w4 = wsrc[j]; + const unsigned uu0 = __float_as_uint(w4.x); + const unsigned uu1 = __float_as_uint(w4.y); + const unsigned uu2 = __float_as_uint(w4.z); + const unsigned uu3 = __float_as_uint(w4.w); + wv[2 * j][0] = __uint_as_float(uu0 << 16); + wv[2 * j][1] = __uint_as_float(uu0 & 0xffff0000u); + wv[2 * j][2] = __uint_as_float(uu1 << 16); + wv[2 * j][3] = __uint_as_float(uu1 & 0xffff0000u); + wv[2 * j + 1][0] = __uint_as_float(uu2 << 16); + wv[2 * j + 1][1] = __uint_as_float(uu2 & 0xffff0000u); + wv[2 * j + 1][2] = __uint_as_float(uu3 << 16); + wv[2 * j + 1][3] = __uint_as_float(uu3 & 0xffff0000u); + } + } + float acc[CPT]; + { + const uint2 b2 = *reinterpret_cast(L.conv_b + c); + acc[0] = __uint_as_float(b2.x << 16); + acc[1] = __uint_as_float(b2.x & 0xffff0000u); + acc[2] = __uint_as_float(b2.y << 16); + acc[3] = __uint_as_float(b2.y & 0xffff0000u); + } + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v[CPT]; + if (src < 0) { + #pragma unroll + for (int q = 0; q < CPT; ++q) + v[q] = __bfloat162float( + state[(size_t)q * (MB_CONV_K - 1) + pos + k]); + } else { + const uint2 p = *reinterpret_cast( + zxbcdt + (size_t)src * MB_IN_PROJ + MB_INNER + c); + v[0] = __uint_as_float(p.x << 16); + v[1] = __uint_as_float(p.x & 0xffff0000u); + v[2] = __uint_as_float(p.y << 16); + v[3] = __uint_as_float(p.y & 0xffff0000u); + } + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] += wv[q][k] * v[q]; + } + uint2 out; + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] = acc[q] / (1.f + __expf(-acc[q])); + out.x = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[0])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[1])) << 16); + out.y = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[2])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[3])) << 16); + *reinterpret_cast(xbc_act + (size_t)t * MB_CONV_DIM + c) = out; + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_dtda_prefill( + const MambaLayer &L, const bf16 *zxbcdt, float *dtda, int n_tokens) { + const int total = n_tokens * MB_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_HEADS; + const int head = idx - t * MB_HEADS; + const float dt = mamba_softplus( + __bfloat162float(zxbcdt[(size_t)t * MB_IN_PROJ + MB_INNER + + MB_CONV_DIM + head]) + L.dt_bias[head]); + dtda[(size_t)idx * 2] = dt; + dtda[(size_t)idx * 2 + 1] = __expf(L.A[head] * dt); + } +} +template +__device__ __forceinline__ void mamba_ssm_prefill_impl( + const MambaLayer &L, const bf16 *xbc_act, const float *dtda, bf16 *gated, + int n_pos, int n_seq, int slot, const int *seq_slot) { + constexpr int NS = MB_STATE / 32; + constexpr int PG = MB_HEAD_DIM / PW; + constexpr bool RS = (PW % 4) == 0; + constexpr int H1 = RS ? PW / 2 : 1; + constexpr int H2 = RS ? PW / 4 : 1; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const bool keep_hi = (lane & 16) != 0; + const bool keep_hi2 = (lane & 8) != 0; + const int warps_total = gridDim.x * WARPS; + const int units = n_seq * MB_HEADS * PG; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int pg = unit % PG; + const int rest = unit / PG; + const int head = rest % MB_HEADS; + const int seq = rest / MB_HEADS; + const int p0 = pg * PW; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Dh = L.D[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p0 * MB_STATE; + float sreg[PW][NS]; + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + sreg[q][j] = st[(size_t)q * MB_STATE + lane + j * 32]; + } + auto fetch = [&](int p, float (&oB)[NS], float (&oC)[NS], + float (&oX)[PW], float2 &od) { + const size_t tt = (size_t)p * n_seq + seq; + const bf16 *xa2 = xbc_act + tt * MB_CONV_DIM; + const bf16 *Bs = xa2 + MB_INNER + group * MB_STATE; + const bf16 *Cs = xa2 + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + od = *reinterpret_cast( + dtda + (tt * MB_HEADS + head) * 2); + #pragma unroll + for (int j = 0; j < NS; ++j) { + oB[j] = __bfloat162float(Bs[lane + j * 32]); + oC[j] = __bfloat162float(Cs[lane + j * 32]); + } + if constexpr (RS) { + const uint2 *xp = reinterpret_cast( + xa2 + head * MB_HEAD_DIM + p0); + #pragma unroll + for (int q = 0; q < PW; q += 4) { + const uint2 w = xp[q >> 2]; + oX[q] = __uint_as_float(w.x << 16); + oX[q + 1] = __uint_as_float(w.x & 0xffff0000u); + oX[q + 2] = __uint_as_float(w.y << 16); + oX[q + 3] = __uint_as_float(w.y & 0xffff0000u); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) + oX[q] = __bfloat162float(xa2[head * MB_HEAD_DIM + p0 + q]); + } + }; + float pB[NS], pC[NS], pX[PW]; + float2 pd; + fetch(0, pB, pC, pX, pd); + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const float dt = pd.x; + const float dA = pd.y; + float Bv[NS], Cv[NS]; + #pragma unroll + for (int j = 0; j < NS; ++j) { + Bv[j] = pB[j] * dt; + Cv[j] = pC[j]; + } + float acc[PW], xvs[PW]; + #pragma unroll + for (int q = 0; q < PW; ++q) xvs[q] = pX[q]; + if (pos + 1 < n_pos) fetch(pos + 1, pB, pC, pX, pd); + #pragma unroll + for (int q = 0; q < PW; ++q) { + const float xv = xvs[q]; + float a = 0.f; + #pragma unroll + for (int j = 0; j < NS; ++j) { + sreg[q][j] = sreg[q][j] * dA + Bv[j] * xv; + a += sreg[q][j] * Cv[j]; + } + acc[q] = a; + } + if constexpr (RS) { + float m1[H1], t1[H1], x1[H1]; + #pragma unroll + for (int q = 0; q < H1; ++q) { + m1[q] = keep_hi ? acc[q + H1] : acc[q]; + t1[q] = keep_hi ? acc[q] : acc[q + H1]; + x1[q] = keep_hi ? xvs[q + H1] : xvs[q]; + } + #pragma unroll + for (int q = 0; q < H1; ++q) + m1[q] += __shfl_xor_sync(0xffffffff, t1[q], 16); + float m2[H2], t2[H2], x2[H2]; + #pragma unroll + for (int q = 0; q < H2; ++q) { + m2[q] = keep_hi2 ? m1[q + H2] : m1[q]; + t2[q] = keep_hi2 ? m1[q] : m1[q + H2]; + x2[q] = keep_hi2 ? x1[q + H2] : x1[q]; + } + #pragma unroll + for (int q = 0; q < H2; ++q) + m2[q] += __shfl_xor_sync(0xffffffff, t2[q], 8); + #pragma unroll + for (int q = 0; q < H2; ++q) { + #pragma unroll + for (int off = 4; off > 0; off >>= 1) + m2[q] += __shfl_xor_sync(0xffffffff, m2[q], off); + } + if ((lane & 7) == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + + p0 + (keep_hi ? H1 : 0) + (keep_hi2 ? H2 : 0); + #pragma unroll + for (int q = 0; q < H2; ++q) + dst[q] = __float2bfloat16(m2[q] + x2[q] * Dh); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc[q] += __shfl_xor_sync(0xffffffff, acc[q], off); + } + if (lane == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + p0; + #pragma unroll + for (int q = 0; q < PW; ++q) + dst[q] = __float2bfloat16(acc[q] + xvs[q] * Dh); + } + } + } + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + st[(size_t)q * MB_STATE + lane + j * 32] = sreg[q][j]; + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, + const float *dtda, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + (void)zxbcdt; + if (n_seq > 1) + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens / n_seq, n_seq, slot, seq_slot); + else + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens, 1, slot, nullptr); +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16( + held[j] * inv * __bfloat162float(L.gnorm_w[idx])); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = BATCH * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / BB_ATT_HEADS; + const int head = pair % BB_ATT_HEADS; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const int kv = head / kv_per_q; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *q = base + head * BB_HEAD_DIM; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (head % kv_per_q == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + for (int tpos = 0; tpos < T; ++tpos) { +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new + : kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#endif + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m, s); + const float rescale = __expf(m - m_new); + const float w = __expf(s - m_new); + denom = denom * rescale + w; +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *vt = kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#else + const bf16 *vt = cur_token + ? v_new + : kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + w * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + constexpr int CHUNKS = BB_HEAD_DIM / 8; + const int total = n_tokens * BB_KV_HEADS * CHUNKS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int c = idx % CHUNKS; + const int rest = idx / CHUNKS; + const int t = rest / BB_KV_HEADS; + const int kv = rest - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM + c * 8; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM + c * 8; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv) + c * 8; + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv) + c * 8; + *reinterpret_cast(kc) = *reinterpret_cast(k_new); + *reinterpret_cast(vc) = *reinterpret_cast(v_new); + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float qv[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qv[j] = __bfloat162float(q[lane + j * 32]); + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + const int *const krow = + L.block_table + ((size_t)seq * 2 + 0) * L.max_blocks_per_seq; + const int *const vrow = + L.block_table + ((size_t)seq * 2 + 1) * L.max_blocks_per_seq; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + const int head_bias = kv * L.tokens_per_block * BB_HEAD_DIM; + int page = -1; + const bf16 *kpage = nullptr, *vpage = nullptr; + float kn[PER_LANE], vn[PER_LANE]; + { + page = 0; + kpage = L.pool_base + (size_t)(krow[0] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + (size_t)(vrow[0] & 0x7FFFFFFF) * elems_per_block; + const int at = head_bias + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + for (int s = 0; s <= pos; ++s) { + float kc[PER_LANE], vc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { kc[j] = kn[j]; vc[j] = vn[j]; } + if (s < pos) { + const int np = (s + 1) / L.tokens_per_block; + if (np != page) { + page = np; + kpage = L.pool_base + + (size_t)(krow[np] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + + (size_t)(vrow[np] & 0x7FFFFFFF) * elems_per_block; + } + const int at = head_bias + + ((s + 1) % L.tokens_per_block) * BB_HEAD_DIM + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) dot += qv[j] * kc[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[j] = acc[j] * rescale + wt * vc[j]; + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH, unsigned int *tile_masks = nullptr, + bf16 *gath = nullptr, int *gidx = nullptr, int *ecnt = nullptr, + int estride = 0, bf16 *row_stage = nullptr) { + __shared__ float score[BB_EXPERTS]; + __shared__ int s_pick[BB_TOP_K], s_slot[BB_TOP_K]; + constexpr int CH = D_MODEL / 8; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + { + const float4 *gsrc = + reinterpret_cast(normed + (size_t)row * D_MODEL); + float4 *gdst = reinterpret_cast(row_stage); + for (int i = threadIdx.x; i < CH; i += THREADS) gdst[i] = gsrc[i]; + } + if ((int)threadIdx.x < BB_EXPERTS) + moe_w[(size_t)row * BB_EXPERTS + threadIdx.x] = 0.f; + __syncthreads(); + for (int e = warp; e < BB_EXPERTS; e += WARPS) { + const bf16 *wr = gate_w + (size_t)e * D_MODEL; + float acc = 0.f; + for (int i = lane; i < D_MODEL; i += 32) + acc += __bfloat162float(row_stage[i]) * + __bfloat162float(wr[i]); + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) score[e] = 1.f / (1.f + __expf(-acc)); + } + __syncthreads(); + if (threadIdx.x == 0) { + float remaining[BB_EXPERTS]; + int pick[BB_TOP_K]; + for (int e = 0; e < BB_EXPERTS; ++e) + remaining[e] = score[e] + __bfloat162float(gate_bias[e]); + float total = 0.f; + for (int k = 0; k < BB_TOP_K; ++k) { + int best = 0; + for (int e = 1; e < BB_EXPERTS; ++e) + if (remaining[e] > remaining[best]) best = e; + pick[k] = best; + total += score[best]; + remaining[best] = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + unsigned int bits = 0u; + for (int k = 0; k < BB_TOP_K; ++k) { + moe_w[row * BB_EXPERTS + pick[k]] = score[pick[k]] * norm; + bits |= 1u << pick[k]; + s_pick[k] = pick[k]; + } + if (tile_masks) atomicOr(&tile_masks[row / BATCH_PAD], bits); + } + __syncthreads(); + if (gath) { + if ((int)threadIdx.x < BB_TOP_K) { + const int k = threadIdx.x; + const int slot = atomicAdd(&ecnt[s_pick[k]], 1); + s_slot[k] = slot; + gidx[(size_t)s_pick[k] * estride + slot] = row; + } + __syncthreads(); + const float4 *src = reinterpret_cast(row_stage); + for (int k = 0; k < BB_TOP_K; ++k) { + float4 *dst = reinterpret_cast( + gath + ((size_t)s_pick[k] * estride + s_slot[k]) * D_MODEL); + for (int i = threadIdx.x; i < CH; i += THREADS) dst[i] = src[i]; + } + } + __syncthreads(); + } +} +__device__ __forceinline__ unsigned moe_row_mask(const float *routes, int row) { + unsigned m = 0u; + const float *rw = routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) if (rw[e] != 0.f) m |= 1u << e; + return m; +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx, + unsigned mask) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + for (unsigned m = mask; m; m &= m - 1u) + value = __fadd_rn(value, parts[base + (size_t)(__ffs(m) - 1) * D_MODEL]); + return __float2bfloat16( + __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL])); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + uint64_t *keys, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int n = threadIdx.x; n < SORT_N; n += THREADS) { + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + keys[n] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + __syncthreads(); + for (int k = 2; k <= SORT_N; k <<= 1) { + if (k >= 64) __syncthreads(); + for (int j = k >> 1; j > 0; j >>= 1) { + for (int i = threadIdx.x; i < SORT_N; i += THREADS) { + const int partner = i ^ j; + if (partner > i) { + const bool descending = ((i & k) == 0); + if ((keys[i] < keys[partner]) == descending) { + const uint64_t t = keys[i]; + keys[i] = keys[partner]; + keys[partner] = t; + } + } + } + if (j >= 32) __syncthreads(); + else __syncwarp(); + } + } + __syncthreads(); + if (threadIdx.x == 0) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float v; + int n; + if (!eos_used && eos_v > av) { + v = eos_v; + n = AUDIO_EOS; + eos_used = true; + } else { + v = av; + n = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = + v + gumbel[((size_t)row * N_POS + pos) * TOP_K + j]; + if (noised > best_noised) { best_noised = noised; best_code = n; } + } + out_codes[row] = best_code; + } + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void prefill_expert_masks(const float *routes, unsigned int *masks, int rows) { + const int lane=threadIdx.x&31, warp=threadIdx.x>>5; + const int nt=(rows+BATCH_PAD-1)/BATCH_PAD; + for (int t=blockIdx.x*WARPS+warp;t(gath + g * D_MODEL + c * 8); + if (local < ecnt[e]) + *dst = *reinterpret_cast( + normed + (size_t)gidx[g] * D_MODEL + c * 8); + else + *dst = make_float4(0.f, 0.f, 0.f, 0.f); + } +} +__device__ __forceinline__ int gemm_eff(int tiles, int reps) { + if (NCONS <= 1) return 1; + const int g = (int)gridDim.x; + const int rounds1 = (tiles * reps + g - 1) / g; + if constexpr (NCONS == 3) { + const int max_eff = reps <= BB_EXPERTS ? 2 : 3; + const int rounds2 = (((tiles + 1) / 2) * reps + g - 1) / g; + if (max_eff == 2) return rounds2 < rounds1 ? 2 : 1; + const int rounds3 = (((tiles + 2) / 3) * reps + g - 1) / g; + if (rounds3 < rounds2 && rounds3 < rounds1) return 3; + return rounds2 < rounds1 ? 2 : 1; + } else { + const int roundsN = (((tiles + NCONS - 1) / NCONS) * reps + g - 1) / g; + return roundsN < rounds1 ? NCONS : 1; + } +} +__device__ __forceinline__ int gemm_groups(int tiles, int reps) { + const int eff = gemm_eff(tiles, reps); + return (tiles + eff - 1) / eff; +} +__device__ __forceinline__ int gemm_consumer_id() { + const int g = warpgroup::groupid(); + if (g == 1) return -1; + const int c = g == 0 ? 0 : g - 1; + return c < NCONS ? c : -1; +} +__device__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][BATCH_PAD]; + __shared__ int moe_gidx_stage[NCONS][BATCH_PAD]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + a_s[s], a, {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + a_s[s], a, {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / BATCH_PAD) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, BATCH_PAD> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], a_s[s]); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * BATCH_PAD; + if (epi == EPI_BIAS) { + warpgroup::store(o_sf, accum); + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + dest[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(o_sf[int2{col, row}] + + __bfloat162float(bias[base + col])); + } + warpgroup::sync(1 + cons); + continue; + } + warpgroup::store(o_s[cons], accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_s[cons][int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__device__ void gemm_narrow( + const weight_gl &w, const act_gl_n &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = PAD_N, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][PAD_N]; + __shared__ int moe_gidx_stage[NCONS][PAD_N]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + PAD_N - 1) / PAD_N; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + PAD_N - 1) / PAD_N; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / PAD_N) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, PAD_N> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], + *reinterpret_cast(&a_s[s])); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * PAD_N; + out_tile_n &o_sn = *reinterpret_cast(&o_s[cons]); + warpgroup::store(o_sn, accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_sn[int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&w_s)[STAGES * NCONS] = + alloc.allocate(); + act_tile (&a_s)[STAGES] = alloc.allocate(); + out_tile (&o_s)[NCONS] = alloc.allocate(); +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 + static_assert(NCONS * sizeof(out_tile) >= sizeof(out_tile_f), + "aliased head staging must fit the output tiles"); + out_tile_f &o_sf = *reinterpret_cast(&o_s[0]); +#else + out_tile_f (&o_sf) = alloc.allocate(); +#endif + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_WARPS + 2]; +#if MK2_PREFILL_ONLY + bf16 *const route_stage = reinterpret_cast(&o_s[0]); + static_assert(sizeof(o_s) >= D_MODEL * sizeof(bf16), + "router row stage must fit the output tiles"); +#else + __shared__ __align__(16) uint64_t sort_keys[SORT_N]; + bf16 *const route_stage = reinterpret_cast(&sort_keys[0]); + static_assert(sizeof(sort_keys) >= D_MODEL * sizeof(bf16), + "router row stage must fit the sampler workspace"); +#endif + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS * NCONS, 0); + for (int c = 0; c < CONSUMER_WARPS * NCONS; ++c) arrive(drained[s]); + } + } + int gemm_iter = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map); + grid_barrier(g.barrier); + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map); + grid_barrier(g.barrier); + mamba_gated_norm(L, bb->zxbcdt, bb->gated); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, nullptr, 1, BATCH, + bb->moe_w, BB_MOE_INTER / OUT_ATOM); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + 1, BATCH, bb->moe_w, D_MODEL / OUT_ATOM); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_deterministic(bb->moe_out, idx, + (1u << BB_EXPERTS) - 1u); + grid_barrier(g.barrier); + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int original_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + bool identical_cohort = false; + if constexpr (BATCH > 1) { + if (bb->native_pf_cohort_check && bb->native_prefill_reset && + bb->pf_n_seq == BATCH && !bb->pf_fuse_decode && + !bb->pf_layer_out && !bb->moe_w_inject && !bb->moe_w_audit) { + if (blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[0] = 0; + grid_barrier(g.barrier); + const auto *bits = reinterpret_cast(bb->prompt); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < original_total * D_MODEL; idx += gridDim.x * THREADS) { + const int canonical = (idx / (BATCH * D_MODEL)) * BATCH * D_MODEL + idx % D_MODEL; + if (bits[idx] != bits[canonical]) atomicOr(bb->native_pf_cohort_check, 1u); + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + if (bb->pf_seq_slot[row] != row || bb->slot_map[row] != row) + atomicOr(bb->native_pf_cohort_check, 1u); + grid_barrier(g.barrier); + identical_cohort = bb->native_pf_cohort_check[0] == 0; + } + } + const int n_total = identical_cohort ? original_total / BATCH : original_total; + if (bb->native_pf_cohort_check && blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[1] = identical_cohort ? 1 : BATCH; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int reset_slots = identical_cohort ? 1 : BATCH; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * reset_slots * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % reset_slots; + const MambaLayer &L = bb->mamba[row_layer / reset_slots]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = original_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + const int src = identical_cohort ? (idx / D_MODEL) * BATCH * D_MODEL + idx % D_MODEL : idx; + bb->pf_h[idx] = bb->prompt[src]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = identical_cohort ? 1 : (bb->pf_n_seq > 0 ? bb->pf_n_seq : 1); + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T, + (!bb->pf_layer_out && layer > 0 && + bb->moe_of_layer[layer - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + if (bb->native_pf_expert_masks) + for (int i = blockIdx.x * THREADS + threadIdx.x; i < n_tok; + i += gridDim.x * THREADS) + bb->native_pf_expert_masks[i] = 0u; + if (bb->pf_moe_ecnt) + for (int e = blockIdx.x * THREADS + threadIdx.x; e < BB_EXPERTS; + e += gridDim.x * THREADS) + bb->pf_moe_ecnt[e] = 0; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + for(size_t j=blockIdx.x*THREADS+threadIdx.x;jpf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + mamba_dtda_prefill(L, bb->pf_zxbcdt, bb->pf_dtda, T); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_dtda, + bb->pf_gated, T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + for(size_t j=blockIdx.x*THREADS+threadIdx.x;jpf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, + T, bb->native_pf_expert_masks, bb->pf_moe_gath, + bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, route_stage); + } + grid_barrier(g.barrier); + int max_ecnt = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + max_ecnt = max(max_ecnt, bb->pf_moe_ecnt[e]); + const bool narrow = BATCH_PAD > PAD_N && max_ecnt <= PAD_N; + const int etile = narrow ? PAD_N : BATCH_PAD; + int nt_sum = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_sum += (bb->pf_moe_ecnt[e] + etile - 1) / etile; + const int up_steps = + gemm_groups(BB_MOE_INTER / OUT_ATOM, nt_sum) * + nt_sum * (D_MODEL / KCHUNK); + const int shup_steps = + gemm_groups(BB_SHARED_INTER / OUT_ATOM, n_tok) * + n_tok * (D_MODEL / KCHUNK); + const int down_steps = + gemm_groups(D_MODEL / OUT_ATOM, nt_sum) * + nt_sum * (BB_MOE_INTER / KCHUNK); + const int shdown_steps = + gemm_groups(D_MODEL / OUT_ATOM, n_tok) * + n_tok * (BB_SHARED_INTER / KCHUNK); + if (narrow) + gemm_narrow(L.exp_up, bb->pf_moe_gath_vn, + BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, PAD_N, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + else + gemm(L.exp_up, bb->pf_moe_gath_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, BATCH_PAD, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + up_steps, up_steps + shup_steps); + grid_barrier(g.barrier); + if (narrow) + gemm_narrow(L.exp_down, bb->pf_moe_hc_vn, + BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, PAD_N, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + else + gemm(L.exp_down, bb->pf_moe_hc_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + down_steps, down_steps + shdown_steps); + grid_barrier(g.barrier); + if (bb->pf_layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * D_MODEL; idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = + bb_moe_fold_deterministic(bb->pf_moe_out, idx, + moe_row_mask(bb->pf_moe_w, idx / D_MODEL)); + grid_barrier(g.barrier); + } + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, + identical_cohort ? bb->pf_normed : bb->pf_final, ln_scratch, + n_total, + (!bb->pf_layer_out && n_backbone_layers > 0 && + bb->moe_of_layer[n_backbone_layers - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + grid_barrier(g.barrier); + if (identical_cohort) { + const int worker = blockIdx.x * THREADS + threadIdx.x; + const int workers = gridDim.x * THREADS; + for (int i = worker; i < original_total * D_MODEL; i += workers) + bb->pf_final[i] = bb->pf_normed[(i / (BATCH * D_MODEL)) * D_MODEL + i % D_MODEL]; + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = worker; i < 14 * (BATCH - 1) * CV; i += workers) { + const int cell = i % CV, pair = i / CV; + const int row = pair % (BATCH - 1) + 1; + const MambaLayer &L = bb->mamba[pair / (BATCH - 1)]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = L.conv_state[cell]; + } + const int kv_columns = BB_KV_HEADS * BB_HEAD_DIM; + for (size_t i = worker; i < size_t(3) * 2 * (BATCH - 1) * n_total * kv_columns; i += workers) { + size_t q = i / kv_columns; + const int col = i % kv_columns; + const int pos = q % n_total; q /= n_total; + const int row = q % (BATCH - 1) + 1; q /= BATCH - 1; + const int kind = q % 2, layer = q / 2; + const AttnLayer &L = bb->attn[layer]; + const int offset = kv_offset(L, pos, col / BB_HEAD_DIM) + col % BB_HEAD_DIM; + kv_slot(L, row, kind, pos)[offset] = kv_slot(L, 0, kind, pos)[offset]; + } + grid_barrier(g.barrier); + } + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } +#if MK2_PREFILL_ONLY + return; +#else + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, bb->phoneme_b); + grid_barrier(g.barrier); + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + if (threadIdx.x == 0) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = 0; n < BB_PHONEME; ++n) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + bb->phoneme_out[row] = arg; + } + __syncthreads(); + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + for (int i = lane; i < D_HEAD; i += 32) { + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + for (int t = 0; t <= pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + for (int i = lane; i < D_HEAD; i += 32) { + float acc = 0.f; + for (int t = 0; t <= pos; ++t) + acc += sc[t] * __bfloat162float( + g.vcache[cache_base + t * D_HEAD + i]); + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + i] = + __float2bfloat16(acc); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, g.sampled_pos, + sort_keys, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { +#ifdef ABLATE_SAMPLE + __syncthreads(); +#endif + } else { + grid_barrier(g.barrier); + } + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x < BATCH && blockIdx.x == 0) { + const int code = g.sampled_pos[threadIdx.x]; + g.sampled[threadIdx.x * N_POS + pos] = code; + if (g.native_codes) + g.native_codes[threadIdx.x * N_POS + pos] = (long long)code; + } + grid_barrier(g.barrier); + } + if (g.native_seq_len && blockIdx.x == 0 && threadIdx.x < BATCH) + g.native_seq_len[threadIdx.x] += 1; + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +#endif +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 2 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK), + 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = MK2_PREFILL_ONLY + ? out->sms * out->blocks_per_sm + : std::min(out->sms * out->blocks_per_sm, D_FFN / OUT_ATOM); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + 0, + nullptr, + nullptr, + nullptr, nullptr, nullptr, nullptr, nullptr, 0, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "dedup_native_prefill_impl.inc" diff --git a/deploy/nemotron_native/kernels/b512/native/prefill/mk2_abi.cuh b/deploy/nemotron_native/kernels/b512/native/prefill/mk2_abi.cuh new file mode 100644 index 0000000..95a253b --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/native/prefill/mk2_abi.cuh @@ -0,0 +1,82 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} diff --git a/deploy/nemotron_native/kernels/b512/native/prefill/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/b512/native/prefill/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/native/prefill/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/b512/native/prefill/prefill.cu b/deploy/nemotron_native/kernels/b512/native/prefill/prefill.cu new file mode 100644 index 0000000..4d82c95 --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/native/prefill/prefill.cu @@ -0,0 +1,18 @@ +#define MK2_BATCH 512 +#define MK2_BATCH_TILE 64 +#define MK2_PREFILL_ONLY 1 +#define STAGES 3 +#define KCHUNK 128 +#define MK2_GEMM_CONSUMERS 3 +#define MK2_THREADS 512 +#define MK2_NO_MAIN 1 +#define MK2_DIRECT_INPUT 1 +#define MK2_B1_LOCAL_BARRIERS 1 +#define MK2_SKIP_HIDDEN_SNAPSHOT 1 +#define MK2_ACTIVE_EXPERTS 1 +#define MK2_SSM_SPLITS 1 +#define PF_SSM_PW_MULTI 12 +#define PF_ACTIVE_TOKEN_EXPERTS 1 +#include +#include +#include "dedup_predictor_mk2.cuh" diff --git a/deploy/nemotron_native/kernels/b512/native/stage0/mk2_abi.cuh b/deploy/nemotron_native/kernels/b512/native/stage0/mk2_abi.cuh new file mode 100644 index 0000000..bf1e415 --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/native/stage0/mk2_abi.cuh @@ -0,0 +1,85 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_set_native_predictor_active_mask(MK2Predictor *p, const unsigned char *mask); +int mk2_native_startup_skip_abi_version(); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *, int grid); diff --git a/deploy/nemotron_native/kernels/b512/native/stage0/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/b512/native/stage0/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/native/stage0/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/b512/native/stage0/native_prefill_impl.inc b/deploy/nemotron_native/kernels/b512/native/stage0/native_prefill_impl.inc new file mode 100644 index 0000000..a254308 --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/native/stage0/native_prefill_impl.inc @@ -0,0 +1,198 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 16 || BATCH == 32 || BATCH == 64)) + throw std::runtime_error("native prefill requires matching batch 1/16/32/64"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (padded > INT_MAX / ((BB_EXPERTS + 1) * D_MODEL)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + float *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(float))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(float))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} diff --git a/deploy/nemotron_native/kernels/b512/native/stage0/predictor_mk2.cu b/deploy/nemotron_native/kernels/b512/native/stage0/predictor_mk2.cu new file mode 100644 index 0000000..db2404d --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/native/stage0/predictor_mk2.cu @@ -0,0 +1,5563 @@ +#include "../../../inflight/progress.cuh" +#if MK2_INFLIGHT +#include +#endif +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#define MK2_FUSE_CONV (MK2_SSM_SPLITS == 1) +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +constexpr int BATCH_PAD = 64; +constexpr int DECODE_TILES = BATCH / BATCH_PAD; +constexpr int DECODE_ROWS = BATCH_PAD; +static_assert(BATCH == 512 || BATCH == 1024); +static_assert(BATCH % BATCH_PAD == 0); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#ifdef MK2_NO_MAIN +#define MK2_BISECTION_HOOKS 0 +#else +#define MK2_BISECTION_HOOKS 1 +#endif +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +#ifndef NMULT +#define NMULT 1 +#endif +#ifndef SSM_PREFETCH_MAX +#define SSM_PREFETCH_MAX 8 +#endif +constexpr int NM_MAX = (NMULT > 1 && DECODE_TILES % NMULT == 0) ? NMULT : 1; +#ifndef MK2_OUT_MULT +#define MK2_OUT_MULT 1 +#endif +constexpr int MT = MK2_OUT_MULT; +static_assert(MT >= 1 && MT <= NM_MAX, "an item holds at most NM_MAX accumulators"); +constexpr int SLOTS = 1 + NM_MAX; +static_assert(MT + 1 <= SLOTS); +constexpr int THREADS = 256; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +constexpr int NORM_SCRATCH_STRIDE = 2 * WARPS + 2; +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_ACC_F32 }; +#define EPI_ROW(idx) ((OUT_ATOM == 64 && BATCH_PAD == 64) \ + ? (((idx) & 15) | ((((idx) >> 6) & 3) << 4)) : ((idx) / (OUT_ATOM / 8))) +#define EPI_OCT(idx) ((OUT_ATOM == 64 && BATCH_PAD == 64) \ + ? ((((idx) >> 4) & 3) | ((((idx) >> 8) & 1) << 2)) : ((idx) % (OUT_ATOM / 8))) +struct alignas(16) bf16_oct { bf16 x0, x1, x2, x3, x4, x5, x6, x7; }; +static_assert(sizeof(bf16_oct) == 16); +__device__ __forceinline__ uint32_t pack_bf16_pair(float2 value) { + const bf16_2 packed = + base_types::convertor::convert(value); + static_assert(sizeof(bf16_2) == sizeof(uint32_t)); + return (uint32_t)__bfloat16_as_ushort(packed.x) | + ((uint32_t)__bfloat16_as_ushort(packed.y) << 16); +} +#include +constexpr int R9_STATE_MODE=1; +using r9_state_t=bf16; +__device__ __forceinline__ int r9_scale_index(int head,int channel,int n) {return (((head/8)*4+(head%8)/2)*4+channel/6)*128+n;} +__device__ __forceinline__ float r9_load(r9_state_t v,float scale) { +return __bfloat162float(v); +} +__device__ __forceinline__ r9_state_t r9_store(float x,float scale) { +return __float2bfloat16(x); +} +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; + const float *r8_initial_ssm = nullptr; + const unsigned int *r8_initial_ready = nullptr; + const int *r8_seq_len = nullptr; + int r8_initial_positions = 0; + const float *r9_scales = nullptr; + const int *r9_map = nullptr; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + bf16 *r6_packed = nullptr; + int *r6_map = nullptr,*r6_counts = nullptr; + act_gl *r6_views = nullptr; + int *r6_lut = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +#define BARRIER_GEN_WORD 32 +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + constexpr unsigned long long GRID_BLOCKS = 132ull; + auto *wide_state = reinterpret_cast(state); + unsigned long long arrival; + asm volatile("atom.global.add.release.gpu.u64 %0, [%1], 1;" + : "=l"(arrival) : "l"(wide_state) : "memory"); + const unsigned int rank = (unsigned int)(arrival % GRID_BLOCKS); + if (rank != GRID_BLOCKS - 1u) { + const unsigned long long target = arrival - rank + GRID_BLOCKS; + unsigned long long seen; + for (;;) { + asm volatile("ld.global.relaxed.gpu.u64 %0, [%1];" + : "=l"(seen) : "l"(wide_state) : "memory"); + if (seen >= target) break; + __nanosleep(BARRIER_BACKOFF_NS); + } + } + unsigned long long release_seen; + asm volatile("ld.global.acquire.gpu.u64 %0, [%1];" + : "=l"(release_seen) : "l"(wide_state) : "memory"); + asm volatile("" :: "l"(release_seen)); + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / THREADS; + float wv[PER_THREAD]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) wv[j] = weight[threadIdx.x + j * THREADS]; + float v[PER_THREAD], v_next[PER_THREAD]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + v[j] = __bfloat162float( + src[(size_t)blockIdx.x * D_MODEL + threadIdx.x + j * THREADS]); + for (int row = blockIdx.x, row_iter = 0; row < BATCH; + row += gridDim.x, ++row_iter) { + float *row_scratch = scratch + (row_iter & 1) * NORM_SCRATCH_STRIDE; + const int next_row = row + (int)gridDim.x; + if (next_row < BATCH) { + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + v_next[j] = __bfloat162float( + src[(size_t)next_row * D_MODEL + threadIdx.x + j * THREADS]); + } + float sum = 0.f, sq = 0.f; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) { + row_scratch[warp] = sum; + row_scratch[WARPS + warp] = sq; + } + __syncthreads(); + float s = 0.f, q = 0.f; + for (int w = 0; w < WARPS; ++w) { + s += row_scratch[w]; + q += row_scratch[WARPS + w]; + } + const float mean = s / D_MODEL; + const float inv = rsqrtf(q / D_MODEL - mean * mean + LN_EPS); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * wv[j]); + } + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) v[j] = v_next[j]; + } +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, bf16 *snapshot = nullptr, int snapshot_layer = -1, + float *snapshot_fp32 = nullptr, float *out_fp32 = nullptr, + const float *fold_src = nullptr, const float *fold_routes = nullptr) { + constexpr int PER_THREAD = D_MODEL / THREADS; + bf16 wv[PER_THREAD], hv_cur[PER_THREAD], rv_cur[PER_THREAD]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) wv[j] = weight[threadIdx.x + j * THREADS]; + if ((int)blockIdx.x < n_rows) { + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const size_t at = (size_t)blockIdx.x * D_MODEL + threadIdx.x + j * THREADS; + if (!fold_src) hv_cur[j] = h[at]; + rv_cur[j] = r[at]; + } + } + for (int row = blockIdx.x, row_iter = 0; row < n_rows; + row += gridDim.x, ++row_iter) { + bf16 hv_next[PER_THREAD], rv_next[PER_THREAD]; + const int next_row = row + (int)gridDim.x; + if (next_row < n_rows) { + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const size_t at = (size_t)next_row * D_MODEL + threadIdx.x + j * THREADS; + if (!fold_src) hv_next[j] = h[at]; + rv_next[j] = r[at]; + } + } + float *row_scratch = scratch + (row_iter & 1) * NORM_SCRATCH_STRIDE; + const size_t base = (size_t)row * D_MODEL; + float fv[PER_THREAD]; + if (fold_src) { + const float *rw = fold_routes + (size_t)row * BB_EXPERTS; + const int lane = threadIdx.x & 31; + unsigned mask = __ballot_sync( + 0xffffffffu, lane < BB_EXPERTS && rw[lane] != 0.f); + const float *pe = fold_src + + (size_t)row * (BB_EXPERTS + 1) * D_MODEL + threadIdx.x; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) fv[j] = 0.f; + if (__popc(mask) >= BB_TOP_K) { + #pragma unroll + for (int k = 0; k < BB_TOP_K; ++k) { + const int e = __ffs(mask) - 1; + mask &= mask - 1; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + fv[j] = __fadd_rn( + fv[j], pe[(size_t)e * D_MODEL + j * THREADS]); + } + } + while (mask) { + const int e = __ffs(mask) - 1; + mask &= mask - 1; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + fv[j] = __fadd_rn( + fv[j], pe[(size_t)e * D_MODEL + j * THREADS]); + } + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) + fv[j] = __fadd_rn(fv[j], + pe[(size_t)BB_EXPERTS * D_MODEL + j * THREADS]); + } + float v[PER_THREAD]; + float sq = 0.f; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + const bf16 hv = fold_src ? __float2bfloat16(fv[j]) : hv_cur[j]; + if (fold_src) h[base + i] = hv; + if (snapshot) + snapshot[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = hv; + if (snapshot_fp32) + snapshot_fp32[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = + __bfloat162float(hv); + v[j] = __bfloat162float(hv) + __bfloat162float(rv_cur[j]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) row_scratch[warp] = sq; + __syncthreads(); + float s = 0.f; + for (int w = 0; w < WARPS; ++w) s += row_scratch[w]; + const float inv = rsqrtf(s / D_MODEL + RMS_EPS); + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + const bf16 normed = __float2bfloat16( + __bfloat162float(wv[j]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + out[base + i] = normed; + if (out_fp32) out_fp32[base + i] = __bfloat162float(normed); + } + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + hv_cur[j] = hv_next[j]; + rv_cur[j] = rv_next[j]; + } + } +} +__device__ __forceinline__ bf16 mamba_conv_channel( + const MambaLayer &L, bf16 *state_slot, const bf16 *zrow, int c) { + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = state_slot + (size_t)c * (MB_CONV_K - 1); + const float raw = __bfloat162float(zrow[c]); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + return __float2bfloat16(acc); +} +struct mamba_conv_taps { bf16 v[MB_CONV_K - 1]; float raw; }; +__device__ __forceinline__ mamba_conv_taps mamba_conv_read( + const bf16 *state_slot, const bf16 *zrow, int c) { + const bf16 *state = state_slot + (size_t)c * (MB_CONV_K - 1); + mamba_conv_taps t; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) t.v[k] = state[k]; + t.raw = __bfloat162float(zrow[c]); + return t; +} +__device__ __forceinline__ bf16 mamba_conv_apply( + const MambaLayer &L, bf16 *state_slot, int c, mamba_conv_taps t) { + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = state_slot + (size_t)c * (MB_CONV_K - 1); + const float raw = t.raw; + float acc = __bfloat162float(__ldg(L.conv_b + c)); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(__ldg(w + k)) * __bfloat162float(t.v[k]); + acc += __bfloat162float(__ldg(w + (MB_CONV_K - 1))) * raw; + acc = acc / (1.f + __expf(-acc)); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = t.v[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + return __float2bfloat16(acc); +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map MK2_PROGRESS_PARAMETER(progress_mask)) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + MK2_PROGRESS_GUARD(progress_mask, row); + const int c = idx - row * MB_CONV_DIM; + xbc_act[(size_t)row * MB_CONV_DIM + c] = mamba_conv_channel( + L, L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride, + zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER, c); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map MK2_PROGRESS_PARAMETER(progress_mask)) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + constexpr bool FUSE_GNORM = (split == 1); + extern __shared__ int __shm[]; + bf16 (*const gn_stage)[MB_GROUP_SIZE] = + reinterpret_cast(&__shm[0]); + constexpr bool FUSE_CONV = (MK2_FUSE_CONV != 0); + bf16 (*const conv_x)[MB_GROUP_SIZE] = gn_stage + WARPS; + constexpr int SSM_NSTATE = MB_STATE / 32; + constexpr int SSM_PF = channels < SSM_PREFETCH_MAX ? channels : SSM_PREFETCH_MAX; + constexpr int SSM_GROUP_BYTES = SSM_PF * MB_STATE * (int)sizeof(r9_state_t); + constexpr bool SSM_SMEM_OK = R9_STATE_MODE == 1 && channels == 2 * SSM_PF && + SSM_GROUP_BYTES % 512 == 0; + static_assert(!SSM_SMEM_OK || + 2 * WARPS * MB_GROUP_SIZE * sizeof(bf16) + + WARPS * 2 * SSM_GROUP_BYTES <= STAGES * SLOTS * sizeof(weight_tile)); + char *const ssm_stage = reinterpret_cast(conv_x + WARPS); + const int pairs = BATCH * MB_GROUPS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_GROUPS * split); + MK2_PROGRESS_GUARD(progress_mask, row); + const int group = (pair / split) % MB_GROUPS; + const int first_channel = (pair % split) * channels; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const bool r8_initial = L.r8_initial_ssm && L.r8_initial_ready && + *L.r8_initial_ready == 1 && L.r8_seq_len[row] == L.r8_initial_positions; + r9_state_t *st = reinterpret_cast(L.ssm_state) + + (size_t)slot_map[row] * L.ssm_slot_stride; + char *const warp_ssm = ssm_stage + (size_t)warp * 2 * SSM_GROUP_BYTES; + auto issue_ssm = [&](int head_in_group, int channel0, int buffer) { + const int head = group * MB_HEADS_PER_GROUP + head_in_group; + const r9_state_t *src = st + + (size_t)(head * MB_HEAD_DIM + channel0) * MB_STATE; + const uint32_t dst = (uint32_t)__cvta_generic_to_shared( + warp_ssm + (size_t)buffer * SSM_GROUP_BYTES + lane * 16); + #pragma unroll + for (int chunk = 0; chunk < SSM_GROUP_BYTES / 512; ++chunk) + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: + "r"(dst + (uint32_t)chunk * 512), + "l"(reinterpret_cast(src) + + (size_t)chunk * 512 + lane * 16) : "memory"); + asm volatile("cp.async.commit_group;\n" ::: "memory"); + }; + if constexpr (SSM_SMEM_OK) { + if (!r8_initial) { + issue_ssm(0, 0, 0); + issue_ssm(0, SSM_PF, 1); + } + } + bf16 Braw[MB_STATE / 32]; + float Cv[MB_STATE / 32]; + const bf16 *xrow; + if constexpr (FUSE_CONV) { + bf16 *const cslot = + L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride; + const bf16 *const zrow = + zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER; + constexpr int CONV_BC = MB_STATE / 32; + constexpr int CONV_X = MB_GROUP_SIZE / 32; + mamba_conv_taps tb[CONV_BC], tc[CONV_BC], tx[CONV_X]; + #pragma unroll + for (int j = 0; j < CONV_BC; ++j) { + tb[j] = mamba_conv_read( + cslot, zrow, MB_INNER + group * MB_STATE + lane + j * 32); + tc[j] = mamba_conv_read( + cslot, zrow, MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE + lane + j * 32); + } + #pragma unroll + for (int j = 0; j < CONV_X; ++j) + tx[j] = mamba_conv_read( + cslot, zrow, group * MB_GROUP_SIZE + lane + j * 32); + #pragma unroll + for (int j = 0; j < CONV_BC; ++j) { + Braw[j] = mamba_conv_apply( + L, cslot, MB_INNER + group * MB_STATE + lane + j * 32, + tb[j]); + Cv[j] = __bfloat162float(mamba_conv_apply( + L, cslot, MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE + lane + j * 32, tc[j])); + } + #pragma unroll + for (int j = 0; j < CONV_X; ++j) + conv_x[warp][lane + j * 32] = mamba_conv_apply( + L, cslot, group * MB_GROUP_SIZE + lane + j * 32, tx[j]); + __syncwarp(); + xrow = &conv_x[warp][0]; + } else { + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Braw[j] = Bsec[lane + j * 32]; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + xrow = xa + (size_t)group * MB_GROUP_SIZE; + } + float scalar_dt = 0.f; + float scalar_dA = 0.f; + if (lane < MB_HEADS_PER_GROUP) { + const int scalar_head = group * MB_HEADS_PER_GROUP + lane; + scalar_dt = mamba_softplus( + __bfloat162float(dtrow[scalar_head]) + + L.dt_bias[scalar_head]); + scalar_dA = __expf(L.A[scalar_head] * scalar_dt); + } + #pragma unroll 1 + for (int head_in_group = 0; head_in_group < MB_HEADS_PER_GROUP; + ++head_in_group) { + const int head = group * MB_HEADS_PER_GROUP + head_in_group; + const float dt = __shfl_sync( + 0xffffffffu, scalar_dt, head_in_group); + const float dA = __shfl_sync( + 0xffffffffu, scalar_dA, head_in_group); + const float Dh = L.D[head]; + const float *st_in = L.r8_initial_ssm + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float Bv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) + Bv[j] = __bfloat162float(Braw[j]) * dt; + constexpr int NSTATE = SSM_NSTATE; + constexpr int SSM_PREFETCH = SSM_PF; + static_assert(channels % SSM_PREFETCH == 0); + for (int g0 = 0; g0 < channels; g0 += SSM_PREFETCH) { + float sv[SSM_PREFETCH][NSTATE]; + float xvr[SSM_PREFETCH]; + if (r8_initial) { + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const size_t prow = (size_t)(first_channel + g0 + gi) * MB_STATE; + xvr[gi] = __bfloat162float( + xrow[head_in_group * MB_HEAD_DIM + + first_channel + g0 + gi]); + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = st_in[prow + lane + j * 32]; + } + } else if constexpr (SSM_SMEM_OK) { + asm volatile("cp.async.wait_group 1;\n" ::: "memory"); + __syncwarp(); + const int buffer = g0 / SSM_PREFETCH; + const r9_state_t *sm = reinterpret_cast( + warp_ssm + (size_t)buffer * SSM_GROUP_BYTES); + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + xvr[gi] = __bfloat162float( + xrow[head_in_group * MB_HEAD_DIM + + first_channel + g0 + gi]); + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = r9_load( + sm[(size_t)gi * MB_STATE + lane + j * 32], 1.f); + } + __syncwarp(); + if (head_in_group + 1 < MB_HEADS_PER_GROUP) + issue_ssm(head_in_group + 1, g0, buffer); + else + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } else { + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + xvr[gi] = __bfloat162float( + xrow[head_in_group * MB_HEAD_DIM + + first_channel + g0 + gi]); + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + { + int p = first_channel+g0+gi; + int mapped = R9_STATE_MODE==4 ? L.r9_map[head*24+p] : head*24+p; + int n = lane+j*32; + float scale = R9_STATE_MODE==1 ? 1.f : L.r9_scales[r9_scale_index(mapped/24,mapped%24,n)]; + sv[gi][j] = r9_load(st[(size_t)mapped*128+n],scale); + } + } + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const int p = first_channel + g0 + gi; + const float xv = xvr[gi]; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < NSTATE; ++j) { + const int n = lane + j * 32; + const size_t at = (size_t)p * MB_STATE + n; + const float s = sv[gi][j] * dA + Bv[j] * xv; + int mapped = R9_STATE_MODE==4 ? L.r9_map[head*24+p] : head*24+p; + float scale = R9_STATE_MODE==1 ? 1.f : L.r9_scales[r9_scale_index(mapped/24,mapped%24,n)]; + st[(size_t)mapped*128+n] = r9_store(s,scale); + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) { + const bf16 outv = __float2bfloat16(acc + xv * Dh); + if constexpr (FUSE_GNORM) + gn_stage[warp][head_in_group * MB_HEAD_DIM + p] = outv; + else + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = outv; + } + } + } + } + if constexpr (SSM_SMEM_OK) { + if (!r8_initial) { + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); + __syncwarp(); + } + } + if constexpr (FUSE_GNORM) { + __syncwarp(); + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32], wg[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float gv = __bfloat162float(gn_stage[warp][lane + j * 32]) * + (z / (1.f + __expf(-z))); + wg[j] = __bfloat162float(L.gnorm_w[idx]); + held[j] = gv; + sq += gv * gv; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16(held[j] * inv * wg[j]); + } + __syncwarp(); + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_CONV_DIM; + const int c = idx - t * MB_CONV_DIM; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v; +#ifdef PF_CONV_ZERO_STATE + if (src < 0) v = 0.f; +#else + if (src < 0) v = __bfloat162float(state[pos + k]); +#endif + else v = __bfloat162float( + zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]); + acc += __bfloat162float(w[k]) * v; + } +#ifndef PF_CONV_NO_SILU + acc = acc / (1.f + __expf(-acc)); +#endif + xbc_act[(size_t)t * MB_CONV_DIM + c] = __float2bfloat16(acc); + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + const int units = n_seq * MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int seq = unit / (MB_HEADS * MB_HEAD_DIM); + const int rem = unit - seq * (MB_HEADS * MB_HEAD_DIM); + const int head = rem / MB_HEAD_DIM; + const int p = rem - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p * MB_STATE; + float sreg[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) sreg[j] = st[lane + j * 32]; + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const bf16 *xa = xbc_act + t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + const bf16 *dtrow = zxbcdt + t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + sreg[j] = sreg[j] * dA + __bfloat162float(Bsec[n]) * dt * xv; + acc += sreg[j] * __bfloat162float(Csec[n]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) st[lane + j * 32] = sreg[j]; + } + return; + } +#if PF_SSM_SPLIT_P + const int units = MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; unit += warps_total) { + const int head = unit / MB_HEAD_DIM; + const int p_only = unit - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; +#else + for (int head = blockIdx.x * WARPS + warp; head < MB_HEADS; + head += warps_total) { + const int group = head / MB_HEADS_PER_GROUP; +#endif + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *st = L.ssm_state + (size_t)slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE +#if PF_SSM_SPLIT_P + + (size_t)p_only * MB_STATE +#endif + ; + for (int t = 0; t < n_tokens; ++t) { + const bf16 *xa = xbc_act + (size_t)t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = + zxbcdt + (size_t)t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } +#if PF_SSM_SPLIT_P + { + const int p = p_only; +#else + for (int p = 0; p < MB_HEAD_DIM; ++p) { +#endif + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { +#if PF_SSM_SPLIT_P + const size_t at = (size_t)lane + j * 32; +#else + const size_t at = (size_t)p * MB_STATE + lane + j * 32; +#endif + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } + } +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32], wg[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + wg[j] = __bfloat162float(L.gnorm_w[idx]); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16(held[j] * inv * wg[j]); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len MK2_PROGRESS_PARAMETER(progress_mask)) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + constexpr int Q_PER_WARP = 3; + constexpr int GROUPS_PER_KV = 1; + const int pairs = BATCH * BB_KV_HEADS * GROUPS_PER_KV; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x + warp * gridDim.x; pair < pairs; + pair += gridDim.x * WARPS) { + const int row = pair / (BB_KV_HEADS * GROUPS_PER_KV); + MK2_PROGRESS_GUARD(progress_mask, row); + const int within_row = pair % (BB_KV_HEADS * GROUPS_PER_KV); + const int kv = within_row / GROUPS_PER_KV; + const int q_group = within_row % GROUPS_PER_KV; + const int q_first = q_group * Q_PER_WARP; + const int q_count = q_group == 0 ? Q_PER_WARP : 1; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (q_group == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[h][j] = 0.f; + } + float m[Q_PER_WARP], denom[Q_PER_WARP]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + m[h] = -CUDART_INF_F; + denom[h] = 0.f; + } + constexpr int PF = 8; + const int tpb = L.tokens_per_block; + const size_t elems_per_block = + (size_t)tpb * BB_KV_HEADS * BB_HEAD_DIM; + const int *bt_k = L.block_table + (size_t)(row * 2) * L.max_blocks_per_seq; + const int *bt_v = bt_k + L.max_blocks_per_seq; + const int head_base = kv * tpb * BB_HEAD_DIM; + float qr[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const bf16 *q = base + (kv * kv_per_q + q_first + h) * BB_HEAD_DIM; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qr[h][j] = __bfloat162float(q[lane + j * 32]); + } + for (int page0 = 0; page0 < T; page0 += tpb) { + const int k_page = bt_k[page0 / tpb] & 0x7FFFFFFF; + const int v_page = bt_v[page0 / tpb] & 0x7FFFFFFF; + const bf16 *kbase = L.pool_base + + (size_t)k_page * elems_per_block + + head_base; + const bf16 *vbase = L.pool_base + + (size_t)v_page * elems_per_block + + head_base; + const int tend = page0 + tpb < T ? page0 + tpb : T; + for (int t0 = page0; t0 < tend; t0 += PF) { + const int navail = tend - t0 < PF ? tend - t0 : PF; + float kr[PF][PER_LANE], vr[PF][PER_LANE]; + #pragma unroll + for (int u = 0; u < PF; ++u) { + const int tpos = t0 + (u < navail ? u : navail - 1); +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = vbase + (tpos - page0) * BB_HEAD_DIM; +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new : kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = cur_token + ? v_new : vbase + (tpos - page0) * BB_HEAD_DIM; +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + kr[u][j] = __bfloat162float(kt[i]); + vr[u][j] = __bfloat162float(vt[i]); + } + } + #pragma unroll + for (int u = 0; u < PF; ++u) { + if (u >= navail) break; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + dot += qr[h][j] * kr[u][j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m[h], s); + const float rescale = __expf(m[h] - m_new); + const float w = __expf(s - m_new); + denom[h] = denom[h] * rescale + w; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[h][j] = acc[h][j] * rescale + w * vr[u][j]; + m[h] = m_new; + } + } + } + } + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const float inv = 1.f / denom[h]; + const int head = kv * kv_per_q + q_first + h; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[h][j] * inv); + } + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * BB_KV_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / BB_KV_HEADS; + const int kv = idx - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv); + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv); + for (int i = 0; i < BB_HEAD_DIM; ++i) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + for (int s = 0; s <= pos; ++s) { + const bf16 *kt = kv_slot(L, seq, 0, s) + kv_offset(L, s, kv); + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + const bf16 *vt = kv_slot(L, seq, 1, s) + kv_offset(L, s, kv); + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + wt * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH, bf16 *row_stage = nullptr) { + constexpr int RB = 4; + __shared__ float score[RB][BB_EXPERTS]; + constexpr int CH = D_MODEL / 8; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + for (int row0 = blockIdx.x; row0 < n_rows; row0 += RB * gridDim.x) { + const bf16 *act[RB]; + int rows[RB]; + bool live[RB]; +#pragma unroll + for (int r = 0; r < RB; ++r) { + const int rr = row0 + r * gridDim.x; + live[r] = rr < n_rows; + rows[r] = live[r] ? rr : row0; + act[r] = normed + (size_t)rows[r] * D_MODEL; + } + if (row_stage != nullptr) { + float4 *gdst = reinterpret_cast(row_stage); +#pragma unroll + for (int r = 0; r < RB; ++r) { + const float4 *gsrc = reinterpret_cast(act[r]); + for (int i = threadIdx.x; i < CH; i += THREADS) + gdst[r * CH + i] = gsrc[i]; + } +#pragma unroll + for (int r = 0; r < RB; ++r) act[r] = row_stage + r * D_MODEL; + } + if ((int)threadIdx.x < BB_EXPERTS) { +#pragma unroll + for (int r = 0; r < RB; ++r) + if (live[r]) + moe_w[(size_t)rows[r] * BB_EXPERTS + threadIdx.x] = 0.f; + } + __syncthreads(); + float acc[3][RB]; +#pragma unroll + for (int q = 0; q < 3; ++q) +#pragma unroll + for (int r = 0; r < RB; ++r) acc[q][r] = 0.f; + const bf16 *w0 = gate_w + (size_t)warp * D_MODEL; + const bf16 *w1 = gate_w + (size_t)(warp + WARPS) * D_MODEL; + const bf16 *w2 = gate_w + (size_t)(warp + 2 * WARPS) * D_MODEL; + for (int i = lane; i < D_MODEL; i += 32) { + const float wv0 = __bfloat162float(w0[i]); + const float wv1 = __bfloat162float(w1[i]); + const float wv2 = __bfloat162float(w2[i]); +#pragma unroll + for (int r = 0; r < RB; ++r) { + const float x = __bfloat162float(act[r][i]); + acc[0][r] += x * wv0; + acc[1][r] += x * wv1; + acc[2][r] += x * wv2; + } + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) { +#pragma unroll + for (int q = 0; q < 3; ++q) +#pragma unroll + for (int r = 0; r < RB; ++r) + acc[q][r] += __shfl_xor_sync(0xffffffffu, acc[q][r], off); + } + if (lane == 0) { +#pragma unroll + for (int r = 0; r < RB; ++r) { + score[r][warp] = 1.f / (1.f + __expf(-acc[0][r])); + score[r][warp + WARPS] = 1.f / (1.f + __expf(-acc[1][r])); + score[r][warp + 2 * WARPS] = 1.f / (1.f + __expf(-acc[2][r])); + } + } + __syncthreads(); + const int my_row = row0 + warp * gridDim.x; + if (warp < RB && my_row < n_rows) { + const float *sc = &score[warp][0]; + float remaining = (lane < BB_EXPERTS) + ? sc[lane] + __bfloat162float(gate_bias[lane]) + : -CUDART_INF_F; + int pick[BB_TOP_K]; + float total = 0.f; + const bool any_nan = __any_sync(0xffffffffu, isnan(remaining)); + if (any_nan) { +#pragma unroll + for (int k = 0; k < BB_TOP_K; ++k) { + float bv = remaining; + int bi = lane; +#pragma unroll + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffffu, bv, off); + const int oi = __shfl_xor_sync(0xffffffffu, bi, off); + if (ov > bv || (ov == bv && oi < bi)) { + bv = ov; + bi = oi; + } + } + pick[k] = bi; + total += sc[bi]; + if (lane == bi) remaining = -CUDART_INF_F; + } + } else { + unsigned bits = __float_as_uint(remaining); + if ((bits & 0x7fffffffu) == 0u) bits = 0u; + unsigned remaining_key = + (bits & 0x80000000u) ? ~bits : (bits ^ 0x80000000u); + constexpr unsigned NEG_INF_KEY = 0x007fffffu; +#pragma unroll + for (int k = 0; k < BB_TOP_K; ++k) { + const unsigned best_key = + __reduce_max_sync(0xffffffffu, remaining_key); + const unsigned tied = + __ballot_sync(0xffffffffu, remaining_key == best_key); + const int bi = __ffs(tied) - 1; + pick[k] = bi; + total += sc[bi]; + if (lane == bi) remaining_key = NEG_INF_KEY; + } + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); +#pragma unroll + for (int k = 0; k < BB_TOP_K; ++k) + if (lane == pick[k]) + moe_w[(size_t)my_row * BB_EXPERTS + lane] = sc[lane] * norm; + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert <= BB_EXPERTS; ++expert) + value = __fadd_rn(value, parts[base + (size_t)expert * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ bf16 bb_moe_fold_selected( + const float *parts, const float *weights, int idx) { + const int row = idx / D_MODEL; + const size_t base = (size_t)row * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert < BB_EXPERTS; ++expert) { + const float part = weights[row * BB_EXPERTS + expert] != 0.f + ? parts[base + (size_t)expert * D_MODEL] : 0.f; + value = __fadd_rn(value, part); + } + value = __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ __forceinline__ void sample_argmax_merge( + float &value, int &rank, float other, int other_rank) { + if (other > value || (other == value && other_rank < rank)) { + value = other; + rank = other_rank; + } +} +__device__ __forceinline__ int phoneme_warp_argmax(const bf16 *lg) { + const int lane = threadIdx.x & 31; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = lane; n < BB_PHONEME; n += 32) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, best, off); + const int other_arg = __shfl_down_sync(0xffffffffu, arg, off); + sample_argmax_merge(best, arg, other, other_arg); + } + return arg; +} +__device__ __forceinline__ int sample_block_argmax(float value, int rank) { + __shared__ float warp_value[WARPS]; + __shared__ int warp_rank[WARPS]; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + if (lane == 0) { warp_value[warp] = value; warp_rank[warp] = rank; } + __syncthreads(); + if (warp == 0) { + value = lane < WARPS ? warp_value[lane] : -CUDART_INF_F; + rank = lane < WARPS ? warp_rank[lane] : INT_MAX; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + } + return rank; +} +__device__ __forceinline__ int sample_eos_insertion_rank( + const uint64_t *keys, float eos) { + int lo = 0, hi = TOP_K; + while (lo < hi) { + const int mid = (lo + hi) >> 1; + if (eos > ordered_to_float((uint32_t)(keys[mid] >> 32))) hi = mid; + else lo = mid + 1; + } + return lo; +} +__device__ __forceinline__ void sample_topk_serial( + const uint64_t *keys, float eos, const float *noise, int *out) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float value; + int code; + if (!eos_used && eos > av) { + value = eos; + code = AUDIO_EOS; + eos_used = true; + } else { + value = av; + code = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = value + noise[j]; + if (noised > best_noised) { + best_noised = noised; + best_code = code; + } + } + *out = best_code; +} +__device__ __forceinline__ void sample_topk_parallel( + uint64_t top_key, uint64_t *keys, float eos, + const float *noise, int *out) { + constexpr int TOP_WARPS = (TOP_K + 31) / 32; + __shared__ int warp_nan[TOP_WARPS]; + const int audio_rank = threadIdx.x; + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + float audio_value = -CUDART_INF_F; + bool sorted_nan = false; + if (audio_rank < TOP_K) { + keys[audio_rank] = top_key; + audio_value = ordered_to_float((uint32_t)(top_key >> 32)); + sorted_nan = isnan(audio_value); + } + const unsigned nan_lanes = __ballot_sync(0xffffffffu, sorted_nan); + if (lane == 0 && warp < TOP_WARPS) + warp_nan[warp] = nan_lanes != 0; + const int insertion = __syncthreads_count( + audio_rank < TOP_K && !(eos > audio_value)); + bool any_nan = false; + #pragma unroll + for (int w = 0; w < TOP_WARPS; ++w) any_nan |= warp_nan[w] != 0; + if (any_nan) { + if (threadIdx.x == 0) sample_topk_serial(keys, eos, noise, out); + return; + } + float value = -CUDART_INF_F; + int rank = INT_MAX; + if (audio_rank < TOP_K) { + const int shifted_rank = audio_rank + (audio_rank >= insertion); + if (shifted_rank < TOP_K) { + const float noised = audio_value + noise[shifted_rank]; + if (noised > -CUDART_INF_F) { + value = noised; + rank = shifted_rank; + } + } + } + if (threadIdx.x == 0 && insertion < TOP_K) { + const float noised = eos + noise[insertion]; + if (noised > -CUDART_INF_F) + sample_argmax_merge(value, rank, noised, insertion); + } + const int winner = sample_block_argmax(value, rank); + if (threadIdx.x == 0) { + if (winner == INT_MAX) *out = 0; + else if (winner == insertion) *out = AUDIO_EOS; + else { + const int a = winner - (winner > insertion); + *out = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + } + } +} +constexpr int SORT_ITEMS_PER_THREAD = SORT_N / THREADS; +struct sampler_storage { + uint64_t keys[2][SORT_N]; +}; +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + sampler_storage &storage, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + uint64_t thread_keys[SORT_ITEMS_PER_THREAD]; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + thread_keys[item] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + constexpr int SELECT_N = 128; + int cross_stage = 0; + for (int k = 2; k <= SELECT_N; k <<= 1) { + for (int j = k >> 1; j > 0; j >>= 1) { + if (j >= 32) { + const int bank = cross_stage++ & 1; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + storage.keys[bank][n] = thread_keys[item]; + } + __syncthreads(); + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + const uint64_t other = storage.keys[bank][n ^ j]; + const bool take_hi = ((n & k) == 0) == ((n & j) == 0); + const uint64_t mine = thread_keys[item]; + thread_keys[item] = take_hi + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } else { + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + const uint64_t other = __shfl_xor_sync( + 0xffffffffu, thread_keys[item], j); + const bool take_hi = ((n & k) == 0) == ((n & j) == 0); + const uint64_t mine = thread_keys[item]; + thread_keys[item] = take_hi + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } + } + } + constexpr int run_bank = 1; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + const int run = n / SELECT_N; + const int rank = n & (SELECT_N - 1); + const int descending_rank = (run & 1) ? SELECT_N - 1 - rank : rank; + storage.keys[run_bank][run * SELECT_N + descending_rank] = + thread_keys[item]; + } + __syncthreads(); + uint64_t level4[2]; + #pragma unroll + for (int item = 0; item < 2; ++item) { + const int q = threadIdx.x + item * THREADS; + const int pair = q / SELECT_N; + const int rank = q & (SELECT_N - 1); + const int a = (2 * pair) * SELECT_N + rank; + const int b = (2 * pair + 2) * SELECT_N - 1 - rank; + const uint64_t ka = storage.keys[run_bank][a]; + const uint64_t kb = storage.keys[run_bank][b]; + level4[item] = ka > kb ? ka : kb; + } + for (int j = SELECT_N >> 1, stage = 0; j >= 32; + j >>= 1, ++stage) { + const int bank = stage & 1; + #pragma unroll + for (int item = 0; item < 2; ++item) { + const int q = threadIdx.x + item * THREADS; + storage.keys[bank][q] = level4[item]; + } + __syncthreads(); + #pragma unroll + for (int item = 0; item < 2; ++item) { + const int q = threadIdx.x + item * THREADS; + const uint64_t other = storage.keys[bank][q ^ j]; + const uint64_t mine = level4[item]; + level4[item] = (q & j) == 0 + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } + for (int j = 16; j > 0; j >>= 1) { + #pragma unroll + for (int item = 0; item < 2; ++item) { + const int q = threadIdx.x + item * THREADS; + const uint64_t other = __shfl_xor_sync( + 0xffffffffu, level4[item], j); + const uint64_t mine = level4[item]; + level4[item] = (q & j) == 0 + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } + #pragma unroll + for (int item = 0; item < 2; ++item) + storage.keys[0][threadIdx.x + item * THREADS] = level4[item]; + __syncthreads(); + const int pair2 = threadIdx.x / SELECT_N; + const int rank2 = threadIdx.x & (SELECT_N - 1); + const uint64_t l2a = + storage.keys[0][(2 * pair2) * SELECT_N + rank2]; + const uint64_t l2b = + storage.keys[0][(2 * pair2 + 2) * SELECT_N - 1 - rank2]; + uint64_t level2 = l2a > l2b ? l2a : l2b; + for (int j = SELECT_N >> 1, stage = 0; j >= 32; + j >>= 1, ++stage) { + const int bank = 1 - (stage & 1); + storage.keys[bank][threadIdx.x] = level2; + __syncthreads(); + const uint64_t other = storage.keys[bank][threadIdx.x ^ j]; + const uint64_t mine = level2; + level2 = (threadIdx.x & j) == 0 + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + for (int j = 16; j > 0; j >>= 1) { + const uint64_t other = __shfl_xor_sync( + 0xffffffffu, level2, j); + const uint64_t mine = level2; + level2 = (threadIdx.x & j) == 0 + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + storage.keys[1][threadIdx.x] = level2; + __syncthreads(); + uint64_t top_key = 0; + if (threadIdx.x < SELECT_N) { + const uint64_t ka = storage.keys[1][threadIdx.x]; + const uint64_t kb = storage.keys[1][2 * SELECT_N - 1 - threadIdx.x]; + top_key = ka > kb ? ka : kb; + } + for (int j = SELECT_N >> 1, stage = 0; j >= 32; + j >>= 1, ++stage) { + const int bank = stage & 1; + if (threadIdx.x < SELECT_N) + storage.keys[bank][threadIdx.x] = top_key; + __syncthreads(); + if (threadIdx.x < SELECT_N) { + const uint64_t other = + storage.keys[bank][threadIdx.x ^ j]; + const uint64_t mine = top_key; + top_key = (threadIdx.x & j) == 0 + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } + if (threadIdx.x < SELECT_N) { + for (int j = 16; j > 0; j >>= 1) { + const uint64_t other = __shfl_xor_sync( + 0xffffffffu, top_key, j); + const uint64_t mine = top_key; + top_key = (threadIdx.x & j) == 0 + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } + sample_topk_parallel( + top_key, storage.keys[0], eos_v, + gumbel + ((size_t)row * N_POS + pos) * TOP_K, + out_codes + row); + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void r6_sparse_map(const float* routes,int* map,int* counts + MK2_PROGRESS_PARAMETER(progress_mask)) { + __shared__ int r7_warp_counts[THREADS/32]; + const int lane=threadIdx.x%32, warp=threadIdx.x/32; + const unsigned lower=(1u<maxlim?lim[q]:maxlim; + } + for(int i=i0;i=0; + v[q]=live + ? reinterpret_cast(source)[ + (size_t)row[q]*row_vectors+cv] + : make_uint4(0,0,0,0); + } + #pragma unroll + for(int q=0;q(packed)[(size_t)(e0+q)*BATCH*row_vectors+i]=v[q]; + } + } +} +__device__ void r6_tile(int work,int output_tiles,const int* lut,int &tile,int &tok, + int mtg) { + const int groups=output_tiles/mtg; + tok=lut[1+work/groups]; + const int e=tok/DECODE_TILES; + tile=e*output_tiles+(work%groups)*mtg; +} +__device__ __forceinline__ void mk2_group_tile( + int work, int grouped, int &tile, int &tok, int &here) { + here = work < grouped ? NM_MAX : 1; + const int unit = work < grouped ? work * NM_MAX : grouped * (NM_MAX - 1) + work; + tile = unit / DECODE_TILES; + tok = unit - tile * DECODE_TILES; +} +__device__ __forceinline__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&p_s)[STAGES][SLOTS], out_tile &o_s, + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + int &tile_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = DECODE_TILES, + int rows_valid = DECODE_ROWS, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const int *r6_lut = nullptr,const int *r6_map = nullptr) { + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const bool two_d = !r6_lut && out_tiles_per_expert == 0 + && n_tok == DECODE_TILES && DECODE_TILES % NM_MAX == 0; + int n_tiles = compact_tiles ? __popc(active_mask)*compact_tiles : n_out/OUT_ATOM; + int grouped = 0; + if (two_d) { + const int g = (int)gridDim.x; + const int units = n_tiles * DECODE_TILES; + const int rounds = (units + g - 1) / g; + const int cap = units / NM_MAX; + grouped = g * (rounds / NM_MAX); + if (grouped > cap) grouped = cap; + } + const int r6_output_tiles = r6_lut ? n_out/BB_EXPERTS/OUT_ATOM : 0; + const int r6_mtg = (MT > 1 && r6_lut && r6_output_tiles % MT == 0) ? MT : 1; + if(r6_lut) { + n_tiles=r6_lut[0]*(r6_output_tiles/r6_mtg); + n_tok=1; + } else if(two_d) { n_tiles = grouped + (n_tiles*DECODE_TILES - grouped*NM_MAX); + n_tok=1; } + const int iters = k_total / KCHUNK; + int tile_start = (int)blockIdx.x - tile_io % (int)gridDim.x; + if (tile_start < 0) tile_start += (int)gridDim.x; + tile_io += n_tiles; + int tiles_here = 0; + for (int t = tile_start; t < n_tiles; t += gridDim.x) ++tiles_here; + const int iter_base = iter_io; + iter_io += tiles_here * iters * n_tok; + const int my_tiles = (n_tiles - tile_start + gridDim.x - 1) / gridDim.x; + if (warpgroup::groupid() == 1) { + if (warpgroup::laneid() == 0) { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) { + int tile_id = mk2_expert_tile(work, active_mask, compact_tiles), r6_tok=0; + int nm = 1, mt = 1; + if(r6_lut) { r6_tile(work,r6_output_tiles,r6_lut,tile_id,r6_tok,r6_mtg); + mt = r6_mtg; } + else if(two_d) mk2_group_tile(work,grouped,tile_id,r6_tok,nm); + for (int token_iteration = 0; token_iteration < n_tok; ++token_iteration) + for (int i = 0; i < iters; ++i, ++iter) { + const int tok = (r6_lut || two_d) ? r6_tok : token_iteration; + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(drained[s], phase); + tma::expect_bytes(filled[s], + mt * sizeof(weight_tile) + nm * sizeof(act_tile)); + if (MT == 1 || mt == 1) { + tma::load_async<2, cache_policy::EVICT_FIRST>( + p_s[s][0], w, {tile_id, i}, filled[s]); + } else { + #pragma unroll + for (int c = 0; c < MT; ++c) + tma::load_async<2, cache_policy::EVICT_FIRST>( + p_s[s][c], w, {tile_id + c, i}, filled[s]); + } + const int a_col = out_tiles_per_expert && !r6_lut + ? (tile_id / out_tiles_per_expert) * act_k_tiles_per_expert + i + : i; + if (NM_MAX == 1 || nm == NM_MAX) { + #pragma unroll + for (int m = 0; m < NM_MAX; ++m) + tma::load_async<2, cache_policy::EVICT_LAST>( + p_s[s][mt + m], a, {tok + m, a_col}, filled[s]); + } else { + tma::load_async<2, cache_policy::EVICT_LAST>( + p_s[s][mt], a, {tok, a_col}, filled[s]); + } + } + } + } + } else { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) + for (int token_iteration = 0; token_iteration < n_tok; ++token_iteration) { + int tile_id=mk2_expert_tile(work, active_mask, compact_tiles),tok=token_iteration; + int nm = 1, mt = 1; + if(r6_lut) { r6_tile(work,r6_output_tiles,r6_lut,tile_id,tok,r6_mtg); + mt = r6_mtg; } + else if(two_d) mk2_group_tile(work,grouped,tile_id,tok,nm); + rt_fl<16, BATCH_PAD> accum[NM_MAX]; + if (MT > 1 && mt == MT) { + #pragma unroll + for (int c = 0; c < MT; ++c) warp::zero(accum[c]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], p_s[s][MT], p_s[s][0]); + #pragma unroll + for (int c = 1; c < MT; ++c) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[c], p_s[s][MT], p_s[s][c]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else if (NM_MAX == 1 || nm == NM_MAX) { + #pragma unroll + for (int m = 0; m < NM_MAX; ++m) warp::zero(accum[m]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], p_s[s][1], p_s[s][0]); + #pragma unroll + for (int m = 1; m < NM_MAX; ++m) + warpgroup::mma_ABt, weight_tile, + weight_tile, 0>( + accum[m], p_s[s][1 + m], p_s[s][0]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } else { + warp::zero(accum[0]); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum[0], p_s[s][1], p_s[s][0]); + warpgroup::mma_async_wait(); + if (warpgroup::laneid() == 0) arrive(drained[s], CONSUMER_WARPS); + } + } + #pragma unroll + for (int u = 0; u < NM_MAX; ++u) { + if (u >= mt * nm) break; + const int c = (mt > 1) ? u : 0; + const int m = (mt > 1) ? 0 : u; + const int base = (r6_lut ? (tile_id + c)%r6_output_tiles + : tile_id) * OUT_ATOM; + const int grow0 = (tok + m) * BATCH_PAD; + if (epi == EPI_STORE) { + static_assert(OUT_ATOM == 64 && BATCH_PAD == 64); + const int tid = warpgroup::laneid(); + const int lane = tid & 31; + const int warp = tid >> 5; + const int quad_lane = lane & 3; + #pragma unroll + for (int j = 0; j < 4; ++j) { + uint32_t r0 = pack_bf16_pair( + accum[u].tiles[0][j].data[0]); + uint32_t r1 = pack_bf16_pair( + accum[u].tiles[0][j].data[2]); + uint32_t r2 = pack_bf16_pair( + accum[u].tiles[0][j].data[1]); + uint32_t r3 = pack_bf16_pair( + accum[u].tiles[0][j].data[3]); + uint32_t t = __shfl_xor_sync( + 0xffffffffu, (quad_lane & 1) ? r0 : r1, 1); + if (quad_lane & 1) r0 = t; else r1 = t; + t = __shfl_xor_sync( + 0xffffffffu, (quad_lane & 1) ? r2 : r3, 1); + if (quad_lane & 1) r2 = t; else r3 = t; + t = __shfl_xor_sync( + 0xffffffffu, (quad_lane & 2) ? r0 : r2, 2); + if (quad_lane & 2) r0 = t; else r2 = t; + t = __shfl_xor_sync( + 0xffffffffu, (quad_lane & 2) ? r1 : r3, 2); + if (quad_lane & 2) r1 = t; else r3 = t; + const int row = warp * 16 + (lane >> 2) + + 8 * (quad_lane >> 1); + const int col = 16 * j + 8 * (quad_lane & 1); + if (row < rows_valid) { + const size_t at = + (size_t)(grow0 + row) * dest_stride + base + col; + *reinterpret_cast(dest + at) = + make_uint4(r0, r1, r2, r3); + } + } + continue; + } + if (epi == EPI_BIAS) { + warpgroup::store(o_sf, accum[u]); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + constexpr int OUT_OCTS = OUT_ATOM / 8; + for (int idx = warpgroup::laneid(); idx < OUT_OCTS * rows_valid; + idx += 128) { + const int row = EPI_ROW(idx), col = 8 * EPI_OCT(idx); + const size_t at = (size_t)(grow0 + row) * dest_stride + base + col; + const float4 a0 = *reinterpret_cast( + &o_sf[int2{row, col}]); + const float4 a1 = *reinterpret_cast( + &o_sf[int2{row, col + 4}]); + const float v0 = a0.x + + __bfloat162float(bias[base + col]); + const float v1 = a0.y + + __bfloat162float(bias[base + col + 1]); + const float v2 = a0.z + + __bfloat162float(bias[base + col + 2]); + const float v3 = a0.w + + __bfloat162float(bias[base + col + 3]); + const float v4 = a1.x + + __bfloat162float(bias[base + col + 4]); + const float v5 = a1.y + + __bfloat162float(bias[base + col + 5]); + const float v6 = a1.z + + __bfloat162float(bias[base + col + 6]); + const float v7 = a1.w + + __bfloat162float(bias[base + col + 7]); + *reinterpret_cast(dest + at) = bf16_oct{ + __float2bfloat16(v0), __float2bfloat16(v1), + __float2bfloat16(v2), __float2bfloat16(v3), + __float2bfloat16(v4), __float2bfloat16(v5), + __float2bfloat16(v6), __float2bfloat16(v7)}; + } + warpgroup::sync(1); + continue; + } + warpgroup::store(o_s, accum[u]); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + constexpr int OUT_OCTS = OUT_ATOM / 8; + for (int idx = warpgroup::laneid(); idx < OUT_OCTS * rows_valid; + idx += 128) { + const int row = EPI_ROW(idx); + const int col = 8 * EPI_OCT(idx); + const bf16_oct acc = *reinterpret_cast( + &o_s[int2{row, col}]); + const float v0 = __bfloat162float(acc.x0); + const float v1 = __bfloat162float(acc.x1); + const float v2 = __bfloat162float(acc.x2); + const float v3 = __bfloat162float(acc.x3); + const float v4 = __bfloat162float(acc.x4); + const float v5 = __bfloat162float(acc.x5); + const float v6 = __bfloat162float(acc.x6); + const float v7 = __bfloat162float(acc.x7); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: + *reinterpret_cast(dest + at) = acc; + break; + case EPI_SILU: + *reinterpret_cast(dest + at) = bf16_oct{ + __float2bfloat16(v0 / (1.f + __expf(-v0))), + __float2bfloat16(v1 / (1.f + __expf(-v1))), + __float2bfloat16(v2 / (1.f + __expf(-v2))), + __float2bfloat16(v3 / (1.f + __expf(-v3))), + __float2bfloat16(v4 / (1.f + __expf(-v4))), + __float2bfloat16(v5 / (1.f + __expf(-v5))), + __float2bfloat16(v6 / (1.f + __expf(-v6))), + __float2bfloat16(v7 / (1.f + __expf(-v7)))}; + break; + case EPI_MOE_ACC: { + const int expert = (tile_id + c) / out_tiles_per_expert; + const int local = ((tile_id + c) % out_tiles_per_expert) * OUT_ATOM + + col; + const int original_row=r6_map?r6_map[grow]:grow; + if(original_row<0) break; + const float wt = moe_w[original_row * BB_EXPERTS + expert]; + if (wt != 0.f) { + const size_t dst = + ((size_t)original_row * (BB_EXPERTS + 1) + expert) * + D_MODEL + local; + *reinterpret_cast(moe_out + dst) = make_float4( + __fmul_rn(wt, v0), __fmul_rn(wt, v1), + __fmul_rn(wt, v2), __fmul_rn(wt, v3)); + *reinterpret_cast(moe_out + dst + 4) = make_float4( + __fmul_rn(wt, v4), __fmul_rn(wt, v5), + __fmul_rn(wt, v6), __fmul_rn(wt, v7)); + } + break; + } + case EPI_ACC_F32: + *reinterpret_cast(moe_out + + ((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * + D_MODEL + base + col) = make_float4(v0, v1, v2, v3); + *reinterpret_cast(moe_out + + ((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * + D_MODEL + base + col + 4) = make_float4(v4, v5, v6, v7); + break; + case EPI_BIAS: break; + case EPI_GELU: + *reinterpret_cast(dest + at) = bf16_oct{ + __float2bfloat16(gelu_tanh_fp32(v0)), + __float2bfloat16(gelu_tanh_fp32(v1)), + __float2bfloat16(gelu_tanh_fp32(v2)), + __float2bfloat16(gelu_tanh_fp32(v3)), + __float2bfloat16(gelu_tanh_fp32(v4)), + __float2bfloat16(gelu_tanh_fp32(v5)), + __float2bfloat16(gelu_tanh_fp32(v6)), + __float2bfloat16(gelu_tanh_fp32(v7))}; + break; + case EPI_RESIDUAL: { + const bf16_oct prior = + *reinterpret_cast(residual + at); + *reinterpret_cast(residual + at) = bf16_oct{ + __float2bfloat16(__bfloat162float(prior.x0) + v0), + __float2bfloat16(__bfloat162float(prior.x1) + v1), + __float2bfloat16(__bfloat162float(prior.x2) + v2), + __float2bfloat16(__bfloat162float(prior.x3) + v3), + __float2bfloat16(__bfloat162float(prior.x4) + v4), + __float2bfloat16(__bfloat162float(prior.x5) + v5), + __float2bfloat16(__bfloat162float(prior.x6) + v6), + __float2bfloat16(__bfloat162float(prior.x7) + v7)}; + break; + } + } + } + warpgroup::sync(1); + } + } + (void)my_tiles; + } +} +template +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2_impl( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr, + const unsigned char *native_predictor_active_mask = nullptr + MK2_PROGRESS_PARAMETER(native_progress_mask)) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&p_s)[STAGES][SLOTS] = alloc.allocate(); + out_tile_f (&o_sf) = alloc.allocate(); + out_tile &o_s = *reinterpret_cast(&o_sf); + bf16 *const route_stage = reinterpret_cast(&o_sf); + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_SCRATCH_STRIDE]; + static_assert(STAGES * SLOTS * sizeof(weight_tile) >= sizeof(sampler_storage)); + sampler_storage &sort_storage = *reinterpret_cast(&__shm[0]); + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS, 0); + arrive(drained[s], CONSUMER_WARPS); + } + } + int gemm_iter = 0; + int gemm_tile = 0; + __syncthreads(); + const bool prefilling = MK2_PF && (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + const float *fold_src = nullptr, *fold_routes = nullptr; + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + fold_src = nullptr; + fold_routes = nullptr; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + grid_barrier(g.barrier); +#endif + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); +#endif + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (MK2_BISECTION_HOOKS && bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, BATCH, + route_stage); + } + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + } + grid_barrier(g.barrier); + r6_sparse_map(bb->moe_w,bb->r6_map,bb->r6_counts + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + r6_sparse_gather(bb->normed,bb->r6_packed,bb->r6_map,bb->r6_counts, + bb->r6_lut); + grid_barrier(g.barrier); + gemm(L.exp_up, bb->r6_views[0], BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0,0,nullptr,nullptr,DECODE_TILES,DECODE_ROWS, + nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->r6_views[1], BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + DECODE_TILES, DECODE_ROWS, nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + if (MK2_BISECTION_HOOKS && bb->layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_selected(bb->moe_out, bb->moe_w, idx); + grid_barrier(g.barrier); + } else { + fold_src = bb->moe_out; + fold_routes = bb->moe_w; + } + } + if (MK2_BISECTION_HOOKS && bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if constexpr (MK2_PF) { + if (bb && bb->n_prefill_tokens > 0) { + const int n_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = n_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + bb->pf_h[idx] = bb->prompt[idx]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = bb->pf_n_seq > 0 ? bb->pf_n_seq : 1; + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_gated, + T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (MK2_BISECTION_HOOKS && bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, T, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * (BB_EXPERTS + 1) * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->pf_normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->pf_moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = bb_moe_fold_deterministic(bb->pf_moe_out, idx); + grid_barrier(g.barrier); + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, bb->pf_final, ln_scratch, + n_total); + grid_barrier(g.barrier); + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } + } + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, bb->phoneme_b); + grid_barrier(g.barrier); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + for (int row = blockIdx.x + warp * gridDim.x; row < BATCH; + row += gridDim.x * WARPS) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + MK2_PROGRESS_GUARD(native_progress_mask, row); + const int arg = phoneme_warp_argmax(lg); + if (lane == 0) bb->phoneme_out[row] = arg; + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + __shared__ int native_run_predictor; + if (native_predictor_active_mask && bb && g.native_codes) { + if (threadIdx.x == 0) { + int any_active = 0; + for (int row = 0; row < BATCH; ++row) + any_active |= int(native_predictor_active_mask[row] != 0 && + MK2_PROGRESS_ENABLED(native_progress_mask, row)); + native_run_predictor = any_active; + } + __syncthreads(); + } + const bool skip_native_predictor = native_predictor_active_mask && bb && + g.native_codes && native_run_predictor == 0; + if (skip_native_predictor) { + for (int index = blockIdx.x * THREADS + threadIdx.x; + index < BATCH * N_POS; index += gridDim.x * THREADS) { + MK2_PROGRESS_GUARD(native_progress_mask, index / N_POS); + g.sampled[index] = 0; + g.native_codes[index] = 0; + } + grid_barrier(g.barrier); + } else { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + constexpr int DPL = D_HEAD / 32; + float qv[DPL], kn[DPL], vn[DPL]; + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + qv[j] = __bfloat162float(q[i]); + kn[j] = __bfloat162float(k_new[i]); + vn[j] = __bfloat162float(v_new[i]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float my_sc = 0.f; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += qv[j] * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == t) my_sc = acc * scale; + } + { + float acc = 0.f; + #pragma unroll + for (int j = 0; j < DPL; ++j) acc += qv[j] * kn[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == pos) my_sc = acc * scale; + } + float m = __shfl_sync(0xffffffff, my_sc, 0); +#pragma unroll 4 + for (int t = 1; t <= pos; ++t) + m = fmaxf(m, __shfl_sync(0xffffffff, my_sc, t)); + float my_e = __expf(my_sc - m); + float s = 0.f; +#pragma unroll 4 + for (int t = 0; t <= pos; ++t) + s += __shfl_sync(0xffffffff, my_e, t); + const float inv = 1.f / s; + const float my_p = my_e * inv; + float ov[DPL] = {}; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *vt = g.vcache + cache_base + t * D_HEAD; + const float s = __shfl_sync(0xffffffff, my_p, t); + #pragma unroll + for (int j = 0; j < DPL; ++j) + ov[j] += s * __bfloat162float(vt[lane + j * 32]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + ov[j] += __shfl_sync(0xffffffff, my_p, pos) * vn[j]; + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + + lane + j * 32] = __float2bfloat16(ov[j]); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, + g.sampled_pos, sort_storage, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif +#ifdef ABLATE_SAMPLE + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { + __syncthreads(); + } else { + grid_barrier(g.barrier); + } +#endif + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x == 0) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + g.sampled[row * N_POS + pos] = code; + if (g.native_codes && MK2_PROGRESS_ENABLED(native_progress_mask, row)) { +#if MK2_INFLIGHT + g.native_codes[row * N_POS + pos] = + native_predictor_active_mask && !native_predictor_active_mask[row] + ? 0 : (long long)code; +#else + g.native_codes[row * N_POS + pos] = (long long)code; +#endif + } + } + } +#ifdef ABLATE_SAMPLE + grid_barrier(g.barrier); +#else + if (pos + 1 == n_pos) grid_barrier(g.barrier); +#endif + } + } + if (g.native_seq_len && blockIdx.x == 0) + for (int row = threadIdx.x; row < BATCH; row += blockDim.x) { + MK2_PROGRESS_GUARD(native_progress_mask, row); + g.native_seq_len[row] += 1; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +} +#define predictor_mk2 predictor_mk2_impl +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2_startup( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr, + const unsigned char *native_predictor_active_mask = nullptr + MK2_PROGRESS_PARAMETER(native_progress_mask)) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&p_s)[STAGES][SLOTS] = alloc.allocate(); + out_tile_f (&o_sf) = alloc.allocate(); + out_tile &o_s = *reinterpret_cast(&o_sf); + bf16 *const route_stage = reinterpret_cast(&o_sf); + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_SCRATCH_STRIDE]; + static_assert(STAGES * SLOTS * sizeof(weight_tile) >= sizeof(sampler_storage)); + sampler_storage &sort_storage = *reinterpret_cast(&__shm[0]); + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS, 0); + arrive(drained[s], CONSUMER_WARPS); + } + } + int gemm_iter = 0; + int gemm_tile = 0; + __syncthreads(); + constexpr bool prefilling = false; + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + const float *fold_src = nullptr, *fold_routes = nullptr; + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + fold_src = nullptr; + fold_routes = nullptr; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + grid_barrier(g.barrier); +#endif + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#if !MK2_FUSE_CONV + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); +#endif + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (MK2_BISECTION_HOOKS && bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, BATCH, + route_stage); + } + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + } + grid_barrier(g.barrier); + r6_sparse_map(bb->moe_w,bb->r6_map,bb->r6_counts + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + r6_sparse_gather(bb->normed,bb->r6_packed,bb->r6_map,bb->r6_counts, + bb->r6_lut); + grid_barrier(g.barrier); + gemm(L.exp_up, bb->r6_views[0], BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0,0,nullptr,nullptr,DECODE_TILES,DECODE_ROWS, + nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->r6_views[1], BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + DECODE_TILES, DECODE_ROWS, nullptr,0,bb->r6_lut,bb->r6_map); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + if (MK2_BISECTION_HOOKS && bb->layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_selected(bb->moe_out, bb->moe_w, idx); + grid_barrier(g.barrier); + } else { + fold_src = bb->moe_out; + fold_routes = bb->moe_w; + } + } + if (MK2_BISECTION_HOOKS && bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if constexpr (false) { + const int n_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = n_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + bb->pf_h[idx] = bb->prompt[idx]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = bb->pf_n_seq > 0 ? bb->pf_n_seq : 1; + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_gated, + T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (MK2_BISECTION_HOOKS && bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, T, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * (BB_EXPERTS + 1) * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->pf_normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->pf_moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = bb_moe_fold_deterministic(bb->pf_moe_out, idx); + grid_barrier(g.barrier); + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, bb->pf_final, ln_scratch, + n_total); + grid_barrier(g.barrier); + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch, + BATCH, nullptr, -1, nullptr, nullptr, fold_src, fold_routes); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, bb->phoneme_b); + grid_barrier(g.barrier); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + for (int row = blockIdx.x + warp * gridDim.x; row < BATCH; + row += gridDim.x * WARPS) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + MK2_PROGRESS_GUARD(native_progress_mask, row); + const int arg = phoneme_warp_argmax(lg); + if (lane == 0) bb->phoneme_out[row] = arg; + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + __shared__ int native_run_predictor; + if (native_predictor_active_mask && bb && g.native_codes) { + if (threadIdx.x == 0) { + int any_active = 0; + for (int row = 0; row < BATCH; ++row) + any_active |= int(native_predictor_active_mask[row] != 0 && + MK2_PROGRESS_ENABLED(native_progress_mask, row)); + native_run_predictor = any_active; + } + __syncthreads(); + } + const bool skip_native_predictor = native_predictor_active_mask && bb && + g.native_codes && native_run_predictor == 0; + if constexpr (true) { + for (int index = blockIdx.x * THREADS + threadIdx.x; + index < BATCH * N_POS; index += gridDim.x * THREADS) { + MK2_PROGRESS_GUARD(native_progress_mask, index / N_POS); + g.sampled[index] = 0; + g.native_codes[index] = 0; + } + grid_barrier(g.barrier); + } else { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + constexpr int DPL = D_HEAD / 32; + float qv[DPL], kn[DPL], vn[DPL]; + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + qv[j] = __bfloat162float(q[i]); + kn[j] = __bfloat162float(k_new[i]); + vn[j] = __bfloat162float(v_new[i]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float my_sc = 0.f; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += qv[j] * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == t) my_sc = acc * scale; + } + { + float acc = 0.f; + #pragma unroll + for (int j = 0; j < DPL; ++j) acc += qv[j] * kn[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == pos) my_sc = acc * scale; + } + float m = __shfl_sync(0xffffffff, my_sc, 0); +#pragma unroll 4 + for (int t = 1; t <= pos; ++t) + m = fmaxf(m, __shfl_sync(0xffffffff, my_sc, t)); + float my_e = __expf(my_sc - m); + float s = 0.f; +#pragma unroll 4 + for (int t = 0; t <= pos; ++t) + s += __shfl_sync(0xffffffff, my_e, t); + const float inv = 1.f / s; + const float my_p = my_e * inv; + float ov[DPL] = {}; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *vt = g.vcache + cache_base + t * D_HEAD; + const float s = __shfl_sync(0xffffffff, my_p, t); + #pragma unroll + for (int j = 0; j < DPL; ++j) + ov[j] += s * __bfloat162float(vt[lane + j * 32]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + ov[j] += __shfl_sync(0xffffffff, my_p, pos) * vn[j]; + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + + lane + j * 32] = __float2bfloat16(ov[j]); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + p_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, + g.sampled_pos, sort_storage, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif +#ifdef ABLATE_SAMPLE + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { + __syncthreads(); + } else { + grid_barrier(g.barrier); + } +#endif + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x == 0) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + g.sampled[row * N_POS + pos] = code; + if (g.native_codes && MK2_PROGRESS_ENABLED(native_progress_mask, row)) { +#if MK2_INFLIGHT + g.native_codes[row * N_POS + pos] = + native_predictor_active_mask && !native_predictor_active_mask[row] + ? 0 : (long long)code; +#else + g.native_codes[row * N_POS + pos] = (long long)code; +#endif + } + } + } +#ifdef ABLATE_SAMPLE + grid_barrier(g.barrier); +#else + if (pos + 1 == n_pos) grid_barrier(g.barrier); +#endif + } + } + if (g.native_seq_len && blockIdx.x == 0) + for (int row = threadIdx.x; row < BATCH; row += blockDim.x) { + MK2_PROGRESS_GUARD(native_progress_mask, row); + g.native_seq_len[row] += 1; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 64 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * SLOTS * sizeof(weight_tile) + + sizeof(out_tile_f) + 2048; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; + const unsigned char *native_predictor_active_mask = nullptr; + MK2_PROGRESS_MEMBER + int maximum_grid = 0; + bool grid_reset_pending = false; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * SLOTS * sizeof(weight_tile) + + sizeof(out_tile_f) + 2048), 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaFuncSetAttribute(predictor_mk2_startup, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaFuncSetAttribute(predictor_mk2_impl, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = std::min(out->sms * out->blocks_per_sm, 4096); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 64 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->maximum_grid = info.sms * info.blocks_per_sm; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + mk2_check(owned.alloc(&hb.r6_packed,(size_t)BB_EXPERTS*BATCH*D_MODEL*sizeof(bf16))); + mk2_check(owned.alloc(&hb.r6_map,(size_t)BB_EXPERTS*BATCH*sizeof(int))); + mk2_check(owned.alloc(&hb.r6_counts,BB_EXPERTS*sizeof(int))); + mk2_check(owned.alloc(&hb.r6_lut,(1+BB_EXPERTS*DECODE_TILES)*sizeof(int))); + act_gl sparse_views[]={ + act_gl{hb.r6_packed,nullptr,nullptr,(size_t)BB_EXPERTS*BATCH,(size_t)D_MODEL}, + act_gl{moe_h,nullptr,nullptr,(size_t)BB_EXPERTS*BATCH,(size_t)BB_MOE_INTER}}; + mk2_check(owned.alloc(&hb.r6_views,sizeof(sparse_views))); + mk2_check(cudaMemcpy(hb.r6_views,sparse_views,sizeof(sparse_views),cudaMemcpyHostToDevice)); + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.maximum_grid = pred->maximum_grid; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native, bool startup = false) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(s, &capture_status)); + if (st->pred.grid_reset_pending || capture_status != cudaStreamCaptureStatusNone) { + mk2_check(cudaMemsetAsync(st->pred.g.barrier, 0, 64 * sizeof(unsigned int), s)); + st->pred.grid_reset_pending = false; + } + if (startup) { + predictor_mk2_startup<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS, nullptr, + st->pred.native_predictor_active_mask + MK2_PROGRESS_ARGUMENT(st->pred.native_progress_mask)); + } else { + predictor_mk2_impl<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS, nullptr, + native ? st->pred.native_predictor_active_mask : nullptr + MK2_PROGRESS_ARGUMENT(native ? st->pred.native_progress_mask : nullptr)); + } + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" int mk2_run_native_startup_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream) { + return mk2_enqueue_stage0(p, b, stream, true, true); +} +extern "C" int mk2_set_native_predictor_active_mask(MK2Predictor *p, + const unsigned char *mask) { + try { + if (!p || reinterpret_cast(p)->kind != 1) + throw std::runtime_error("startup mask requires a full Stage0 handle"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && st->pred.native_predictor_active_mask != mask) + throw std::runtime_error("startup mask pointer cannot change after first Stage0 enqueue"); + if (mask) + mk2_stage0_device_pointer(mask, st->pred.device, "native_predictor_active_mask"); + st->pred.native_predictor_active_mask = mask; + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown startup mask binding failure"; return 1; } +} +extern "C" int mk2_native_startup_skip_abi_version() { return 1; } +#include "../../../inflight/runtime.cuh" +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "native_prefill_impl.inc" +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *p, int grid) { + if (!p) { g_last_error = "grid setter requires full Stage0 handle"; return 1; } + auto *st = reinterpret_cast(p); + if (st->kind != 1 || grid < 1 || grid > st->maximum_grid) { + g_last_error = "requested Stage0 grid outside independently resident capacity"; + return 1; + } + if (st->grid != grid) st->grid_reset_pending = true; + st->grid = grid; + g_last_error.clear(); + return 0; +} +extern "C" int mk2_r8_bind_initial_state(MK2Predictor *p, const void *const *initial, const void *ready, const void *seq_len, int positions) { + try { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || st->bindings_initialized || !initial || !ready || !seq_len || positions <= 0) + throw std::runtime_error("invalid immutable initial-state binding"); + for (int i = 0; i < 14; ++i) { + mk2_stage0_device_pointer(initial[i], st->pred.device, "initial state"); + st->mb_host[i].r8_initial_ssm = static_cast(initial[i]); + st->mb_host[i].r8_initial_ready = static_cast(ready); + st->mb_host[i].r8_seq_len = static_cast(seq_len); + st->mb_host[i].r8_initial_positions = positions; + } + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" int mk2_r9_state_bind(void* p,const void*const* scales,const void*const* maps) {auto*st=(MK2Stage0State*)p;if(st->bindings_initialized)return 1;for(int i=0;i<14;++i){st->mb_host[i].r9_scales=(const float*)scales[i];st->mb_host[i].r9_map=(const int*)maps[i];}return 0;} +extern "C" int mk2_r9_state_bytes(){return sizeof(r9_state_t);} diff --git a/deploy/nemotron_native/kernels/b512/prefill.cu b/deploy/nemotron_native/kernels/b512/prefill.cu new file mode 100644 index 0000000..9b2dcb2 --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/prefill.cu @@ -0,0 +1,20 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace b512_prefill { +#include "native/prefill/prefill.cu" +} diff --git a/deploy/nemotron_native/kernels/b512/stage0.cu b/deploy/nemotron_native/kernels/b512/stage0.cu new file mode 100644 index 0000000..edd33b9 --- /dev/null +++ b/deploy/nemotron_native/kernels/b512/stage0.cu @@ -0,0 +1,35 @@ +#define MK2_BATCH 512 +#define STAGES 3 +#define KCHUNK 128 +#define MK2_NO_MAIN 1 +#define MK2_DIRECT_INPUT 1 +#define MK2_B1_LOCAL_BARRIERS 1 +#define MK2_SKIP_HIDDEN_SNAPSHOT 1 +#define MK2_ACTIVE_EXPERTS 1 +#define MK2_SSM_SPLITS 1 +#define NMULT 2 +#define MK2_OUT_MULT 2 +#define SSM_PREFETCH_MAX 12 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#if MK2_INFLIGHT +#include +#endif +#include +#include +#include +#include +#include +#include +#include +#include +namespace b512_stage0 { +#include "native/stage0/predictor_mk2.cu" +} diff --git a/deploy/nemotron_native/kernels/codec/base/kernel_kf_codec.h b/deploy/nemotron_native/kernels/codec/base/kernel_kf_codec.h new file mode 100644 index 0000000..9a8b4d5 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/base/kernel_kf_codec.h @@ -0,0 +1,13 @@ +#pragma once +#include +struct CodecArgs { + const void *inp[41]; + const void *st[26]; + void *ns[26]; + void *waveform; + void *scratch; + int B; + int F; +}; +void codec_launch(const CodecArgs &a, void *stream); +size_t codec_scratch_bytes(int B, int F); diff --git a/deploy/nemotron_native/kernels/codec/base/legacy_codec.cuh b/deploy/nemotron_native/kernels/codec/base/legacy_codec.cuh new file mode 100644 index 0000000..cb128c4 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/base/legacy_codec.cuh @@ -0,0 +1,1388 @@ +#include "kernel_kf_codec.h" +#include +#include +#include +#include +using namespace nvcuda; +template __device__ void prepare_attribute(F,int) {} +#ifndef TC2_QPT +#define TC2_QPT 4 +#endif +#define DIVUP(a,b) (((a) + (b) - 1) / (b)) +__device__ __forceinline__ float half_snake(float v, int c, int half, + const float *__restrict__ alpha) { + if (c < half) { + const float a = alpha[c]; + const float s = __sinf(a * v); + return __fmaf_rn(s, s / (a + 1e-9f), v); + } + return v > 0.f ? v : 0.01f * v; +} +#ifndef WG_ST3 +#define WG_ST3 1 +#define WG_MB3 3 +#define WG_ST7 1 +#define WG_MB7 3 +#define WG_BN3 128 +#define WG_BN7 128 +#endif +#define PDL_ENTER() do { (void)0; if (threadIdx.x == 0) (void)0; } while (0) +#define PDL_TRIGGER() do { if (threadIdx.x == 0) (void)0; } while (0) +template +__device__ __forceinline__ void pdl_launch(KFn kern, dim3 grid, dim3 blk, size_t shmem, + cudaStream_t s, KArgs... args) { + native_invoke(kern, grid, blk, shmem, args...); +} +template +__device__ __forceinline__ void pdl_launch_pf(KFn kern, dim3 grid, dim3 blk, size_t shmem, + cudaStream_t s, const void *pf, size_t pfb, KArgs... args) { + native_invoke_pf(kern, grid, blk, shmem, pf, pfb, args...); +} +#include "wg_kernel.inc" +#include "wg2_kernel.inc" +#include "wg3_kernel.inc" +#include "wg4_step_widths.inc" +#include "wg5_step_widths.inc" +#include "repack3.inc" +#include "tail_kernel.inc" +template +__device__ __noinline__ void dense_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + PDL_TRIGGER(); + constexpr int BK = 16 * KT; + constexpr int NT = NWARPS * 32; + constexpr int WTM = BM / WM; + constexpr int WTN = BN / WN; + constexpr int FM = WTM / 16, FN = WTN / 16; + constexpr int LDA = BK + 8, LDB = BN + 8; + constexpr int GEMM_SM = BM * LDA + BK * LDB; + constexpr int EPI_SM = NWARPS * WTM * WTN * 2; + constexpr int SM_H = GEMM_SM > EPI_SM ? GEMM_SM : EPI_SM; + __half *smem=reinterpret_cast<__half*>(native_arena()); + char *colok=reinterpret_cast(native_aux()); + __half *As = smem; + __half *Bs = smem + BM * LDA; + float *Cw = reinterpret_cast(smem); + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int wm = warp / WN, wn = warp % WN; + const int m0 = blockIdx.x * BM; + const int n0 = blockIdx.y * BN; + const int H = KT - 1; + const int Kg = Cin * KT; + if (tid < BN) { + const int n = n0 + tid; + colok[tid] = (n < N && (n % L) < T) ? 1 : 0; + } + wmma::fragment acc[FM][FN]; +#pragma unroll + for (int i = 0; i < FM; ++i) +#pragma unroll + for (int j = 0; j < FN; ++j) wmma::fill_fragment(acc[i][j], 0.f); + constexpr int ATOT = BM * (BK / 8); + constexpr int AREG = (ATOT + NT - 1) / NT; + constexpr int BTOT = 16 * (BN / 8); + constexpr int BREG = (BTOT + NT - 1) / NT; + int4 ar[AREG]; + int4 br0[BREG], br1[BREG]; +#define CODEC_LOADG(kg0) { const int _ci0 = (kg0) / KT; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * NT; ar[r] = make_int4(0, 0, 0, 0); if (ATOT % NT == 0 || idx < ATOT) { const int m = idx / (BK / 8), v = idx % (BK / 8); const int co = m0 + m, kg = (kg0) + v * 8; if (co < Cout && kg < Kg) ar[r] = *reinterpret_cast(w + (size_t)co * Kg + kg); } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * NT; br0[r] = make_int4(0, 0, 0, 0); br1[r] = make_int4(0, 0, 0, 0); if (BTOT % NT == 0 || idx < BTOT) { const int cil = idx / (BN / 8), nc = idx % (BN / 8); const int ci = _ci0 + cil; if (ci < Cin) { const int4 *src = reinterpret_cast( in + (size_t)ci * SCin + n0 + nc * 8); br0[r] = src[0]; br1[r] = src[1]; } } } } + CODEC_LOADG(0) + for (int kg0 = 0; kg0 < Kg; kg0 += BK) { + __syncthreads(); +#pragma unroll + for (int r = 0; r < AREG; ++r) { + const int idx = tid + r * NT; + if (ATOT % NT == 0 || idx < ATOT) + *reinterpret_cast(As + (idx / (BK / 8)) * LDA + + (idx % (BK / 8)) * 8) = ar[r]; + } +#pragma unroll + for (int r = 0; r < BREG; ++r) { + const int idx = tid + r * NT; + if (BTOT % NT != 0 && idx >= BTOT) continue; + const int cil = idx / (BN / 8), nc = idx % (BN / 8); + unsigned int u[8]; + u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; + u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; +#pragma unroll + for (int k = 0; k < KT; ++k) { + const int kh = k >> 1; + int4 o; + if (k & 1) { + o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); + o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); + o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); + o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); + } else { + o.x = u[kh + 0]; o.y = u[kh + 1]; + o.z = u[kh + 2]; o.w = u[kh + 3]; + } + *reinterpret_cast(Bs + (cil * KT + k) * LDB + nc * 8) = o; + } + } + __syncthreads(); + if (kg0 + BK < Kg) CODEC_LOADG(kg0 + BK) +#pragma unroll + for (int kf = 0; kf < BK; kf += 16) { + wmma::fragment af[FM]; + wmma::fragment bf[FN]; +#pragma unroll + for (int i = 0; i < FM; ++i) + wmma::load_matrix_sync(af[i], As + (wm * WTM + i * 16) * LDA + kf, LDA); +#pragma unroll + for (int j = 0; j < FN; ++j) + wmma::load_matrix_sync(bf[j], Bs + kf * LDB + wn * WTN + j * 16, LDB); +#pragma unroll + for (int i = 0; i < FM; ++i) +#pragma unroll + for (int j = 0; j < FN; ++j) + wmma::mma_sync(acc[i][j], af[i], bf[j], acc[i][j]); + } + } +#undef CODEC_LOADG + __syncthreads(); + float *Cme = Cw + warp * WTM * WTN; +#pragma unroll + for (int i = 0; i < FM; ++i) +#pragma unroll + for (int j = 0; j < FN; ++j) + wmma::store_matrix_sync(Cme + (i * 16) * WTN + j * 16, acc[i][j], WTN, + wmma::mem_row_major); + __syncwarp(); + const int half = Cout >> 1; + const int lane = tid & 31; + for (int e = lane; e < WTM * WTN; e += 32) { + const int mm = e / WTN, nn = e % WTN; + const int co = m0 + wm * WTM + mm; + const int nl = wn * WTN + nn; + if (co >= Cout || !colok[nl]) continue; + const int n = n0 + nl; + float v = Cme[e] + bias[co]; + if (resid != nullptr) + v += __half2float(resid[(size_t)co * SCout + n + H]); + if (do_snake) v = half_snake(v, co, half, alpha); + out[(size_t)co * SCout + n + H] = __float2half(v); + } + if (blockIdx.y == 0) { + const int nbx = gridDim.x; + if (ns != nullptr) { + const int B = N / L; + const int total = Cin * B * H; + for (int idx = blockIdx.x * NT + tid; idx < total; idx += nbx * NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int B = N / L; + const int total = Cout * B * pad_len; + for (int idx = blockIdx.x * NT + tid; idx < total; idx += nbx * NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} +template +__device__ __noinline__ void transpose_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + __half *__restrict__ out, const __half *__restrict__ pad_src, + __half *__restrict__ ns, int Cin, int Cout, int Tin, + int Lin, int SCin, int jbase, int Lout, int SCout, int Hout, int pad_off, + int pad_len, int B) { + if constexpr (PDL) PDL_ENTER(); + constexpr int Ktaps = 2 * STRIDE; + constexpr int QPT = QPT_; + constexpr int VPT = QPT * STRIDE; + static_assert(QPT % 2 == 0, "half2 input loads need an even frame count"); + static_assert(VPT % 2 == 0, "half2 staging needs an even slot"); + float *ws=reinterpret_cast(native_aux()); + __half *sm=reinterpret_cast<__half*>(native_arena()); + const int co = blockIdx.y; + const int b = blockIdx.z; + const int tid = threadIdx.x, BLK = blockDim.x; + for (int i = tid; i < G * Ktaps; i += BLK) + ws[i] = __half2float(w[(size_t)(co * G + i / Ktaps) * Ktaps + (i % Ktaps)]); + __syncthreads(); + const float bs = bias[0 + co]; + const int half = Cout >> 1; + const bool snake = co < half; + const float av = snake ? alpha[co] : 0.f; + const float ri = snake ? __frcp_rn(av + 1e-9f) : 0.f; + const __half *ic = in + (size_t)(co * G) * SCin + b * Lin + jbase; + __half *oc = out + (size_t)co * SCout + b * Lout + Hout; + constexpr bool VEC16 = (VPT % 8) == 0; + const int ophase = + VEC16 ? 0 : (int)((((uintptr_t)oc) >> 1) & 7u); + __half *const sbase = sm + ophase; + const int hd = (8 - ophase) & 7; + const int lane = tid & 31; + for (int qb = 1 + blockIdx.x * BLK * QPT; qb <= Tin; + qb += gridDim.x * BLK * QPT) { + const int q = qb + QPT * tid; + const bool a0 = (q <= Tin); + float jq[QPT][G], jp0[G]; +#pragma unroll + for (int g = 0; g < G; ++g) { +#pragma unroll + for (int e2 = 0; e2 < QPT / 2; ++e2) { + const __half2 v = a0 ? *reinterpret_cast( + ic + (size_t)g * SCin + q + 2 * e2) + : __floats2half2_rn(0.f, 0.f); + jq[2 * e2][g] = __half2float(v.x); + jq[2 * e2 + 1][g] = __half2float(v.y); + } + } +#pragma unroll + for (int g = 0; g < G; ++g) { + const float up = __shfl_up_sync(0xffffffffu, jq[QPT - 1][g], 1); + jp0[g] = (lane == 0) + ? (a0 ? __half2float(ic[(size_t)g * SCin + q - 1]) : 0.f) + : up; + } + __half *const sp = sbase + tid * VPT; + __align__(16) __half vreg[VPT]; +#pragma unroll + for (int r = 0; r < STRIDE; ++r) { + float wc[G], wp[G]; +#pragma unroll + for (int g = 0; g < G; ++g) { + wc[g] = ws[g * Ktaps + r]; + wp[g] = ws[g * Ktaps + r + STRIDE]; + } +#pragma unroll + for (int e = 0; e < QPT; ++e) { + float v = bs; +#pragma unroll + for (int g = 0; g < G; ++g) { + const float cur = jq[e][g]; + const float prv = (e == 0) ? jp0[g] : jq[e - 1][g]; + v = __fmaf_rn(wc[g], cur, __fmaf_rn(wp[g], prv, v)); + } + if (snake) { + const float sn = __sinf(av * v); + v = __fmaf_rn(sn * sn, ri, v); + } else { + v = v > 0.f ? v : 0.01f * v; + } + vreg[e * STRIDE + r] = __float2half(v); + } + } + if constexpr (VEC16) { +#pragma unroll + for (int m = 0; m < VPT / 8; ++m) + *reinterpret_cast(sp + m * 8) = + *reinterpret_cast(vreg + m * 8); + } else { +#pragma unroll + for (int i = 0; i < VPT; ++i) sp[i] = vreg[i]; + } + __syncthreads(); + const int nval = min(BLK * QPT, Tin - qb + 1); + const int tot = nval * STRIDE; + __half *dst = oc + (size_t)(qb - 1) * STRIDE; + if constexpr (VEC16) { + const int tot2 = tot >> 1; + for (int i = tid; i < tot2; i += BLK) + reinterpret_cast<__half2 *>(dst)[i] = + reinterpret_cast(sbase)[i]; + if ((tot & 1) && tid == 0) dst[tot - 1] = sbase[tot - 1]; + } else { + const int head = hd < tot ? hd : tot; + for (int i = tid; i < head; i += BLK) dst[i] = sbase[i]; + const int nb4 = (tot - head) >> 3; + const int4 *s4 = reinterpret_cast(sbase + head); + int4 *d4 = reinterpret_cast(dst + head); + for (int i = tid; i < nb4; i += BLK) d4[i] = s4[i]; + for (int i = head + (nb4 << 3) + tid; i < tot; i += BLK) + dst[i] = sbase[i]; + } + __syncthreads(); + } + if (blockIdx.x == 0) { + if (ns != nullptr) { + for (int ci = co + tid * Cout; ci < Cin; ci += BLK * Cout) + ns[(size_t)b * Cin + ci] = + in[(size_t)ci * SCin + b * Lin + jbase + native_true_tokens(b,Tin)]; + } + if (pad_src != nullptr) { + for (int j = tid; j < pad_len; j += BLK) + out[(size_t)co * SCout + b * Lout + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } +} +template +__device__ __noinline__ void transpose_conv_early_bpack_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + __half *__restrict__ out, const __half *__restrict__ pad_src, + __half *__restrict__ ns, int Cin, int Cout, int Tin, + int Lin, int SCin, int jbase, int Lout, int SCout, int Hout, int pad_off, + int pad_len, int B) { + if constexpr (PDL) PDL_ENTER(); + static_assert(BPACK <= 32); + static_assert(LPB == 8 || LPB == 16 || LPB == 32); + constexpr int Ktaps = 2 * STRIDE; + constexpr int WARP = 32; + constexpr int NT = WARP * BPACK; + constexpr int BPB = WARP / LPB; + float *ws=reinterpret_cast(native_aux()); + __half *sm=reinterpret_cast<__half*>(native_arena()); + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int lane = tid & 31; + const int sub = lane % LPB; + const int bsub = lane / LPB; + const int co = blockIdx.y; + const int b = blockIdx.z * (BPACK * BPB) + warp * BPB + bsub; + const bool active = b < B; + for (int i = tid; i < G * Ktaps; i += NT) + ws[i] = __half2float(w[(size_t)(co * G + i / Ktaps) * Ktaps + + (i % Ktaps)]); + __syncthreads(); + const int qb = 1 + blockIdx.x * LPB; + const int q = qb + sub; + __half *sp = sm + (warp * WARP + lane) * STRIDE; + const bool live = active && q <= Tin; + { + const float bs = bias[co]; + const int half = Cout >> 1; + const bool snake = co < half; + const float av = snake ? alpha[co] : 0.f; + const float ri = snake ? __frcp_rn(av + 1e-9f) : 0.f; + const __half *ic = in + (size_t)(co * G) * SCin + b * Lin + jbase; + float jq[G], jp[G]; +#pragma unroll + for (int g = 0; g < G; ++g) + jq[g] = live ? __half2float(ic[(size_t)g * SCin + q]) : 0.f; +#pragma unroll + for (int g = 0; g < G; ++g) { + const float up = __shfl_up_sync(0xffffffffu, jq[g], 1, LPB); + jp[g] = (sub == 0) + ? (live ? __half2float(ic[(size_t)g * SCin + q - 1]) : 0.f) + : up; + } + if (live) +#pragma unroll + for (int r = 0; r < STRIDE; ++r) { + float v = bs; +#pragma unroll + for (int g = 0; g < G; ++g) + v = __fmaf_rn(ws[g * Ktaps + r], jq[g], + __fmaf_rn(ws[g * Ktaps + r + STRIDE], jp[g], v)); + if (snake) { + const float sn = __sinf(av * v); + v = __fmaf_rn(sn * sn, ri, v); + } else { + v = v > 0.f ? v : 0.01f * v; + } + sp[r] = __float2half(v); + } + } + __syncwarp(); + if (active) { + const int nval = min(LPB, Tin - (qb - 1)); + const int tot2 = (nval * STRIDE) >> 1; + const __half *src = sm + (warp * WARP + bsub * LPB) * STRIDE; + __half *dst = out + (size_t)co * SCout + b * Lout + Hout + + (size_t)(qb - 1) * STRIDE; + for (int i = sub; i < tot2; i += LPB) + reinterpret_cast<__half2 *>(dst)[i] = + reinterpret_cast(src)[i]; + if ((nval * STRIDE) & 1) + if (sub == 0) dst[nval * STRIDE - 1] = src[nval * STRIDE - 1]; + } + if (blockIdx.x == 0 && active) { + if (ns != nullptr) { + for (int ci = co + sub * Cout; ci < Cin; ci += LPB * Cout) + ns[(size_t)b * Cin + ci] = + in[(size_t)ci * SCin + b * Lin + jbase + native_true_tokens(b,Tin)]; + } + if (pad_src != nullptr) { + for (int j = sub; j < pad_len; j += LPB) + out[(size_t)co * SCout + b * Lout + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } +} +template +__device__ __noinline__ void transpose_conv_bpack2_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + __half *__restrict__ out, const __half *__restrict__ pad_src, + __half *__restrict__ ns, int Cin, int Cout, int Tin, + int Lin, int SCin, int jbase, int Lout, int SCout, int Hout, int pad_off, + int pad_len, int B) { + if constexpr (PDL) PDL_ENTER(); + static_assert(BPACK <= 32); + static_assert(LPB == 8 || LPB == 16 || LPB == 32); + constexpr int Ktaps = 2 * STRIDE; + constexpr int WARP = 32; + constexpr int NT = WARP * BPACK; + constexpr int BPB = WARP / LPB; + constexpr int QPT = 2; + constexpr int VPT = QPT * STRIDE; + float *ws=reinterpret_cast(native_aux()); + __half *sm=reinterpret_cast<__half*>(native_arena()); + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int lane = tid & 31; + const int sub = lane % LPB; + const int bsub = lane / LPB; + const int co = blockIdx.y; + const int b = blockIdx.z * (BPACK * BPB) + warp * BPB + bsub; + const bool active = b < B; + for (int i = tid; i < G * Ktaps; i += NT) + ws[i] = __half2float(w[(size_t)(co * G + i / Ktaps) * Ktaps + + (i % Ktaps)]); + __syncthreads(); + const int qb = 1 + blockIdx.x * (LPB * QPT); + const int q = qb + sub * QPT; + __half *sp = sm + (warp * WARP + lane) * VPT; + const bool a0 = active && (q <= Tin); + { + const float bs = bias[co]; + const int half = Cout >> 1; + const bool snake = co < half; + const float av = snake ? alpha[co] : 0.f; + const float ri = snake ? __frcp_rn(av + 1e-9f) : 0.f; + const __half *ic = in + (size_t)(co * G) * SCin + b * Lin + jbase; + float jq0[G], jq1[G], jp0[G]; +#pragma unroll + for (int g = 0; g < G; ++g) { + const __half2 v = a0 ? *reinterpret_cast( + ic + (size_t)g * SCin + q) + : __floats2half2_rn(0.f, 0.f); + jq0[g] = __half2float(v.x); + jq1[g] = __half2float(v.y); + } +#pragma unroll + for (int g = 0; g < G; ++g) { + const float up = __shfl_up_sync(0xffffffffu, jq1[g], 1, LPB); + jp0[g] = (sub == 0) + ? (a0 ? __half2float(ic[(size_t)g * SCin + q - 1]) + : 0.f) + : up; + } + if (a0) { + float vv[VPT]; +#pragma unroll + for (int r = 0; r < STRIDE; ++r) { + float wc[G], wp[G]; +#pragma unroll + for (int g = 0; g < G; ++g) { + wc[g] = ws[g * Ktaps + r]; + wp[g] = ws[g * Ktaps + r + STRIDE]; + } +#pragma unroll + for (int e = 0; e < QPT; ++e) { + float v = bs; +#pragma unroll + for (int g = 0; g < G; ++g) { + const float cur = (e == 0) ? jq0[g] : jq1[g]; + const float prv = (e == 0) ? jp0[g] : jq0[g]; + v = __fmaf_rn(wc[g], cur, __fmaf_rn(wp[g], prv, v)); + } + if (snake) { + const float sn = __sinf(av * v); + v = __fmaf_rn(sn * sn, ri, v); + } else { + v = v > 0.f ? v : 0.01f * v; + } + vv[e * STRIDE + r] = v; + } + } + __half2 *sp2 = reinterpret_cast<__half2 *>(sp); +#pragma unroll + for (int i = 0; i < VPT / 2; ++i) + sp2[i] = __floats2half2_rn(vv[2 * i], vv[2 * i + 1]); + } + } + __syncwarp(); + if (active) { + const int nval = min(LPB * QPT, Tin - (qb - 1)); + const int tot = nval * STRIDE; + const __half *src = sm + (warp * WARP + bsub * LPB) * VPT; + __half *dst = out + (size_t)co * SCout + b * Lout + Hout + + (size_t)(qb - 1) * STRIDE; + for (int i = sub; i < (tot >> 1); i += LPB) + reinterpret_cast<__half2 *>(dst)[i] = + reinterpret_cast(src)[i]; + if (tot & 1) + if (sub == 0) dst[tot - 1] = src[tot - 1]; + } + if (blockIdx.x == 0 && active) { + if (ns != nullptr) { + for (int ci = co + sub * Cout; ci < Cin; ci += LPB * Cout) + ns[(size_t)b * Cin + ci] = + in[(size_t)ci * SCin + b * Lin + jbase + native_true_tokens(b,Tin)]; + } + if (pad_src != nullptr) { + for (int j = sub; j < pad_len; j += LPB) + out[(size_t)co * SCout + b * Lout + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } +} +#define TCONV_BPACK(G,STRIDE,BPACK,LPB,TIN,COUT,BB,...) pdl_launch(transpose_conv_early_bpack_kernel, dim3(DIVUP(TIN, LPB), COUT, DIVUP(BB, BPACK * (32 / LPB))), dim3(32 * BPACK), (size_t)(32 * (BPACK)) * (STRIDE) * 2 + 64, s, __VA_ARGS__) +#define TCONV_BP2(G,STRIDE,LPB,TIN,COUT,BB,...) pdl_launch(transpose_conv_bpack2_kernel, dim3(DIVUP(TIN, (LPB) * 2), COUT, DIVUP(BB, 32)), dim3(32 * (LPB)), (size_t)(32 * (LPB)) * (4 * (STRIDE)) + 64, s, __VA_ARGS__) +#define TCONV_BPACK_DISPATCH1(G,STRIDE,BPACK,TIN,COUT,BB,...) do { if ((TIN) >= 32) TCONV_BP2(G, STRIDE, 16, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) >= 16) TCONV_BP2(G, STRIDE, 8, TIN, COUT, BB, __VA_ARGS__); else TCONV_BPACK(G, STRIDE, BPACK, 16, TIN, COUT, BB, __VA_ARGS__); } while (0) +#define TCONV_BPACK_DISPATCH(G,STRIDE,BPACK,TIN,COUT,BB,...) do { if ((TIN) >= 32) TCONV_BP2(G, STRIDE, 16, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) >= 16) TCONV_BP2(G, STRIDE, 8, TIN, COUT, BB, __VA_ARGS__); else TCONV_BPACK(G, STRIDE, BPACK, 16, TIN, COUT, BB, __VA_ARGS__); } while (0) +#define TCONV_BPACK_DISPATCH8(G,STRIDE,BPACK,TIN,COUT,BB,...) do { if ((TIN) >= 32) TCONV_BP2(G, STRIDE, 16, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) >= 16) TCONV_BP2(G, STRIDE, 8, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) <= 8) TCONV_BPACK(G, STRIDE, BPACK, 8, TIN, COUT, BB, __VA_ARGS__); else TCONV_BPACK(G, STRIDE, BPACK, 16, TIN, COUT, BB, __VA_ARGS__); } while (0) +__device__ __noinline__ void pack_latent_kernel(const __half *__restrict__ latent, + const __half *__restrict__ state, + __half *__restrict__ out, int Cin, int T, + int L, int SC, int H, int B) { + const int total = Cin * B * L; + for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; + idx += gridDim.x * blockDim.x) { + const int j = idx % L, r = idx / L; + const int b = r % B, ci = r / B; + out[(size_t)ci * SC + b * L + j] = + (j < H) ? state[((size_t)b * Cin + ci) * H + j] + : latent[((size_t)b * Cin + ci) * T + (j - H)]; + } +} +#define PC_VEC 8 +__device__ __noinline__ void post_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, __half *__restrict__ y, + __half *__restrict__ ns, int Cin, int T, int L, int SC, int KT, int B, + int N) { + float *pw=reinterpret_cast(native_arena()); + for (int i = threadIdx.x; i < Cin * KT; i += 256) pw[i] = __half2float(w[i]); + __syncthreads(); + const int H = KT - 1; + const float bs = bias[0]; + for (int n = (blockIdx.x * 256 + threadIdx.x) * PC_VEC; n < N; + n += gridDim.x * 256 * PC_VEC) { + float v[PC_VEC]; +#pragma unroll + for (int i = 0; i < PC_VEC; ++i) v[i] = bs; + for (int ci = 0; ci < Cin; ++ci) { + const __half *xr = in + (size_t)ci * SC + n; + const float *wr = pw + ci * KT; + const int4 lo = *reinterpret_cast(xr); + const int4 hi = *reinterpret_cast(xr + 8); + unsigned u[8]; + u[0] = lo.x; u[1] = lo.y; u[2] = lo.z; u[3] = lo.w; + u[4] = hi.x; u[5] = hi.y; u[6] = hi.z; u[7] = hi.w; + float x[PC_VEC + 6]; +#pragma unroll + for (int j = 0; j < (PC_VEC + 6) / 2; ++j) { + const float2 f = __half22float2( + *reinterpret_cast(&u[j])); + x[2 * j] = f.x; + x[2 * j + 1] = f.y; + } +#pragma unroll + for (int i = 0; i < PC_VEC; ++i) + for (int k = 0; k < 7; ++k) + if (k < KT) v[i] = __fmaf_rn(wr[k], x[i + k], v[i]); + } + const int b0 = n / L, t0 = n - (n / L) * L; +#pragma unroll + for (int i = 0; i < PC_VEC; ++i) { + int t = t0 + i, b = b0; + if (t >= L) { t -= L; b = b0 + 1; } + if (t < T && b < B) + y[(size_t)b * T + t] = + __float2half(fminf(fmaxf(v[i], -1.f), 1.f)); + } + } + if (ns != nullptr) { + const int total = B * Cin * H; + for (int idx = blockIdx.x * 256 + threadIdx.x; idx < total; + idx += gridDim.x * 256) { + const int j = idx % H, r = idx / H; + const int ci = r % Cin, b = r / Cin; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SC + b * L + native_true_tokens(b,T) + j]; + } + } +} +#define RP_JOBS 6 +struct RepackSet { + const __half *src[RP_JOBS]; + __half *dst[RP_JOBS]; + int Cout[RP_JOBS], CoutP[RP_JOBS], Cin[RP_JOBS]; + int base[RP_JOBS], nch[RP_JOBS], rs[RP_JOBS]; + int k3base; + const __half *k3src[2]; + __half *k3dst[2]; +}; +__host__ __device__ inline int native_RP_COUT(int j) {const int v[6]={768, 384, 256, 128, 64, 32};return v[j];} +__host__ __device__ inline int native_RP_CIN(int j) {const int v[6]={384, 768, 128, 256, 32, 64};return v[j];} +__host__ __device__ inline int native_RP_COUTP(int j) {const int v[6]={768, 384, 256, 128, 128, 128};return v[j];} +template +__device__ __noinline__ void repack_kernel(RepackSet js) { + const int bx = blockIdx.x; + int job, job_bx; + if (bx < 1152) { + job = 0; job_bx = 0; + } else if (bx < 2304) { + job = 1; job_bx = 1152; + } else if (bx < 2432) { + job = 2; job_bx = 2304; + } else if (bx < 2560) { + job = 3; job_bx = 2432; + } else if (bx < 2576) { + job = 4; job_bx = 2560; + } else { + job = 5; job_bx = 2576; + } + const int Cout = js.Cout[job], CoutP = js.CoutP[job], Cin = js.Cin[job]; + const int idx = (bx - job_bx) * blockDim.x + threadIdx.x; + const int co = idx / Cin, ci = idx - co * Cin; + const int c = ci / CIB, cl = ci - c * CIB; + __half *d = js.dst[job] + (size_t)c * (KT * CoutP * CIB) + + (size_t)(co >> 3) * (CIB * 8) + (cl >> 3) * 64 + (co & 7) * 8 + + (cl & 7); + const __half *sp = js.src[job] + (size_t)(co * Cin + ci) * KT; + const bool live = co < Cout; +#pragma unroll + for (int k = 0; k < KT; ++k) + d[(size_t)k * CoutP * CIB] = live ? sp[k] : __float2half(0.f); +} +template +__device__ __noinline__ void repack_rs_kernel( + const RepackSet &js, const __half *__restrict__ latent, const __half *__restrict__ latent_state, + __half *__restrict__ latent_out, int latent_C, int latent_T, int latent_L, + int latent_SC, int latent_H, int latent_B) { + PDL_TRIGGER(); + constexpr int CO_TILE = 16; + constexpr int CI_TILE = 64; + constexpr int ROW_H = CI_TILE * KT + 8; + constexpr int ROW4 = CI_TILE * KT / 8; + __half *tile=reinterpret_cast<__half*>(native_arena()); + const int bx = blockIdx.x; + { + const int idx = bx + threadIdx.x * gridDim.x; + const int total = latent_C * latent_B * latent_L; + if (idx < total) { + const int j = idx % latent_L, r = idx / latent_L; + const int b = r % latent_B, ci = r / latent_B; + latent_out[(size_t)ci * latent_SC + b * latent_L + j] = + (j < latent_H) + ? latent_state[((size_t)b * latent_C + ci) * latent_H + j] + : latent[((size_t)b * latent_C + ci) * latent_T + + (j - latent_H)]; + } + } + if (bx >= js.k3base) { + constexpr int K3 = 3; + constexpr int R3_H = CI_TILE * K3 + 8; + constexpr int R3_4 = CI_TILE * K3 / 8; + const int r = bx - js.k3base; + const int conv = r / 1152; + const int loc = r - conv * 1152; + const int Cout3 = conv ? 768 : 1536; + const int Cin3 = conv ? 1536 : 768; + const int nci3 = Cin3 / CI_TILE; + const int cot3 = loc / nci3; + if (cot3 * CO_TILE >= Cout3) return; + const int co03 = cot3 * CO_TILE; + const int ci03 = (loc - cot3 * nci3) * CI_TILE; + const int NCH3 = Cin3 >> 4; + const __half *const s3 = js.k3src[conv]; + __half *const d3 = js.k3dst[conv]; + const int t3 = threadIdx.x; +#pragma unroll + for (int n = 0; n < CO_TILE * R3_4 / 128; ++n) { + const int i = t3 + n * 128; + const int co = i / R3_4, q = i - co * R3_4; + reinterpret_cast(tile + co * R3_H)[q] = + *reinterpret_cast( + s3 + ((size_t)(co03 + co) * Cin3 + ci03) * K3 + q * 8); + } + __syncthreads(); + const int w3 = t3 >> 5, l3 = t3 & 31; + const int rA3 = l3 >> 2, q43 = l3 & 3; + const int cb3 = co03 >> 4; + __half *const db3 = d3 + ((size_t)(cb3 >> 2) * NCH3 * K3) * 1024 + + ((cb3 & 3) * 32 + l3) * 8; +#pragma unroll + for (int u = w3; u < (CI_TILE / 16) * K3; u += 4) { + const int cl = u / K3, tap = u - cl * K3; + const int cgi = (ci03 >> 4) + cl; + const __half *sa = tile + rA3 * R3_H + (cl * 16 + 2 * q43) * K3 + tap; + const __half *sb = sa + 8 * R3_H; + uint32_t o[4]; + o[0] = __half_as_ushort(sa[0]) | + ((uint32_t)__half_as_ushort(sa[K3]) << 16); + o[1] = __half_as_ushort(sb[0]) | + ((uint32_t)__half_as_ushort(sb[K3]) << 16); + o[2] = __half_as_ushort(sa[8 * K3]) | + ((uint32_t)__half_as_ushort(sa[9 * K3]) << 16); + o[3] = __half_as_ushort(sb[8 * K3]) | + ((uint32_t)__half_as_ushort(sb[9 * K3]) << 16); + *reinterpret_cast(db3 + (size_t)(cgi * K3 + tap) * 1024) = + make_int4(o[0], o[1], o[2], o[3]); + } + return; + } + int job = 0; +#pragma unroll + for (int q = 1; q < RP_JOBS; ++q) if (bx >= js.base[q]) job = q; + const int local = bx - js.base[job]; + const int Cout = js.Cout[job], Cin = js.Cin[job]; + const int nci = (Cin + CI_TILE - 1) / CI_TILE; + const int cot = local / nci; + const int co0 = cot * CO_TILE; + const int ci0 = (local - cot * nci) * CI_TILE; + const int NCH = js.nch[job]; + const int t = threadIdx.x; + const int valid4 = (min(Cin - ci0, CI_TILE) * KT) >> 3; + for (int i = t; i < CO_TILE * ROW4; i += 128) { + const int co = i / ROW4, q = i - co * ROW4; + int4 v = make_int4(0, 0, 0, 0); + if (co0 + co < Cout && q < valid4) + v = *reinterpret_cast( + js.src[job] + ((size_t)(co0 + co) * Cin + ci0) * KT + q * 8); + reinterpret_cast(tile + co * ROW_H)[q] = v; + } + __syncthreads(); + const int warp = t >> 5, lane = t & 31; + const int rA = lane >> 2, q4 = lane & 3; + const int cb = co0 >> 4; + const int mtile = cb >> 2, wsub = cb & 3; + __half *const dbase = js.dst[job] + + ((size_t)mtile * NCH * KT) * 1024 + (wsub * 32 + lane) * 8; + if (!js.rs[job]) { + const int CoutP = js.CoutP[job]; + constexpr int UN2 = 16 * (CI_TILE / 8) * KT; +#pragma unroll + for (int n = 0; n < UN2 / 128; ++n) { + const int u = t + n * 128; + const int col = u & 15, cig = (u >> 4) & (CI_TILE / 8 - 1); + const int tap = u / (16 * (CI_TILE / 8)); + const int gci = ci0 + cig * 8; + if (gci >= Cin) continue; + const int gco = co0 + col; + uint32_t o[4]; + const __half *sp = tile + col * ROW_H; +#pragma unroll + for (int e = 0; e < 4; ++e) + o[e] = __half_as_ushort(sp[(cig * 8 + 2 * e) * KT + tap]) | + ((uint32_t)__half_as_ushort( + sp[(cig * 8 + 2 * e + 1) * KT + tap]) << 16); + __half *d = js.dst[job] + + (size_t)(gci >> 4) * ((size_t)KT * CoutP * 16) + + (size_t)tap * (CoutP * 16) + (size_t)(gco >> 3) * 128 + + ((gci >> 3) & 1) * 64 + (gco & 7) * 8; + *reinterpret_cast(d) = make_int4(o[0], o[1], o[2], o[3]); + } + return; + } + constexpr int UNITS = (CI_TILE / 16) * KT; +#pragma unroll + for (int u = warp; u < UNITS; u += 4) { + const int cl = u / KT, tap = u - cl * KT; + const int cgi = (ci0 >> 4) + cl; + if (cgi >= NCH) continue; + const __half *sa = tile + rA * ROW_H + (cl * 16 + 2 * q4) * KT + tap; + const __half *sb = sa + 8 * ROW_H; + uint32_t o[4]; + o[0] = __half_as_ushort(sa[0]) | + ((uint32_t)__half_as_ushort(sa[KT]) << 16); + o[1] = __half_as_ushort(sb[0]) | + ((uint32_t)__half_as_ushort(sb[KT]) << 16); + o[2] = __half_as_ushort(sa[8 * KT]) | + ((uint32_t)__half_as_ushort(sa[9 * KT]) << 16); + o[3] = __half_as_ushort(sb[8 * KT]) | + ((uint32_t)__half_as_ushort(sb[9 * KT]) << 16); + *reinterpret_cast(dbase + (size_t)(cgi * KT + tap) * 1024) = + make_int4(o[0], o[1], o[2], o[3]); + } +} +__device__ __noinline__ void repack3_rs_kernel( + const __half *__restrict__ win, const __half *__restrict__ wsk, + __half *__restrict__ din, __half *__restrict__ dsk, + const __half *__restrict__ asrc, __half *__restrict__ adst, int SCA, + int ncol, int ntile) { + PDL_TRIGGER(); + if (blockIdx.y == 12) { + if (adst == nullptr) return; + for (int u = blockIdx.x; u < 96 * ntile; u += gridDim.x) { + const int g = u / ntile; + const int col = (u - g * ntile) * 128 + (int)threadIdx.x; + if (col >= ncol) continue; + __half v[8]; +#pragma unroll + for (int h = 0; h < 8; ++h) + v[h] = asrc[(size_t)(g * 8 + h) * SCA + col]; + *reinterpret_cast(adst + (size_t)g * ((size_t)SCA * 8) + + (size_t)col * 8) = + *reinterpret_cast(v); + } + return; + } + constexpr int KT = 3; + constexpr int CO_TILE = 16; + constexpr int CI_TILE = 64; + constexpr int ROW_H = CI_TILE * KT + 8; + constexpr int ROW4 = CI_TILE * KT / 8; + __half *tile=reinterpret_cast<__half*>(native_arena()); + const int jy = blockIdx.y; + const int conv = jy & 1, b = 5 - (jy >> 1); + const int Cout = conv ? 768 : 1536; + const int Cin = conv ? 1536 : 768; + const __half *src = (conv ? wsk : win) + (size_t)(b + 1) * Cout * Cin * KT; + __half *dst = (conv ? dsk : din) + (size_t)b * ((size_t)Cout * Cin * KT); + const int nci = Cin / CI_TILE; + const int cot = blockIdx.x / nci; + const int co0 = cot * CO_TILE; + const int ci0 = (blockIdx.x - cot * nci) * CI_TILE; + const int NCH = Cin >> 4; + const int t = threadIdx.x; +#pragma unroll + for (int n = 0; n < CO_TILE * ROW4 / 128; ++n) { + const int i = t + n * 128; + const int co = i / ROW4, q = i - co * ROW4; + reinterpret_cast(tile + co * ROW_H)[q] = + *reinterpret_cast( + src + ((size_t)(co0 + co) * Cin + ci0) * KT + q * 8); + } + __syncthreads(); + const int warp = t >> 5, lane = t & 31; + const int rA = lane >> 2, q4 = lane & 3; + const int cb = co0 >> 4; + __half *const dbase = dst + ((size_t)(cb >> 2) * NCH * KT) * 1024 + + ((cb & 3) * 32 + lane) * 8; +#pragma unroll + for (int u = warp; u < (CI_TILE / 16) * KT; u += 4) { + const int cl = u / KT, tap = u - cl * KT; + const int cgi = (ci0 >> 4) + cl; + const __half *sa = tile + rA * ROW_H + (cl * 16 + 2 * q4) * KT + tap; + const __half *sb = sa + 8 * ROW_H; + uint32_t o[4]; + o[0] = __half_as_ushort(sa[0]) | + ((uint32_t)__half_as_ushort(sa[KT]) << 16); + o[1] = __half_as_ushort(sb[0]) | + ((uint32_t)__half_as_ushort(sb[KT]) << 16); + o[2] = __half_as_ushort(sa[8 * KT]) | + ((uint32_t)__half_as_ushort(sa[9 * KT]) << 16); + o[3] = __half_as_ushort(sb[8 * KT]) | + ((uint32_t)__half_as_ushort(sb[9 * KT]) << 16); + *reinterpret_cast(dbase + (size_t)(cgi * KT + tap) * 1024) = + make_int4(o[0], o[1], o[2], o[3]); + } +} +__device__ __noinline__ void repack_activation_only(const __half *asrc,__half *adst,int SCA,int ncol,int ntile){ + if (adst == nullptr) return; + for (int u = blockIdx.x; u < 96 * ntile; u += gridDim.x) { + const int g = u / ntile; + const int col = (u - g * ntile) * 128 + (int)threadIdx.x; + if (col >= ncol) continue; + __half v[8]; +#pragma unroll + for (int h = 0; h < 8; ++h) + v[h] = asrc[(size_t)(g * 8 + h) * SCA + col]; + *reinterpret_cast(adst + (size_t)g * ((size_t)SCA * 8) + + (size_t)col * 8) = + *reinterpret_cast(v); + } + return; +} +namespace { +struct Stack { + int T, L, N, SC, H; +}; +inline __host__ __device__ int tc_blk(int Tin, int qpt = 2) { + int best = 256, bw = 1 << 30; + for (int b = 256; b >= 32; b >>= 1) { + const int per = qpt * b; + const int waste = (Tin + per - 1) / per * per - Tin; + if (waste < bw) { bw = waste; best = b; } + } + return best; +} +inline __host__ __device__ int roundup(int a, int b) { return (a + b - 1) / b * b; } +__host__ __device__ Stack mk(int B, int T, int K) { + Stack s; + s.T = T; + s.H = K - 1; + s.L = s.H + T; + s.N = B * s.L; + s.SC = roundup(s.N, 128) + 128; + return s; +} +struct Plan { + Stack s0, s1, s2, s3, s4; + __half *LB, *X0, *M0, *Y0; + __half *A, *Ac, *Bb, *M1; + __half *P, *Q, *M2; + __half *R, *S, *M3; + __half *U; + __half *W7[RP_JOBS]; + __half *W3i, *W3s; + size_t total; +}; +static constexpr int kTuneB = 32; +inline __host__ __device__ int nref(const Stack &st) { return kTuneB * st.L; } +inline __host__ __device__ bool use_rp3i(int, int F) { return (size_t)kTuneB * (4 * F + 2) > 400; } +inline __host__ __device__ bool use_rp3s(int, int F) { return (size_t)kTuneB * (4 * F + 2) > 400; } +__host__ __device__ Plan build_plan(int B, int F, __half *base) { + Plan p; + p.s0 = mk(B, 2 * F, 3); + p.s1 = mk(B, 4 * F, 3); + p.s2 = mk(B, 36 * F, 7); + p.s3 = mk(B, 252 * F, 7); + p.s4 = mk(B, 1764 * F, 7); + size_t off = 0; + auto alloc = [&](int ch, const Stack &s) -> __half * { + __half *r = base ? base + off : nullptr; + off += (size_t)ch * s.SC; + return r; + }; + p.LB = alloc(40, p.s0); + p.X0 = alloc(768, p.s0); + p.M0 = alloc(1536, p.s0); + p.Y0 = alloc(768, p.s0); + p.A = alloc(768, p.s1); + p.Ac = use_rp3i(B, F) ? alloc(768, p.s1) : nullptr; + p.Bb = alloc(768, p.s1); + p.M1 = alloc(1536, p.s1); + p.P = alloc(384, p.s2); + p.Q = alloc(384, p.s2); + p.M2 = alloc(768, p.s2); + p.R = alloc(128, p.s3); + p.S = alloc(128, p.s3); + p.M3 = alloc(256, p.s3); + p.U = alloc(32, p.s4); + for (int j = 0; j < RP_JOBS; ++j) { + p.W7[j] = base ? base + off : nullptr; + off += (size_t)native_RP_COUTP(j) * native_RP_CIN(j) * 7; + } + p.W3i = nullptr; + if (use_rp3i(B, F)) { + p.W3i = base ? base + off + (size_t)3 * 1536 * 768 : nullptr; + off += (size_t)7 * 3 * 1536 * 768; + } + p.W3s = nullptr; + if (use_rp3s(B, F)) { + p.W3s = base ? base + off + (size_t)3 * 768 * 1536 : nullptr; + off += (size_t)7 * 3 * 768 * 1536; + } + p.total = off; + return p; +} +__host__ __device__ const __half *H16(const void *p) { return static_cast(p); } +__host__ __device__ const float *F32(const void *p) { return static_cast(p); } +} +size_t codec_scratch_bytes(int B, int F) { + return build_plan(B, F, nullptr).total * sizeof(__half); +} +template +__device__ __forceinline__ void wg_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, bool prep = false) { + constexpr int WG_BM = 64 * NWG, WG_NT = 128 * NWG, WG_CLD = WG_BN + 8; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int GEMM_B = STAGES * (A_E + B_E) * 2; + constexpr int EPI_B = WG_BM * WG_CLD * 4; + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute(wg_conv_kernel, SH); + return; + } + native_invoke(wg_conv_kernel,dim3(dim3(DIVUP(Cout, WG_BM), DIVUP(st.N, WG_BN))),dim3(WG_NT),SH, + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg2_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, + bool prep = false) { + constexpr int BK = CIB * KT; + constexpr int WG_BM = 64 * MSPLIT; + constexpr int WG_NT = 128 * MSPLIT * KSPLIT * NSPLIT; + constexpr int WG_BNT = WG_BN * NSPLIT; + constexpr int NCRAW = WG_BNT / 4; + constexpr int EPN = + (WG_NT % NCRAW == 0) ? WG_BNT : (((WG_BNT + 127) / 128) * 128); + constexpr int WG_CLD = EPN + 8; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int GEMM_B = KSPLIT * (SA * A_E + SB * NSPLIT * B_E) * 2; + constexpr int EPI_B = KSPLIT * WG_BM * WG_CLD * 4; + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute( + wg2_conv_kernel, + SH); + return; + } + native_invoke_pf(wg2_conv_kernel,dim3(dim3(DIVUP(Cout, WG_BM), DIVUP(st.N, WG_BNT))),dim3(WG_NT),SH, + w, (size_t)KT * Cin * Cout * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg3_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, + bool prep = false) { + constexpr int WG_BM = 64 * NWG, WG_NT = 128 * NWG * NSG; + constexpr int WG_BNT = WG_BN * NSG, WG_CLD = WG_BNT + 8; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int GEMM_B = STAGES * (A_E + NSG * B_E) * 2; + constexpr int EPI_B = WG_BM * WG_CLD * 4; + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute(wg3_conv_kernel, SH); + return; + } + native_invoke(wg3_conv_kernel,dim3(dim3(DIVUP(Cout, WG_BM), DIVUP(st.N, WG_BNT))),dim3(WG_NT),SH, + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg4_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, int CoutP, + const Stack &st, int SCin, int do_snake, int pad_off, + int pad_len, bool prep = false) { + constexpr int NT = 128 * MSPLIT * NSPLIT; + constexpr int BM = 64 * MSPLIT, BNT = WG_BN * NSPLIT; + constexpr int NBR = BNT + KT - 1; + constexpr int NBRP = BC8 ? ((NBR + 7) & ~7) : ((NBR + 31) & ~31); + constexpr int NCRAW = BNT / 4; + constexpr int EPN = + (NT % NCRAW == 0) ? BNT : (((BNT + 127) / 128) * 128); + constexpr int B_E = CIB * NBRP; + constexpr int GEMM_B = STAGES * B_E * 2; + constexpr int EPI_B = BM * (EPN + 8) * (PREACT ? 2 : 4); + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute( + wg4_conv_kernel, + SH); + return; + } + pdl_launch_pf(wg4_conv_kernel, + dim3(DIVUP(Cout, BM), DIVUP(st.N, BNT)), dim3(NT), SH, s, + w, (size_t)KT * CoutP * Cin * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, + st.T, st.L, SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg5_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, + bool prep = false) { + constexpr int BM = 64 * MSPLIT; + constexpr int NT = 128 * MSPLIT * KSPLIT; + constexpr int NCRAW = WG_BN / 4; + constexpr int EPN = (NT % NCRAW == 0) ? WG_BN : (((WG_BN + 127) / 128) * 128); + constexpr int NBRP = BC8 ? (((WG_BN + 2) + 7) & ~7) + : (((WG_BN + 2) + 31) & ~31); + constexpr int B_E = CIB * NBRP; + constexpr int GEMM_B = KSPLIT * SB * B_E * 2; + constexpr int EPI_COLS = OC8 ? WG_BN : EPN; + constexpr int EPI_CLD_R = EPI_COLS + 8; + constexpr int EPI_CLD = (EPI_CLD_R % 16 == 0) ? EPI_CLD_R + 4 : EPI_CLD_R; + constexpr int EPI_B = KSPLIT * BM * EPI_CLD * (OC8 ? 2 : 4); + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + static_assert(SH <= 227 * 1024, "wg5 shared over budget"); + if (prep) { + prepare_attribute( + wg5_conv_kernel, + SH); + return; + } + if constexpr (PDL) { + cudaLaunchConfig_t cfg = {}; + cudaLaunchAttribute attr[1] = {}; + native_config_grid(cfg) = dim3(DIVUP(Cout, BM), DIVUP(st.N, WG_BN)); + native_config_threads(cfg) = dim3(NT); + cfg.dynamicSmemBytes = SH; + cfg.stream = s; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = 1; + cfg.attrs = attr; + cfg.numAttrs = 1; + native_cudaLaunchKernelEx_pf( + &cfg, + wg5_conv_kernel, + w, (size_t)3 * Cin * Cout * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); + } else { + native_invoke_pf(wg5_conv_kernel,dim3(dim3(DIVUP(Cout, BM), DIVUP(st.N, WG_BN))),dim3(NT),SH, + w, (size_t)3 * Cin * Cout * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, + st.L, SCin, st.SC, st.N, do_snake, pad_off, pad_len); + } +} +#ifndef D7W +#define D7W 96 +#endif +#define D7_M2_C8I(BN,LATE) wg4_launch<7, 16, 2, 1, D7W, 4, 1, false, 2, true, false, LATE>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_M2P_C8O(BN,LATE) wg4_launch<7, 16, 2, 1, D7W, 4, 1, true, 2, false, true, LATE>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_M2N2(BN) wg4_launch<7, 16, 2, 2, BN, 4, 1, false, 2>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_M1(BN) wg4_launch<7, 16, 1, 1, BN, 4, 2, false, 2>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_C8I(W) \ + wg4_launch<7, 16, 2, 1, W, 4, 1, false, 2, true, false, false>( \ + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, \ + SCin, 1, pad_off, pad_len, prep) +#define D7_C8O(W) \ + wg4_launch<7, 16, 2, 1, W, 4, 1, true, 2, false, true, false>( \ + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, \ + SCin, 1, pad_off, pad_len, prep) +template +__device__ __forceinline__ int d7_band_cost(int N, int mtiles, int resident) { + const int v = mtiles * DIVUP(N, BN); + const int sub = (DIVUP(v, resident) >= 2) ? 2 : 1; + return DIVUP(DIVUP(v, sub), resident) * sub * (128 + BN + 6); +} +__device__ __forceinline__ int d7_band(int N, int mtiles, int resident) { + int best = 64, bc = d7_band_cost<64>(N, mtiles, resident); + const int c96 = d7_band_cost<96>(N, mtiles, resident); + if (c96 < bc) { bc = c96; best = 96; } + if (d7_band_cost<112>(N, mtiles, resident) < bc) best = 112; + return best; +} +__device__ __forceinline__ void dense7(cudaStream_t s, int job, const __half *in, const __half *w, + const float *bias, const float *alpha, const __half *resid, + __half *out, const __half *pad_src, __half *ns, + const Stack &st, int SCin, int pad_off, int pad_len, + bool prep = false) { + const int Cin = native_RP_CIN(job), Cout = native_RP_COUT(job), CoutP = native_RP_COUTP(job); + const int b_ = d7_band(st.N, DIVUP(Cout, 128), native_resident()); + if ((job & 1) == 0) { + if (b_ == 64) D7_C8O(64); + else if (b_ == 96) D7_C8O(96); + else D7_C8O(112); + } else { + if (b_ == 64) D7_C8I(64); + else if (b_ == 96) D7_C8I(96); + else D7_C8I(112); + } +} +__device__ __forceinline__ int wg5_band_skip(int N, int mtiles, int resident) { + const int wmin = DIVUP(mtiles * DIVUP(N, 104), resident); + if (DIVUP(mtiles * DIVUP(N, 24), resident) == wmin) return 24; + if (DIVUP(mtiles * DIVUP(N, 40), resident) == wmin) return 40; + if (DIVUP(mtiles * DIVUP(N, 48), resident) == wmin) return 48; + if (DIVUP(mtiles * DIVUP(N, 64), resident) == wmin) return 64; + return 104; +} +__device__ __forceinline__ int wg5_band_main_long(int N, int mtiles, + int resident) { + const int wmin = DIVUP(mtiles * DIVUP(N, 104), resident); + if (DIVUP(mtiles * DIVUP(N, 64), resident) == wmin) return 64; + return 104; +} +__device__ __forceinline__ int wg5_band_main_short(int N, int mtiles, + int resident) { + const int wmin = DIVUP(mtiles * DIVUP(N, 96), resident); + if (DIVUP(mtiles * DIVUP(N, 56), resident) == wmin) return 56; + if (DIVUP(mtiles * DIVUP(N, 80), resident) == wmin) return 80; + return 96; +} +#define DENSE_LAUNCH(KT,BM,BN,NW,WM,WN) native_invoke(dense_conv_kernel,dim3(dim3(DIVUP(Cout, BM), DIVUP(st.N, BN))),dim3(NW * 32),0, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, SCin, st.SC, st.N, do_snake, pad_off, pad_len) +__device__ __forceinline__ void dense(cudaStream_t s, int KT, const __half *in, const __half *w, + const float *bias, const float *alpha, const __half *resid, + __half *out, const __half *pad_src, __half *ns, int Cin, + int Cout, const Stack &st, int SCin, int do_snake, + int pad_off, int pad_len, const __half *w3 = nullptr, + int layout_mode = 0, bool prep = false) { + if (KT == 3 && w3 != nullptr) { + if (Cin == 1536) { +#define WG5S(BN,O,R,PDL) wg5_launch<16, 1, 4, BN, 4, 3, 1, 1, true, O, R, 1536, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5SW(BN,O,R,PDL) wg5_launch<32, 1, 2, BN, 4, 3, 1, 1, true, O, R, 1536, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5S3(O,R) do { const int n_ = nref(st); if (n_ > 500 && n_ <= 640) { WG5SW(56, O, R, false); } else { const int b_ = wg5_band_skip(st.N, DIVUP(Cout, 64), native_resident()); if (b_ == 24) WG5S(24, O, R, true); else if (b_ == 40) WG5S(40, O, R, true); else if (b_ == 48) WG5S(48, O, R, true); else if (b_ == 64) WG5S(64, O, R, true); else WG5S(104, O, R, true); } } while (0) + if (layout_mode == 1) WG5S3(true, false); + else if (layout_mode == 2) WG5S3(true, true ); + else if (layout_mode == 3) WG5S3(false, true ); + else WG5S3(false, false); +#undef WG5S3 +#undef WG5SW +#undef WG5S + } else { +#define WG5I(BN,BC,MS,PDL) wg5_launch<16, MS, 2, BN, 4, 3, 1, 1, BC, true, false, 768, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5I4(BN,BC,PDL) wg5_launch<16, 1, 4, BN, 4, 3, 1, 1, BC, true, false, 768, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5I2W(BN,BC,PDL) wg5_launch<32, 1, 2, BN, 4, 3, 1, 1, BC, true, false, 768, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5I3(BC) do { const int n_ = nref(st); if (n_ > 640) { const int b_ = wg5_band_main_long(st.N, DIVUP(Cout, 128), native_resident()); if (b_ == 64) WG5I(64, BC, 2, true); else WG5I(104, BC, 2, true); } else if (n_ > 500) { WG5I2W(64, BC, false); } else { const int b_ = wg5_band_main_short(st.N, DIVUP(Cout, 64), native_resident()); if (b_ == 56) WG5I4(56, BC, true); else if (b_ == 80) WG5I4(80, BC, true); else WG5I4(96, BC, true); } } while (0) + if (layout_mode == 1) WG5I3(true); + else WG5I3(false); +#undef WG5I3 +#undef WG5I2W +#undef WG5I4 +#undef WG5I + } + } else if (KT == 3) { + const int n_ = nref(st); + if (Cin == 1536 && Cout == 768) { + if (n_ <= 704) + wg2_launch<3, 16, 1, 4, 1, 64, 2, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else + wg2_launch<3, 16, 1, 4, 1, 128, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + } else if (Cin == 768 && Cout == 1536) { + if (n_ > 1000) + wg2_launch<3, 16, 1, 2, 1, 128, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ > 640 && n_ <= 768) + wg2_launch<3, 16, 1, 2, 1, 64, 2, 2, 2>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ <= 320) + wg2_launch<3, 16, 1, 4, 1, 64, 2, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ > 384 && n_ <= 512) + wg2_launch<3, 16, 1, 4, 1, 112, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ > 512 && n_ <= 640) + wg2_launch<3, 16, 1, 4, 1, 120, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else + wg2_launch<3, 16, 1, 4, 1, 128, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + } else { + wg2_launch<3, 16, 1, 1, 1, 64, 3, 2, 2>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, + SCin, do_snake, pad_off, pad_len, prep); + } + } else if (Cout >= 128 && + ((Cin >= 384 && st.T == 288) || + (Cin < 384 && st.T == 2016))) { + wg3_launch<7, 112, 2, 2, 1, 128, 2>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, + SCin, do_snake, pad_off, pad_len); + } else { + wg3_launch<7, 112, 1, 4, 1, 128, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, + SCin, do_snake, pad_off, pad_len); + } +} +__device__ void codec_program(const CodecArgs &a, const Plan &p, const RepackSet &js, int repack_grid) { + cudaStream_t s = nullptr; + const int B = a.B; + const void *const *W = a.inp; + const __half *const *ST = reinterpret_cast(a.st); + __half *const *NS = reinterpret_cast<__half *const *>(a.ns); + { + native_invoke(pack_latent_kernel,dim3(DIVUP(40*B*p.s0.L,256)),dim3(256),0, + H16(W[0]), ST[0], p.LB, 40, p.s0.T, p.s0.L, p.s0.SC, 2, B); + } + dense(s, 3, p.LB, H16(W[1]), F32(W[2]), nullptr, nullptr, p.X0, ST[1], NS[0], + 40, 768, p.s0, p.s0.SC, 0, 0, 2); + const __half *k3iw = H16(W[3]); + const float *k3ib = F32(W[4]); + const float *k3ia = F32(W[5]); + const __half *k3sw = H16(W[6]); + const float *k3sb = F32(W[7]); + const float *k3sa = F32(W[8]); + const size_t IW = (size_t)1536 * 768 * 3, SW = (size_t)768 * 1536 * 3; + dense(s, 3, p.X0, k3iw, k3ib, k3ia, nullptr, p.M0, ST[2], NS[1], 768, 1536, + p.s0, p.s0.SC, 1, 0, 2, p.W3i ? p.W3i - IW : nullptr, 0); + dense(s, 3, p.M0, k3sw, k3sb, k3sa, p.X0, p.Y0, ST[3], NS[2], 1536, 768, + p.s0, p.s0.SC, 1, 1, 1, p.W3s ? p.W3s - SW : nullptr, 0); + TCONV_BPACK_DISPATCH8(1, 2, 8, p.s0.T, 768, B, + p.Y0, H16(W[9]), F32(W[10]), F32(W[11]), p.A, ST[4], NS[3], 768, 768, + p.s0.T, p.s0.L, p.s0.SC, 1, p.s1.L, p.s1.SC, 2, 0, 2, B); + if (p.W3i != nullptr) + native_invoke(repack_activation_only,dim3(min(1152,96*DIVUP(p.s1.SC,128))),dim3(128),0, + p.A, p.Ac, p.s1.SC, p.s1.SC, DIVUP(p.s1.SC,128)); + __half *cur = p.A, *nxt = p.Bb; + for (int i = 0; i < 6; ++i) { + const int bi = i + 1; + const __half *ain = (i == 0 && p.Ac) ? p.Ac : cur; + dense(s, 3, ain, k3iw + bi * IW, k3ib + bi * 1536, k3ia + bi * 768, + nullptr, p.M1, ST[5 + 2 * i], NS[4 + 2 * i], 768, 1536, p.s1, + p.s1.SC, 1, 0, 2, p.W3i ? p.W3i + (bi - 1) * IW : nullptr, + (i > 0 || p.Ac) ? 1 : 0); + const __half *ps = (i < 5) ? ST[6 + 2 * i] : ST[16]; + const int po = (i < 5) ? 0 : 1, pl = 1 + (i < 5); + dense(s, 3, p.M1, k3sw + bi * SW, k3sb + bi * 768, k3sa + bi * 384, cur, + nxt, ps, NS[5 + 2 * i], 1536, 768, p.s1, p.s1.SC, 1, po, pl, + p.W3s ? p.W3s + (bi - 1) * SW : nullptr, + i == 0 ? 1 : (i < 5 ? 2 : 3)); + __half *t = cur; cur = nxt; nxt = t; + } + TCONV_BPACK_DISPATCH1(2, 9, 16, p.s1.T, 384, B, + cur, H16(W[12]), F32(W[13]), F32(W[14]), p.P, ST[17], NS[16], 768, 384, + p.s1.T, p.s1.L, p.s1.SC, 1, p.s2.L, p.s2.SC, 6, 0, 6, B); + dense7(s, 0, p.P, p.W7[0], F32(W[22]), F32(W[23]), nullptr, p.M2, ST[18], + NS[17], p.s2, p.s2.SC, 0, 6); + dense7(s, 1, p.M2, p.W7[1], F32(W[25]), F32(W[26]), p.P, p.Q, ST[19], + NS[18], p.s2, p.s2.SC, 5, 1); + TCONV_BPACK_DISPATCH(3, 7, 16, p.s2.T, 128, B, + p.Q, H16(W[15]), F32(W[16]), F32(W[17]), p.R, ST[20], NS[19], 384, 128, + p.s2.T, p.s2.L, p.s2.SC, 5, p.s3.L, p.s3.SC, 6, 0, 6, B); + dense7(s, 2, p.R, p.W7[2], F32(W[28]), F32(W[29]), nullptr, p.M3, ST[21], + NS[20], p.s3, p.s3.SC, 0, 6); + dense7(s, 3, p.M3, p.W7[3], F32(W[31]), F32(W[32]), p.R, p.S, ST[22], + NS[21], p.s3, p.s3.SC, 5, 1); + if (p.s3.T >= 1500) + pdl_launch(transpose_conv_kernel<4, 7, 8, true>, dim3(1, 32, B), + dim3(tc_blk(p.s3.T, 8)), + (size_t)tc_blk(p.s3.T, 8) * (8 * 7) * 2 + 64, s, + p.S, H16(W[18]), F32(W[19]), F32(W[20]), p.U, ST[23], + NS[22], 128, 32, p.s3.T, p.s3.L, p.s3.SC, 5, p.s4.L, + p.s4.SC, 6, 0, 6, B); + else + pdl_launch(transpose_conv_kernel<4, 7, 4, true>, dim3(1, 32, B), + dim3(tc_blk(p.s3.T, 4)), + (size_t)tc_blk(p.s3.T, 4) * (4 * 7) * 2 + 64, s, + p.S, H16(W[18]), F32(W[19]), F32(W[20]), p.U, ST[23], + NS[22], 128, 32, p.s3.T, p.s3.L, p.s3.SC, 5, p.s4.L, + p.s4.SC, 6, 0, 6, B); + { + const int ntile = DIVUP(p.s4.N, TL_BN); + const int grid = ntile < 132 ? ntile : 132; + pdl_launch_pf(tail_fused_kernel, dim3(grid), dim3(TL_NT), TL_SHARED, s, + p.U, (size_t)32 * p.s4.SC * sizeof(__half), + p.U, p.W7[4], p.W7[5], H16(W[39]), F32(W[34]), F32(W[35]), + F32(W[37]), F32(W[38]), F32(W[40]), ST[24], ST[25], + static_cast<__half *>(a.waveform), NS[23], NS[24], NS[25], B, + p.s4.T, p.s4.L, p.s4.SC, p.s4.N, ntile); + } +} diff --git a/deploy/nemotron_native/kernels/codec/base/native_stage1.cu b/deploy/nemotron_native/kernels/codec/base/native_stage1.cu new file mode 100644 index 0000000..56f759c --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/base/native_stage1.cu @@ -0,0 +1,309 @@ +#include "native_stage1.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#ifndef NATIVE_CODEC_PREFETCH +#define NATIVE_CODEC_PREFETCH 1 +#endif +#ifndef NATIVE_CODEC_PACK_SUBCTAS +#define NATIVE_CODEC_PACK_SUBCTAS 1 +#endif +namespace cg=cooperative_groups; +constexpr int PHYSICAL_THREADS=512,MAIN_ARENA=227840,AUX_BYTES=4096,META_OFFSET=231936,TOTAL_SHARED=232192; +constexpr int SLICE_BYTES=MAIN_ARENA+AUX_BYTES; +constexpr int MAX_SUBCTA=8; +struct NativeMeta { + dim3 grid,threads; + uint3 block[MAX_SUBCTA]; + NativeStage1Status *status;const int32_t *lengths;int padded_frames; + int submask,tshift,tmask,slice; + int resident; +}; +static_assert(sizeof(NativeMeta)<=TOTAL_SHARED-META_OFFSET); +__device__ __forceinline__ char *native_base(){extern __shared__ __align__(16) char arena[];return arena;} +__device__ __forceinline__ NativeMeta *native_meta(){return reinterpret_cast(native_base()+META_OFFSET);} +__device__ __forceinline__ int native_subid(){const NativeMeta *m=native_meta();return (int)((threadIdx.x>>m->tshift)&(unsigned)m->submask);} +__device__ __forceinline__ char *native_arena(){const NativeMeta *m=native_meta();return native_base()+native_subid()*m->slice;} +__device__ __forceinline__ char *native_aux(){const NativeMeta *m=native_meta();return native_base()+native_subid()*m->slice+(m->slice-AUX_BYTES);} +__device__ __forceinline__ uint3 native_tid(){uint3 r;r.x=threadIdx.x&(unsigned)native_meta()->tmask;r.y=0u;r.z=0u;return r;} +__device__ __forceinline__ dim3 native_block(){const uint3 b=native_meta()->block[native_subid()];return dim3(b.x,b.y,b.z);} +__device__ __forceinline__ dim3 native_grid(){return native_meta()->grid;} +__device__ __forceinline__ dim3 native_threads(){return native_meta()->threads;} +__device__ __forceinline__ int native_resident(){return native_meta()->resident;} +__device__ __forceinline__ void native_local_sync(){ + const NativeMeta *m=native_meta();unsigned n=m->threads.x,id=15u-(unsigned)native_subid(); + asm volatile("bar.sync %0, %1;"::"r"(id),"r"(n):"memory"); +} +__device__ __forceinline__ dim3 &native_config_grid(cudaLaunchConfig_t &c){return c.gridDim;} +__device__ __forceinline__ dim3 &native_config_threads(cudaLaunchConfig_t &c){return c.blockDim;} +__device__ __forceinline__ int native_true_tokens(int row,int padded_tokens){ + const NativeMeta *m=native_meta(); + return (padded_tokens/m->padded_frames)*m->lengths[row]; +} +__device__ __forceinline__ void native_warm(const char *p,size_t n,int k,int nk,int me){ + if(p==nullptr||n<16||nk<=0)return; + const size_t a=(size_t)p&15u;p-=a;n=(n+a)&~(size_t)15; + const uint4 *q=reinterpret_cast(p); + const size_t cnt=n>>4,step=(size_t)nk*PHYSICAL_THREADS; + uint32_t acc=0u; + for(size_t i=(size_t)k*PHYSICAL_THREADS+(size_t)me;i __device__ __forceinline__ void native_invoke_pf(F fn,dim3 g,dim3 t,size_t shared,const void *pf,size_t pfb,A&&... args){ + const int me=threadIdx.x,owner=blockIdx.x,owners=gridDim.x; + const cg::grid_group all=cg::this_grid();NativeMeta *meta=native_meta(); + if(t.x>512||t.y!=1||t.z!=1||t.x%32||shared>MAIN_ARENA){ + if(me==0&&owner==0)meta->status->code=10; + all.sync();return; + } + int sub=1,tshift=9; + const unsigned tx=t.x; + const int vtotal=(int)(g.x*g.y*g.z); + if(NATIVE_CODEC_PACK_SUBCTAS && (tx&(tx-1))==0u){ + int s=0;while((1u<MAX_SUBCTA)cand=MAX_SUBCTA; + const int want=(vtotal+owners-1)/owners; + while(cand>1&&cand>want)cand>>=1; + while(cand>1&&shared>(size_t)(SLICE_BYTES/cand-AUX_BYTES))cand>>=1; + sub=cand; + } + const int slice=(SLICE_BYTES/sub)&~127; + if(me==0){meta->grid=g;meta->threads=t;meta->submask=sub-1;meta->tshift=tshift; + meta->tmask=(int)((1u<slice=slice;} + const int total=vtotal,grp=me>>tshift; + const bool holder=me=busy&&pf!=nullptr)native_warm((const char*)pf,pfb,owner-busy,owners-busy,me); + for(int base=owner*sub;baseblock[grp]=make_uint3(i%g.x,(i/g.x)%g.y,i/(g.x*g.y)); + __syncthreads(); + if(holder&&i(args)...); + asm volatile("cp.async.wait_all;":::"memory"); + if(t.x>=128)asm volatile("wgmma.wait_group.sync.aligned 0;":::"memory"); + } + __syncthreads(); + } + all.sync(); +} +template __device__ __forceinline__ void native_invoke(F fn,dim3 g,dim3 t,size_t shared,A&&... args){ + native_invoke_pf(fn,g,t,shared,nullptr,(size_t)0,static_cast(args)...); +} +template __device__ __forceinline__ void native_cudaLaunchKernelEx(cudaLaunchConfig_t *c,F fn,A... args){ + native_invoke(fn,c->gridDim,c->blockDim,c->dynamicSmemBytes,args...); +} +template __device__ __forceinline__ void native_cudaLaunchKernelEx_pf(cudaLaunchConfig_t *c,F fn,const void *pf,size_t pfb,A... args){ + native_invoke_pf(fn,c->gridDim,c->blockDim,c->dynamicSmemBytes,pf,pfb,args...); +} +#define blockIdx native_block() +#define gridDim native_grid() +#define blockDim native_threads() +#define threadIdx native_tid() +#define __syncthreads() native_local_sync() +#include "legacy_codec.cuh" +#undef __syncthreads +#undef threadIdx +#undef blockDim +#undef gridDim +#undef blockIdx +namespace { +struct FastArgs { + CodecArgs codec;Plan plan;RepackSet repack;int repack_grid;int input_frames;int blocks; + const int64_t *codes;const int32_t *lengths,*slots; + const __half *state[26];__half *next[26];float *pcm;int32_t *samples; + NativeStage1Status *status;int nslots; + int channels[26],histories[26]; +}; +static_assert(sizeof(FastArgs)<=4096,"precomputed kernel argument must fit the conservative parameter budget"); +struct Handle {FastArgs args{};Plan plans[17]{};RepackSet repacks[17]{};int repack_grids[17]{};void *scratch=nullptr;size_t bytes=0;void *immutable=nullptr;size_t immutable_bytes=0;int setup_launches=0;std::vector packed;int max_frames=0,blocks=0,device=-1;}; +thread_local std::string error; +bool check(cudaError_t e,const char *where){if(e==cudaSuccess)return true;error=std::string(where)+": "+cudaGetErrorString(e);return false;} +__device__ void reject(NativeStage1Status *s,int code,int row,int frame,int detail){ + if(atomicCAS(&s->code,0,code)==0){s->row=row;s->frame=frame;s->detail=detail;} +} +__global__ __launch_bounds__(512,1) void stage1_fast_persistent(const __grid_constant__ FastArgs a){ + const auto all=cg::this_grid();size_t me=(size_t)blockIdx.x*512+threadIdx.x,step=(size_t)gridDim.x*512; + if(threadIdx.x==0){native_meta()->status=a.status;native_meta()->lengths=a.lengths;native_meta()->padded_frames=a.codec.F;native_meta()->resident=a.blocks;} + if(me==0)*a.status=NativeStage1Status{0,-1,-1,0}; + all.sync();const int B=a.codec.B,F=a.codec.F,inputF=a.input_frames; + for(int row=me;rowinputF)reject(a.status,8,row,-1,n); + if(slot<0||slot>=a.nslots)reject(a.status,2,row,-1,slot); + for(int j=0;j=0&&n<=inputF)?n*1764:0; + } + for(size_t i=me;i<(size_t)B*inputF;i+=step){ + int row=i/inputF,t=i%inputF;bool any=false; + for(int c=0;c<16;++c){int64_t v=a.codes[i*16+c];if(v<0||v>=1024)reject(a.status,4,row,t,c);any|=v!=0;} + if(t=a.lengths[row]&&any)reject(a.status,6,row,t,0); + } + for(int p=0;p<26;++p)for(size_t i=me;i<(size_t)a.nslots*a.channels[p]*a.histories[p];i+=step)a.next[p][i]=a.state[p][i]; + for(size_t i=me;i<(size_t)B*inputF*1764;i+=step)a.pcm[i]=0.f; + all.sync();if(a.status->code)return; + for(int p=0;p<26;++p){ + const size_t stride=(size_t)a.channels[p]*a.histories[p]; + __half *dest=(__half*)a.codec.st[p]; + for(size_t i=me;i<(size_t)B*stride;i+=step)dest[i]=a.state[p][(size_t)a.slots[i/stride]*stride+i%stride]; + } + __half *latent=(__half*)a.codec.inp[0];int T=2*F; + for(size_t i=me;i<(size_t)B*40*T;i+=step){ + int t=i%T,c=(i/T)%40,row=i/((size_t)40*T);float v=0.f; + if(t/2>(2*(c%5)))&3)-2)*.5f;} + latent[i]=__float2half_rn(v); + } + all.sync(); + codec_program(a.codec,a.plan,a.repack,a.repack_grid); + if(a.status->code)return; + for(int p=0;p<26;++p){ + const size_t stride=(size_t)a.channels[p]*a.histories[p];const __half *src=(const __half*)a.codec.ns[p]; + for(size_t i=me;i<(size_t)B*stride;i+=step){int row=i/stride;if(a.lengths[row])a.next[p][(size_t)a.slots[row]*stride+i%stride]=src[i];} + } + const __half *wave=(const __half*)a.codec.waveform; + for(size_t i=me;i<(size_t)B*inputF*1764;i+=step){ + const int row=i/(inputF*1764),t=i%(inputF*1764); + if(tstatus=a.status;native_meta()->resident=a.blocks;} + if(blockIdx.x==0&&threadIdx.x==0)*a.status=NativeStage1Status{0,-1,-1,0}; + all.sync(); + native_invoke(repack_rs_kernel<7>,dim3(a.repack_grid),dim3(128),MAIN_ARENA, + a.repack,H16(a.codec.inp[0]),H16(a.codec.st[0]),a.plan.LB, + 40,a.plan.s0.T,a.plan.s0.L,a.plan.s0.SC,2,a.codec.B); + native_invoke(repack3_rs_kernel,dim3(1152,12),dim3(128),MAIN_ARENA, + H16(a.codec.inp[3]),H16(a.codec.inp[6]),a.plan.W3i,a.plan.W3s, + a.plan.A,a.plan.Ac,a.plan.s1.SC,a.plan.s1.SC,DIVUP(a.plan.s1.SC,128)); +} +bool stack(Handle *h,int slot,const NativeStage1Weights *w,int first,int field,size_t elements){ + const int ids[7]={first,first+3,first+5,first+7,first+9,first+11,first+13}; + const size_t stride=elements*(field==0?sizeof(__half):sizeof(float));void *ptr=nullptr; + if(!check(cudaMalloc(&ptr,7*stride),"packed-weight allocation"))return false; + h->packed.push_back(ptr);h->args.codec.inp[slot]=ptr; + for(int i=0;i<7;++i){const void *src=field==0?w[ids[i]].weight:field==1?(const void*)w[ids[i]].bias:(const void*)w[ids[i]].alpha; + if(!check(cudaMemcpy((char*)ptr+i*stride,src,stride,cudaMemcpyDeviceToDevice),"immutable weight stack"))return false;} + return true; +} +void release(Handle *h){if(h->immutable)cudaFree(h->immutable);if(h->scratch)cudaFree(h->scratch);for(void *p:h->packed)cudaFree(p);delete h;} +} +extern "C" void *native_stage1_create(int b,int maxf,int slots,int blocks,const NativeStage1Weights *w){ + error.clear();if((b!=1&&b!=16&&b!=32&&b!=64)||maxf<1||maxf>16||slots4096||blocks<0||!w){error="invalid create arguments";return nullptr;} + for(int p=0;p<26;++p)if(!w[p].weight||!w[p].bias||(p!=0&&p!=25&&!w[p].alpha)){error="missing weight pointer";return nullptr;} + Handle *h=new(std::nothrow)Handle;if(!h){error="host allocation";return nullptr;} + h->max_frames=maxf;h->args.nslots=slots;h->args.codec.B=b; + const int cs[26]={40,768,1536,768,768,1536,768,1536,768,1536,768,1536,768,1536,768,1536,768,384,768,384,128,256,128,32,64,32}; + const int hs[26]={2,2,2,1,2,2,2,2,2,2,2,2,2,2,2,2,1,6,6,1,6,6,1,6,6,6}; + cudaDeviceProp prop{};int occupancy=0,coop=0; + if(!check(cudaGetDevice(&h->device),"device")||!check(cudaGetDeviceProperties(&prop,h->device),"properties")|| + !check(cudaDeviceGetAttribute(&coop,cudaDevAttrCooperativeLaunch,h->device),"cooperative support")|| + !check(cudaFuncSetAttribute(stage1_fast_persistent,cudaFuncAttributeMaxDynamicSharedMemorySize,TOTAL_SHARED),"shared-memory optin")|| + !check(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&occupancy,stage1_fast_persistent,512,TOTAL_SHARED),"occupancy")) {release(h);return nullptr;} + if(prop.major!=9||prop.minor!=0||!coop||occupancy<1){error="requires resident SM90 cooperative kernel";release(h);return nullptr;} + h->blocks=blocks?blocks:occupancy*prop.multiProcessorCount; + if(h->blocks>occupancy*prop.multiProcessorCount){error="grid exceeds cooperative occupancy";release(h);return nullptr;} + h->args.blocks=h->blocks; + const int capacityF=maxf<=3?3:(maxf<=5?5:std::max(8,maxf)); + size_t decoder=codec_scratch_bytes(b,capacityF),halfs=(size_t)b*40*2*capacityF+(size_t)b*1764*capacityF; + for(int p=0;p<26;++p){h->args.channels[p]=cs[p];h->args.histories[p]=hs[p];halfs+=2*(size_t)b*cs[p]*hs[p];} + h->bytes=decoder+halfs*sizeof(__half); + if(!check(cudaMalloc(&h->scratch,h->bytes),"scratch allocation")){release(h);return nullptr;} + if(!check(cudaMemset(h->scratch,0,h->bytes),"initial scratch padding")){release(h);return nullptr;} + h->args.codec.scratch=h->scratch;__half *q=(__half*)((char*)h->scratch+decoder); + h->args.codec.inp[0]=q;q+=(size_t)b*40*2*capacityF;h->args.codec.waveform=q;q+=(size_t)b*1764*capacityF; + for(int p=0;p<26;++p){h->args.codec.st[p]=q;q+=(size_t)b*cs[p]*hs[p];h->args.codec.ns[p]=q;q+=(size_t)b*cs[p]*hs[p];} + const void **in=h->args.codec.inp; + in[1]=w[0].weight;in[2]=w[0].bias; + if(!stack(h,3,w,1,0,(size_t)1536*768*3)||!stack(h,4,w,1,1,1536)||!stack(h,5,w,1,2,768)|| + !stack(h,6,w,2,0,(size_t)768*1536*3)||!stack(h,7,w,2,1,768)||!stack(h,8,w,2,2,384)){release(h);return nullptr;} + const int upops[4]={3,16,19,22};for(int i=0;i<4;++i){in[9+3*i]=w[upops[i]].weight;in[10+3*i]=w[upops[i]].bias;in[11+3*i]=w[upops[i]].alpha;} + for(int i=0;i<3;++i){int p=17+3*i;in[21+6*i]=w[p].weight;in[22+6*i]=w[p].bias;in[23+6*i]=w[p].alpha; + in[24+6*i]=w[p+1].weight;in[25+6*i]=w[p+1].bias;in[26+6*i]=w[p+1].alpha;} + in[39]=w[25].weight;in[40]=w[25].bias; + __half *immutable_w7[6]{};__half *immutable_w3i=nullptr,*immutable_w3s=nullptr; + if(capacityF>=3){ + size_t wh=0;for(int j=0;j<6;++j)wh+=(size_t)native_RP_COUTP(j)*native_RP_CIN(j)*7; + const size_t IW=(size_t)3*1536*768;wh+=14*IW; + h->immutable_bytes=wh*sizeof(__half); + if(!check(cudaMalloc(&h->immutable,h->immutable_bytes+sizeof(NativeStage1Status)),"immutable weight allocation")){release(h);return nullptr;} + if(!check(cudaMemset(h->immutable,0,h->immutable_bytes+sizeof(NativeStage1Status)),"immutable padding initialization")){release(h);return nullptr;} + __half *q=(__half*)h->immutable; + for(int j=0;j<6;++j){immutable_w7[j]=q;q+=(size_t)native_RP_COUTP(j)*native_RP_CIN(j)*7;} + immutable_w3i=q+IW;q+=7*IW;immutable_w3s=q+IW;q+=7*IW; + if((char*)q!=(char*)h->immutable+h->immutable_bytes){error="immutable layout size mismatch";release(h);return nullptr;} + } + for(int f:{3,5,8}) if(f<=capacityF){ + Plan p=build_plan(b,f,static_cast<__half*>(h->scratch)); + for(int j=0;j<6;++j)p.W7[j]=immutable_w7[j];p.W3i=immutable_w3i;p.W3s=immutable_w3s; + const void *const *W=h->args.codec.inp; + RepackSet js{}; + const void *SRC[RP_JOBS] = {W[21], W[24], W[27], W[30], W[33], W[36]}; + for (int j = 0; j < RP_JOBS; ++j) { + js.src[j] = H16(SRC[j]); + js.dst[j] = p.W7[j]; + js.Cout[j] = native_RP_COUT(j); + js.CoutP[j] = native_RP_COUTP(j); + js.Cin[j] = native_RP_CIN(j); + } + int acc = 0; + for (int j = 0; j < RP_JOBS; ++j) { + js.base[j] = acc; + js.nch[j] = native_RP_CIN(j) / 16; + js.rs[j] = (j < 4) ? 1 : 0; + acc += (native_RP_COUTP(j) / 16) * DIVUP(native_RP_CIN(j), 64); + } + js.k3base = acc; + js.k3src[0] = H16(W[3]); + js.k3src[1] = H16(W[6]); + js.k3dst[0] = p.W3i ? p.W3i - (size_t)3 * 1536 * 768 : nullptr; + js.k3dst[1] = p.W3s ? p.W3s - (size_t)3 * 768 * 1536 : nullptr; + const int k3n = p.W3i ? 2304 : 0; + h->plans[f]=p;h->repacks[f]=js;h->repack_grids[f]=acc+k3n; + } + if(capacityF>=3){ + FastArgs a=h->args;a.codec.F=3;a.plan=h->plans[3];a.repack=h->repacks[3];a.repack_grid=h->repack_grids[3]; + a.status=(NativeStage1Status*)((char*)h->immutable+h->immutable_bytes); + int setup_occupancy=0; + if(!check(cudaFuncSetAttribute(stage1_prepare_immutable,cudaFuncAttributeMaxDynamicSharedMemorySize,TOTAL_SHARED),"setup shared-memory optin")|| + !check(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&setup_occupancy,stage1_prepare_immutable,512,TOTAL_SHARED),"setup occupancy")|| + setup_occupancy<1||h->blocks>setup_occupancy*prop.multiProcessorCount){ + if(error.empty())error="setup grid exceeds cooperative occupancy";release(h);return nullptr;} + void *params[]={&a}; + if(!check(cudaLaunchCooperativeKernel((void*)stage1_prepare_immutable,dim3(h->blocks),dim3(512),params,TOTAL_SHARED,nullptr),"immutable setup launch")|| + !check(cudaStreamSynchronize(nullptr),"immutable setup completion")){release(h);return nullptr;} + NativeStage1Status status{}; + if(!check(cudaMemcpy(&status,a.status,sizeof(status),cudaMemcpyDeviceToHost),"immutable setup status")){release(h);return nullptr;} + if(status.code){error="immutable setup device status "+std::to_string(status.code);release(h);return nullptr;} + h->setup_launches=1; + } + return h; +} +extern "C" int native_stage1_run(void *handle,const int64_t *codes,const int32_t *lengths,const int32_t *slots, + const void *const *states,void *const *next,float *pcm,int32_t *samples,NativeStage1Status *status,int f,void *stream){ + error.clear();Handle *h=(Handle*)handle; + if(!h||!codes||!lengths||!slots||!states||!next||!pcm||!samples||!status){error="null run argument";return 1;} + if(f<1||f>8||f>h->max_frames){error="true-state variant supports public F1..8";return 1;} + const int paddedF=f<=3?3:(f<=5?5:8); + int device=-1;if(!check(cudaGetDevice(&device),"device")||device!=h->device){error="device mismatch";return 1;} + FastArgs a=h->args;a.input_frames=f;a.plan=h->plans[paddedF];a.repack=h->repacks[paddedF];a.repack_grid=h->repack_grids[paddedF];a.codec.F=paddedF;a.codes=codes;a.lengths=lengths;a.slots=slots;a.pcm=pcm;a.samples=samples;a.status=status; + for(int p=0;p<26;++p){if(!states[p]||!next[p]||states[p]==next[p]){error="invalid/aliased state";return 1;}a.state[p]=(const __half*)states[p];a.next[p]=(__half*)next[p];} + void *params[]={&a};if(!check(cudaLaunchCooperativeKernel((void*)stage1_fast_persistent,dim3(h->blocks),dim3(512),params,TOTAL_SHARED,(cudaStream_t)stream),"fast Stage1 launch"))return 1; + return check(cudaPeekAtLastError(),"immediate launch status")?0:1; +} +extern "C" int native_stage1_destroy(void *p){error.clear();if(!p)return 0;release((Handle*)p);return check(cudaPeekAtLastError(),"destroy status")?0:1;} +extern "C" const char *native_stage1_last_error(){return error.c_str();} +extern "C" size_t native_stage1_scratch_bytes(void *p){return p?((Handle*)p)->bytes:0;} +extern "C" int native_stage1_resident_blocks(void *p){return p?((Handle*)p)->blocks:0;} +extern "C" int native_stage1_setup_kernel_launches(void *p){return p?((Handle*)p)->setup_launches:0;} +extern "C" size_t native_stage1_immutable_weight_bytes(void *p){return p?((Handle*)p)->immutable_bytes:0;} diff --git a/deploy/nemotron_native/kernels/codec/base/native_stage1.h b/deploy/nemotron_native/kernels/codec/base/native_stage1.h new file mode 100644 index 0000000..a95b5f0 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/base/native_stage1.h @@ -0,0 +1,26 @@ +#pragma once +#include +#include +struct NativeStage1Weights { const void *weight; const float *bias; const float *alpha; }; +struct NativeStage1Status { int32_t code, row, frame, detail; }; +#ifdef __cplusplus +extern "C" { +#endif +void *native_stage1_create(int batch, int max_frames, int state_slots, + int resident_blocks, const NativeStage1Weights *weights26); +int native_stage1_run(void *handle, const int64_t *codes, const int32_t *lengths, + const int32_t *slots, const void *const *states26, + void *const *next_states26, float *pcm, int32_t *samples, + NativeStage1Status *status, int frames, void *stream); +int native_stage1_run_profile(void *handle, const int64_t *codes, const int32_t *lengths, + const int32_t *slots, const void *const *states26, + void *const *next_states26, float *pcm, int32_t *samples, + NativeStage1Status *status, int frames, void *stream, + uint64_t *phase_timestamps30); +int native_stage1_destroy(void *handle); +const char *native_stage1_last_error(); +size_t native_stage1_scratch_bytes(void *handle); +int native_stage1_resident_blocks(void *handle); +#ifdef __cplusplus +} +#endif diff --git a/deploy/nemotron_native/kernels/codec/base/repack3.inc b/deploy/nemotron_native/kernels/codec/base/repack3.inc new file mode 100644 index 0000000..2f61c93 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/base/repack3.inc @@ -0,0 +1,66 @@ +__device__ __noinline__ void repack3_kernel( + const __half *__restrict__ win, const __half *__restrict__ wsk, + __half *__restrict__ din, __half *__restrict__ dsk) { + int4 (*sm)[49]=reinterpret_cast(native_arena()); + const int job = blockIdx.y; + const int Cout = job ? 768 : 1536; + const int Cin = job ? 1536 : 768; + const __half *src = job ? wsk : win; + __half *dst = job ? dsk : din; + const int GPB = Cin >> 3; + const int Kg = Cin * 3; + const int GT = GPB >> 4; + const int RT = Cout >> 5; + const int t = threadIdx.x; + for (int tile = blockIdx.x; tile < 6 * RT * GT; tile += gridDim.x) { + const int gblk = tile % GT; + const int rest = tile / GT; + const int corow = rest % RT; + const int b = rest / RT; + const int co0 = corow * 32; + const int g0 = gblk * 16; + const __half *sbase = src + (size_t)(b + 1) * Cout * Kg + + (size_t)co0 * Kg + g0 * 24; +#pragma unroll + for (int j = 0; j < 6; ++j) { + const int idx = t + j * 256; + const int r = idx / 48, c = idx - (idx / 48) * 48; + sm[r][c] = *reinterpret_cast(sbase + (size_t)r * Kg + + c * 8); + } + __syncthreads(); + __half *dbase = dst + (size_t)b * (3 * (size_t)Cout * Cin); +#pragma unroll + for (int p = 0; p < 2; ++p) { + const int u = t + p * 256; + const int col = (u & 7) + 8 * (u >> 7); + const int g = g0 + ((u >> 3) & 15); + const int co = co0 + col; + const int c = g >> 1, v = g & 1; + const int4 a = sm[col][((u >> 3) & 15) * 3]; + const int4 e = sm[col][((u >> 3) & 15) * 3 + 1]; + const int4 f = sm[col][((u >> 3) & 15) * 3 + 2]; + uint32_t h[12]; + h[0] = a.x; h[1] = a.y; h[2] = a.z; h[3] = a.w; + h[4] = e.x; h[5] = e.y; h[6] = e.z; h[7] = e.w; + h[8] = f.x; h[9] = f.y; h[10] = f.z; h[11] = f.w; + __half *d = dbase + (size_t)c * (48 * Cout) + + (size_t)(co >> 3) * 128 + v * 64 + (co & 7) * 8; +#pragma unroll + for (int tp = 0; tp < 3; ++tp) { + uint32_t o[4]; +#pragma unroll + for (int jj = 0; jj < 4; ++jj) { + const uint32_t g0w = h[3 * jj], g1w = h[3 * jj + 1], + g2w = h[3 * jj + 2]; + o[jj] = (tp == 0) ? __byte_perm(g0w, g1w, 0x7610) + : (tp == 1) ? __byte_perm(g0w, g2w, 0x5432) + : __byte_perm(g1w, g2w, 0x7610); + } + *reinterpret_cast(d + (size_t)tp * (16 * Cout)) = + make_int4(o[0], o[1], o[2], o[3]); + } + } + __syncthreads(); + } +} diff --git a/deploy/nemotron_native/kernels/codec/base/tail_kernel.inc b/deploy/nemotron_native/kernels/codec/base/tail_kernel.inc new file mode 100644 index 0000000..8794f42 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/base/tail_kernel.inc @@ -0,0 +1,492 @@ +#define TL_NT 512 +#define TL_MW 512 +#define TL_MLD TL_MW +#define TL_VW 504 +#define TL_UW 520 +#define TL_BN 496 +#define TL_VLD 520 +#define TL_LBO_U (TL_UW * 8) +#define TL_LBO_M (TL_MLD * 8) +#define TL_LBO_V (TL_VLD * 8) +#define TL_SBO_A 128 +#define TL_PU_E (4 * TL_UW * 8) +#define TL_PM_E (8 * TL_MLD * 8) +#define TL_PV_E (32 * TL_VLD) +#define TL_AI_E (14 * 1024) +#define TL_AS_E (32 * 512 + 512) +#define TL_SHARED ((2 * TL_PU_E + TL_PM_E + TL_PV_E + TL_AI_E + TL_AS_E) * 2) +__device__ __forceinline__ void tl_stmatrix4(void *dst, const uint32_t r[4], + bool trans) { + const uint32_t a = static_cast(__cvta_generic_to_shared(dst)); + if (trans) + asm volatile( + "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1,%2,%3,%4};\n" ::"r"( + a), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3])); + else + asm volatile( + "stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1,%2,%3,%4};\n" ::"r"( + a), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3])); +} +__device__ __forceinline__ void tl_ldmatrix4t(uint32_t r[4], const void *src) { + const uint32_t a = static_cast(__cvta_generic_to_shared(src)); + asm volatile( + "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(a)); +} +__device__ __forceinline__ void tl_stmatrix1t(void *dst, uint32_t r) { + const uint32_t a = static_cast(__cvta_generic_to_shared(dst)); + asm volatile( + "stmatrix.sync.aligned.m8n8.x1.trans.shared.b16 [%0], {%1};\n" :: + "r"(a), "r"(r)); +} +template +__device__ __forceinline__ void tl_skip_band( + const __half *PU, __half *PM, const __half *AS, + __half *PV, int c0, int warp, int lane, int er0, const float *sh_ask, + const float *sh_ari_v, float b0v, float b1v) { + constexpr int NB8 = BW / 8; + float acc[BW / 2]; +#pragma unroll + for (int i = 0; i < BW / 2; ++i) acc[i] = ((i & 2) == 0) ? b0v : b1v; + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int c = 0; c < 4; ++c) { + const uint64_t da0 = + wg_desc(AS + c * 8 * 512, 128u, TL_SBO_A * 2u); + const uint64_t db0 = + wg_desc(PM + (2 * c) * TL_LBO_M + (c0 + 2) * 8, + TL_LBO_M * 2u, 128u); +#pragma unroll + for (int p = 0; p < 4; ++p) { + const uint64_t da = da0 + 64 * (2 * p); + const uint64_t db = db0 + 2 * p; + if constexpr (BW == 120) { + WGMMA_M64N120K16_T(acc, da, db, 0, 0); + } else { + WGMMA_M64N128K16_T(acc, da, db, 0, 0); + } + } + } + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + const int m_addr = lane >> 3, row = lane & 7; + const int cg = (warp & 1) * 2; + if (er0 >= 32) { +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = __floats2half2_rn( + acc[jj * 4 + h * 2 + 0], acc[jj * 4 + h * 2 + 1]); + r[m] = *reinterpret_cast(&hv); + } else { + r[m] = 0; + } + } + const int jj = j + (m_addr >> 1), h = m_addr & 1; + tl_stmatrix4(PV + (cg + h) * TL_LBO_V + (c0 + jj * 8 + row) * 8, + r, true); + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + if (er0 < 32) { +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; + const int jja = j + (m_addr >> 1), ha = m_addr & 1; + tl_ldmatrix4t(r, PV + (cg + ha) * TL_LBO_V + + (c0 + jja * 8 + row + 1) * 8); +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = *reinterpret_cast<__half2 *>(&r[m]); + acc[jj * 4 + h * 2 + 0] += __half2float(hv.x); + acc[jj * 4 + h * 2 + 1] += __half2float(hv.y); + } + } + } +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; + const int jja = j + (m_addr >> 1), ha = m_addr & 1; + tl_ldmatrix4t(r, PU + (cg + ha) * TL_LBO_U + + (c0 + jja * 8 + row + 16) * 8); +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = *reinterpret_cast<__half2 *>(&r[m]); + acc[jj * 4 + h * 2 + 0] += __half2float(hv.x); + acc[jj * 4 + h * 2 + 1] += __half2float(hv.y); + } + } + } + } + if (er0 >= 32) return; + if (er0 < 16) { + const float a0 = sh_ask[er0], r0 = sh_ari_v[er0]; + const float a1 = sh_ask[er0 + 8], r1 = sh_ari_v[er0 + 8]; +#pragma unroll + for (int i = 0; i < BW / 2; ++i) { + const int h = (i & 2) >> 1; + const float v = acc[i]; + const float sn = __sinf((h ? a1 : a0) * v); + acc[i] = __fmaf_rn(sn * sn, h ? r1 : r0, v); + } + } else { +#pragma unroll + for (int i = 0; i < BW / 2; ++i) { + const float v = acc[i]; + acc[i] = v > 0.f ? v : 0.01f * v; + } + } +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = __floats2half2_rn( + acc[jj * 4 + h * 2 + 0], acc[jj * 4 + h * 2 + 1]); + r[m] = *reinterpret_cast(&hv); + } else { + r[m] = 0; + } + } + const int jj = j + (m_addr >> 1), h = m_addr & 1; + tl_stmatrix4(PM + (cg + h) * TL_LBO_V + + (c0 + jj * 8 + row) * 8, + r, true); + } +} +__device__ __noinline__ void tail_fused_kernel( + const __half *__restrict__ U, const __half *__restrict__ win, + const __half *__restrict__ wsk, const __half *__restrict__ wpo, + const float *__restrict__ bin, const float *__restrict__ ain, + const float *__restrict__ bsk, const float *__restrict__ ask, + const float *__restrict__ bpo, const __half *__restrict__ st_m4, + const __half *__restrict__ st_v, __half *__restrict__ wave, + __half *__restrict__ ns_u, __half *__restrict__ ns_m4, + __half *__restrict__ ns_v, int B, int T4, int L4, int SC4, int N4, + int NTILE) { + char *tl_arena=reinterpret_cast(native_arena()); + __half *const PU = reinterpret_cast<__half *>(tl_arena); + __half *const PM = PU + 2 * TL_PU_E; + __half *const PV = PM + TL_PM_E; + __half *const AI = PV + TL_PV_E; + __half *const AS = AI + TL_AI_E; + __half *const PB = AS + 32 * 512; + float *sh_ain=(float*)native_aux(), *sh_ari=sh_ain+32, *sh_ask=sh_ari+32, *sh_ari_v=sh_ask+16; + const int tid = threadIdx.x; + const int wg = tid >> 7; + const int warp = (tid & 127) >> 5; + const int lane = tid & 31; + const int er0 = warp * 16 + (lane >> 2); + const int cb2 = (lane & 3) * 2; + for (int i = tid; i < 32; i += TL_NT) { + const float a = ain[i]; + sh_ain[i] = a; + sh_ari[i] = __frcp_rn(a + 1e-9f); + } + for (int i = tid; i < 16; i += TL_NT) { + const float a = ask[i]; + sh_ask[i] = a; + sh_ari_v[i] = __frcp_rn(a + 1e-9f); + } + for (int i = tid; i < TL_AI_E / 8; i += TL_NT) { + const int blk = i / 128, off = (i - blk * 128) * 8; + const int c = blk / 7, k = blk - c * 7; + *reinterpret_cast(AI + blk * 1024 + off) = + *reinterpret_cast(win + (size_t)c * 14336 + + (size_t)k * 2048 + off); + } + for (int i = tid; i < 32 * 64; i += TL_NT) { + const int blk = i / 64, off = (i - blk * 64) * 8; + const int c = blk >> 3, k = blk & 7; + *reinterpret_cast(AS + blk * 512 + off) = + (k < 7) ? *reinterpret_cast( + wsk + (size_t)c * 14336 + (size_t)k * 2048 + off) + : make_int4(0, 0, 0, 0); + } + for (int i = tid; i < 4 * (TL_VLD - TL_VW) * 8; i += TL_NT) { + const int g = i / ((TL_VLD - TL_VW) * 8); + PV[g * TL_LBO_V + TL_VW * 8 + (i - g * ((TL_VLD - TL_VW) * 8))] = + __float2half(0.f); + } + for (int i = tid; i < 512; i += TL_NT) PB[i] = __float2half(0.f); + for (int i = tid; i < 2 * 16 * 8; i += TL_NT) { + const int c = i / (16 * 8); + const int r = i - c * 16 * 8; + const int ki = r / 8, k = r & 7; + PB[c * 128 + (ki >> 3) * 64 + (ki & 7) * 8 + k] = + (k < 7) ? wpo[(c * 16 + ki) * 7 + k] : __float2half(0.f); + } + const float bias_i0 = bin[er0], bias_i1 = bin[er0 + 8]; + const float bias_s0 = (er0 < 32) ? bsk[er0] : 0.f; + const float bias_s1 = (er0 + 8 < 32) ? bsk[er0 + 8] : 0.f; + const float bpo0 = bpo[0]; + constexpr int UCOL = TL_UW / 32; + constexpr int UUNIT = 4 * UCOL; + constexpr int UWARP = TL_NT / 32; + constexpr int UITER = (UUNIT + UWARP - 1) / UWARP; + const int uwarp = tid >> 5; + const int uq = lane & 3; + int ugc[UITER], uoc[UITER]; +#pragma unroll + for (int t = 0; t < UITER; ++t) { + const int bi = uwarp + t * UWARP; + const int gc = bi / UCOL; + ugc[t] = gc; + uoc[t] = bi - gc * UCOL; + } + const __half *const urow = + U + (size_t)(lane >> 2) * SC4 + 8 * uq; + __syncthreads(); + PDL_ENTER(); +#define TL_STAGE_U(PUD,UB) { _Pragma("unroll") for (int t2 = 0; t2 < UITER; t2 += 2) { uint32_t u[2][4]; _Pragma("unroll") for (int d = 0; d < 2; ++d) { const int t = t2 + d; const int n = (UB) + uoc[t] * 32; const int gn = n + 8 * uq; const bool ok = gn >= 0 && gn < N4; int4 v = make_int4(0, 0, 0, 0); if (ok) v = *reinterpret_cast( urow + (size_t)ugc[t] * (8 * (size_t)SC4) + n); u[d][0] = v.x; u[d][1] = v.y; u[d][2] = v.z; u[d][3] = v.w; } _Pragma("unroll") for (int d = 0; d < 2; ++d) { const int t = t2 + d; uint32_t *q = u[d]; uint32_t x, y; x = (uq & 1) ? q[0] : q[1]; y = (uq & 1) ? q[2] : q[3]; x = __shfl_xor_sync(0xffffffffu, x, 1); y = __shfl_xor_sync(0xffffffffu, y, 1); if (uq & 1) { q[0] = x; q[2] = y; } else { q[1] = x; q[3] = y; } x = (uq & 2) ? q[0] : q[2]; y = (uq & 2) ? q[1] : q[3]; x = __shfl_xor_sync(0xffffffffu, x, 2); y = __shfl_xor_sync(0xffffffffu, y, 2); if (uq & 2) { q[0] = x; q[1] = y; } else { q[2] = x; q[3] = y; } tl_stmatrix4((PUD) + ugc[t] * TL_LBO_U + (uoc[t] * 32 + lane) * 8, q, true); } } if (uwarp < 4) { const int n = (UB) + 512; const int gn = n + 2 * (lane & 3); uint32_t q = 0; if (gn >= 0 && gn + 1 < N4) q = *reinterpret_cast( U + (size_t)(uwarp * 8 + (lane >> 2)) * SC4 + gn); tl_stmatrix1t((PUD) + uwarp * TL_LBO_U + (512 + (lane & 7)) * 8, q); } } + TL_STAGE_U(PU, (int)blockIdx.x * TL_BN - 16) + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + int ubuf = 0; + for (int tile = blockIdx.x; tile < NTILE; tile += gridDim.x, ubuf ^= 1) { + const int g0 = tile * TL_BN; + __half *const PUc = PU + ubuf * TL_PU_E; + __half *const PUn = PU + (ubuf ^ 1) * TL_PU_E; + { + const int b0 = g0 / L4; + for (int s = 0; s < 2; ++s) { + const int sb = b0 + s; + if (sb >= B) break; + const int cv = (int)((long long)sb * L4 + native_true_tokens(sb,T4) - g0); + if (cv < 0 || cv >= TL_BN) continue; + for (int i = tid; i < 32 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + ns_u[((size_t)sb * 32 + co) * 6 + h] = + PUc[(co >> 3) * TL_LBO_U + + (cv + h + 16) * 8 + (co & 7)]; + } + } + } + { + float acc[64]; +#pragma unroll + for (int i = 0; i < 64; ++i) + acc[i] = ((i & 2) == 0) ? bias_i0 : bias_i1; + const int c0 = wg * 128; + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int c = 0; c < 2; ++c) { + const uint64_t da0 = + wg_desc(AI + c * 7 * 1024, 128u, TL_SBO_A * 2u); + const uint64_t db0 = + wg_desc(PUc + (2 * c) * TL_LBO_U + (c0 + 2) * 8, + TL_LBO_U * 2u, 128u); +#pragma unroll + for (int k = 0; k < 7; ++k) { + const uint64_t da = da0 + 128 * k; + const uint64_t db = db0 + k; + WGMMA_M64N128K16_T(acc, da, db, 0, 0); + } + } + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); + if (tile + (int)gridDim.x < NTILE) + TL_STAGE_U(PUn, (tile + (int)gridDim.x) * TL_BN - 16) + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + if (er0 < 32) { + const float a0 = sh_ain[er0], r0 = sh_ari[er0]; + const float a1 = sh_ain[er0 + 8], r1 = sh_ari[er0 + 8]; +#pragma unroll + for (int i = 0; i < 64; ++i) { + const int h = (i & 2) >> 1; + const float v = acc[i]; + const float sn = __sinf((h ? a1 : a0) * v); + acc[i] = __fmaf_rn(sn * sn, h ? r1 : r0, v); + } + } else { +#pragma unroll + for (int i = 0; i < 64; ++i) { + const float v = acc[i]; + acc[i] = v > 0.f ? v : 0.01f * v; + } + } +#pragma unroll + for (int j = 0; j < 16; j += 2) { + uint32_t r[4]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + const __half2 hv = __floats2half2_rn( + acc[jj * 4 + h * 2 + 0], acc[jj * 4 + h * 2 + 1]); + r[m] = *reinterpret_cast(&hv); + } + const int m_addr = lane >> 3, row = lane & 7; + const int jj = j + (m_addr >> 1), h = m_addr & 1; + tl_stmatrix4(PM + (warp * 2 + h) * TL_LBO_M + + (c0 + jj * 8 + row) * 8, + r, true); + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + { + const int base = g0 - 8; + const int b0 = (base >= 0) ? base / L4 : 0; + const int cm0 = (int)((long long)b0 * L4 - base); + const int cm1 = cm0 + L4; + const bool hit0 = (b0 < B) && (cm0 + 6 > 0) && (cm0 < TL_MW); + const bool hit1 = (b0 + 1 < B) && (cm1 + 6 > 0) && (cm1 < TL_MW); + if (hit0 || hit1) { + for (int s = 0; s < 2; ++s) { + if (!(s ? hit1 : hit0)) continue; + const int sb = b0 + s, cm = s ? cm1 : cm0; + for (int i = tid; i < 64 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + const int c = cm + h; + if (c >= 0 && c < TL_MW) + PM[(co >> 3) * TL_LBO_M + c * 8 + (co & 7)] = + st_m4[((size_t)sb * 64 + co) * 6 + h]; + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + } + } + { + const int b0 = g0 / L4; + for (int s = 0; s < 2; ++s) { + const int sb = b0 + s; + if (sb >= B) break; + const int cv = (int)((long long)sb * L4 + native_true_tokens(sb,T4) - g0); + if (cv < 0 || cv >= TL_BN) continue; + for (int i = tid; i < 64 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + ns_m4[((size_t)sb * 64 + co) * 6 + h] = + PM[(co >> 3) * TL_LBO_M + + (cv + h + 8) * 8 + (co & 7)]; + } + } + } + { + const int c0 = wg * 128; + if (wg == 3) + tl_skip_band<120>(PUc, PM, AS, PV, c0, warp, lane, er0, sh_ask, + sh_ari_v, bias_s0, bias_s1); + else + tl_skip_band<128>(PUc, PM, AS, PV, c0, warp, lane, er0, sh_ask, + sh_ari_v, bias_s0, bias_s1); + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + { + const int b0 = g0 / L4; + const int cv0 = (int)((long long)b0 * L4 - g0), cv1 = cv0 + L4; + const bool hv0 = (b0 < B) && (cv0 + 6 > 0) && (cv0 < TL_VW); + const bool hv1 = (b0 + 1 < B) && (cv1 + 6 > 0) && (cv1 < TL_VW); + if (hv0 || hv1) { + for (int s = 0; s < 2; ++s) { + if (!(s ? hv1 : hv0)) continue; + const int sb = b0 + s, cv = s ? cv1 : cv0; + for (int i = tid; i < 32 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + const int c = cv + h; + if (c >= 0 && c < TL_VW) + PM[(co >> 3) * TL_LBO_V + c * 8 + (co & 7)] = + st_v[((size_t)sb * 32 + co) * 6 + h]; + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + } + } + { + const int b0 = g0 / L4; + for (int s = 0; s < 2; ++s) { + const int sb = b0 + s; + if (sb >= B) break; + const int cv = (int)((long long)sb * L4 + native_true_tokens(sb,T4) - g0); + if (cv < 0 || cv >= TL_BN) continue; + for (int i = tid; i < 32 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + ns_v[((size_t)sb * 32 + co) * 6 + h] = + PM[(co >> 3) * TL_LBO_V + (cv + h) * 8 + (co & 7)]; + } + } + } + __syncthreads(); + constexpr int PC_LD = 10; + { + float *const PC = reinterpret_cast(PV); +#pragma unroll + for (int round = 0; round < 2; ++round) { + float acc[4]; + const int p0 = round * 256 + wg * 64; + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); + const uint64_t da0 = + wg_desc(PM + p0 * 8, TL_LBO_V * 2u, 128u); + const uint64_t db0 = wg_desc(PB, 128u, 128u); + const uint64_t da1 = + wg_desc(PM + 2 * TL_LBO_V + p0 * 8, + TL_LBO_V * 2u, 128u); + const uint64_t db1 = wg_desc(PB + 128, 128u, 128u); + WGMMA_M64N8K16_ZERO_T(acc, da0, db0, 0, 1); + WGMMA_M64N8K16_T(acc, da1, db1, 0, 1); + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + const int row0 = p0 + warp * 16 + (lane >> 2); + const int col = (lane & 3) * 2; + *reinterpret_cast(PC + row0 * PC_LD + col) = + make_float2(acc[0], acc[1]); + *reinterpret_cast(PC + (row0 + 8) * PC_LD + col) = + make_float2(acc[2], acc[3]); + } + } + __syncthreads(); + { + const float *const PC = reinterpret_cast(PV); + const int p0 = tid; + const bool live = p0 < TL_BN; + const unsigned live_mask = __ballot_sync(0xffffffffu, live); + if (live) { + const float s0 = __fadd_rn(PC[(p0 + 0) * PC_LD + 0], + PC[(p0 + 1) * PC_LD + 1]); + const float s1 = __fadd_rn(PC[(p0 + 2) * PC_LD + 2], + PC[(p0 + 3) * PC_LD + 3]); + const float s2 = __fadd_rn(PC[(p0 + 4) * PC_LD + 4], + PC[(p0 + 5) * PC_LD + 5]); + const float s3 = PC[(p0 + 6) * PC_LD + 6]; + float y0 = __fadd_rn(__fadd_rn(s0, s1), + __fadd_rn(s2, s3)); + y0 = fminf(fmaxf(y0 + bpo0, -1.f), 1.f); + const float y1 = __shfl_down_sync(live_mask, y0, 1); + if ((p0 & 1) == 0) { + const int b0 = g0 / L4; + const int rr = g0 - b0 * L4; + int t0 = rr + p0, bo0 = b0; + int t1 = t0 + 1, bo1 = bo0; + if (t0 >= L4) { t0 -= L4; ++bo0; } + if (t1 >= L4) { t1 -= L4; ++bo1; } + if (bo0 == bo1 && bo0 < B && t1 < T4) { + *reinterpret_cast<__half2 *>( + wave + (size_t)bo0 * T4 + t0) = + __floats2half2_rn(y0, y1); + } else { + if (bo0 < B && t0 < T4) + wave[(size_t)bo0 * T4 + t0] = __float2half(y0); + if (bo1 < B && t1 < T4) + wave[(size_t)bo1 * T4 + t1] = __float2half(y1); + } + } + } + } + } +#undef TL_STAGE_U +} diff --git a/deploy/nemotron_native/kernels/codec/base/wg2_kernel.inc b/deploy/nemotron_native/kernels/codec/base/wg2_kernel.inc new file mode 100644 index 0000000..4ffa7ad --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/base/wg2_kernel.inc @@ -0,0 +1,280 @@ +__device__ __forceinline__ void cpa16(void *dsm, const void *src, bool pred) { + unsigned a = static_cast(__cvta_generic_to_shared(dsm)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(a), + "l"(src), "r"(pred ? 16 : 0) + : "memory"); +} +template +__device__ __forceinline__ void cpa16o(unsigned a, const void *src, bool pred) { + asm volatile("cp.async.cg.shared.global [%0+%3], [%1], 16, %2;\n" ::"r"(a), + "l"(src), "r"(pred ? 16 : 0), "n"(OFF) + : "memory"); +} +template +__device__ __noinline__ +void wg2_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int BK = CIB * KT; + constexpr int WG_BM = 64 * MSPLIT; + constexpr int WG_BNT = WG_BN * NSPLIT; + constexpr int NTS = 128 * MSPLIT * NSPLIT; + constexpr int WG_NT = NTS * KSPLIT; + constexpr int NCRAW = WG_BNT / 4; + constexpr int EPN = + (WG_NT % NCRAW == 0) ? WG_BNT : (((WG_BNT + 127) / 128) * 128); + constexpr int WG_CLD = EPN + 8; + constexpr bool NB8_P2 = ((WG_BN / 8) & (WG_BN / 8 - 1)) == 0; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int ST_E = SA * A_E + SB * NSPLIT * B_E; + constexpr int SBO_AE = KG * WG_SLOT; + constexpr int LBO_BE = (WG_BN / 8) * WG_SLOT; + constexpr int NB8 = WG_BN / 8; + constexpr bool ALDG = (WG_BN <= 64); + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *smem = reinterpret_cast<__half *>(arena); + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int kgrp = tid / NTS; + const int stid = tid - kgrp * NTS; + const int wgl = stid >> 7; + const int mg = wgl % MSPLIT; + const int ngid = wgl / MSPLIT; + const int warp = (stid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * WG_BM; + const int nb = blockIdx.y * WG_BNT; + const int n0 = nb + ngid * WG_BN; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + const int NCH = (Kg + BK - 1) / BK; + if (tid < WG_BM) { + const int co = m0 + tid; + if (do_snake && co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + for (int i = tid; i < EPN; i += WG_NT) { + const int n = nb + i; + colok[i] = (i < WG_BNT && n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mg * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (kgrp == 0 && eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (kgrp == 0 && eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + constexpr int KGX = ALDG ? (KG <= 8 ? 8 : (KG <= 16 ? 16 : 32)) : KG; + constexpr int ATOT = WG_BM * KGX; + constexpr int AREG = (ATOT + NTS - 1) / NTS; + constexpr int BTOT = NSPLIT * CIB * NB8; + constexpr int BREG = (BTOT + NTS - 1) / NTS; + const __half *aptr[AREG]; + int aoff[AREG]; + int akmax[AREG]; +#pragma unroll + for (int r = 0; r < AREG; ++r) { + const int idx = stid + r * NTS; + const int m = idx / KGX, v = idx - (idx / KGX) * KGX; + const int co = m0 + m; + const bool live = (ATOT % NTS == 0 || idx < ATOT) && (co < Cout) && + v < KG; + aoff[r] = (m >> 3) * SBO_AE + (v < KG ? v : 0) * WG_SLOT + (m & 7) * 8; + akmax[r] = live ? (Kg - v * 8) : 0; + aptr[r] = w + (size_t)(live ? co : 0) * Kg + v * 8; + } + const __half *bptr[BREG]; + int bcimax[BREG], bnc[BREG], bcil[BREG], bgrp[BREG]; +#pragma unroll + for (int r = 0; r < BREG; ++r) { + const int idx = stid + r * NTS; + const int gg = idx / (CIB * NB8); + const int rr = idx - gg * (CIB * NB8); + const int cil = rr / NB8, nc = rr - (rr / NB8) * NB8; + bcil[r] = cil; + bnc[r] = nc; + bgrp[r] = gg; + const bool live = (BTOT % NTS == 0 || idx < BTOT); + bcimax[r] = live ? (Cin - cil) : 0; + bptr[r] = in + (size_t)cil * SCin + nb + gg * WG_BN + nc * 8; + } + int4 br0[BREG], br1[BREG]; + int4 ar[ALDG ? AREG : 1]; + uint32_t edge[BREG]; + __half *const kaB = smem + kgrp * ST_E; + __half *const kbB = kaB + SA * A_E; + const unsigned barid = 1u + kgrp + (unsigned)(KSPLIT * native_subid()); +#define WG2_BAR() asm volatile("bar.sync %0, %1;\n" ::"r"(barid), "r"(NTS)) +#define WG2_LOADB(c) { const int _ci0 = (c) * CIB; _Pragma("unroll") for (int r = 0; r < BREG; ++r) { br0[r] = make_int4(0, 0, 0, 0); if constexpr (KT == 3) edge[r] = 0; else br1[r] = make_int4(0, 0, 0, 0); if (_ci0 < bcimax[r]) { const int4 *src = reinterpret_cast( bptr[r] + (size_t)_ci0 * SCin); br0[r] = src[0]; if constexpr (KT == 3) { if (!NB8_P2 || bnc[r] == NB8 - 1) edge[r] = src[1].x; } else { br1[r] = src[1]; } } } } +#define WG2_LOADA(c) if constexpr (ALDG) { const int _kg0 = (c) * BK; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { ar[r] = make_int4(0, 0, 0, 0); if (_kg0 < akmax[r]) ar[r] = *reinterpret_cast(aptr[r] + _kg0); } } +#define WG2_PUTA(c,AD) { _Pragma("unroll") for (int r = 0; r < AREG; ++r) { if constexpr (ALDG) { if (akmax[r] > 0) *reinterpret_cast((AD) + aoff[r]) = ar[r]; } else { const int _kg0 = (c) * BK; if (akmax[r] > 0) cpa16((AD) + aoff[r], aptr[r] + _kg0, _kg0 < akmax[r]); } } if constexpr (!ALDG) asm volatile("cp.async.commit_group;\n" ::: "memory"); } +#define WG2_STSB(BD) { _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int cil = bcil[r], nc = bnc[r]; unsigned int u[8]; u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; if constexpr (KT == 3) { if constexpr (NB8_P2) { u[4] = __shfl_down_sync(0xffffffffu, u[0], 1, NB8); if (nc == NB8 - 1) u[4] = edge[r]; } else { u[4] = edge[r]; } u[5] = u[6] = u[7] = 0; } else { u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; } if (BTOT % NTS == 0 || (stid + r * NTS) < BTOT) { _Pragma("unroll") for (int k = 0; k < KT; ++k) { const int kh = k >> 1; int4 o; if (k & 1) { o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); } else { o.x = u[kh + 0]; o.y = u[kh + 1]; o.z = u[kh + 2]; o.w = u[kh + 3]; } const int row = cil * KT + k; *reinterpret_cast((BD) + bgrp[r] * B_E + nc * WG_SLOT + (row >> 3) * LBO_BE + (row & 7) * 8) = o; } } } } +#define WG2_MMA(AD,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t da0 = wg_desc((AD) + mg * 8 * SBO_AE, WG_SLOT * 2, SBO_AE * 2); const uint64_t db0 = wg_desc((BD) + ngid * B_E, LBO_BE * 2, WG_SLOT * 2); _Pragma("unroll") for (int kk = 0; kk < BK; kk += 16) { const uint64_t da = da0 + (kk >> 3) * (WG_SLOT / 8); const uint64_t db = db0 + (kk >> 3) * (LBO_BE / 8); if constexpr (WG_BN == 64) { WGMMA_M64N64K16(acc, da, db); } else if constexpr (WG_BN == 256) { WGMMA_M64N256K16(acc, da, db); } else if constexpr (WG_BN == 104) { WGMMA_M64N104K16(acc, da, db); } else if constexpr (WG_BN == 112) { WGMMA_M64N112K16(acc, da, db); } else if constexpr (WG_BN == 120) { WGMMA_M64N120K16(acc, da, db); } else if constexpr (WG_BN == 224) { WGMMA_M64N224K16(acc, da, db); } else { static_assert(WG_BN == 128, "no wgmma macro for WG_BN"); WGMMA_M64N128K16(acc, da, db); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + constexpr int SMAX = SA > SB ? SA : SB; + WG2_LOADA(kgrp) + WG2_LOADB(kgrp) +#pragma unroll 1 + for (int s = 0; s < SMAX; ++s) { + const int c = s * KSPLIT + kgrp; + if (s < SA) { + if (c < NCH) WG2_PUTA(c, kaB + s * A_E) + if (c + KSPLIT < NCH) WG2_LOADA(c + KSPLIT) + } + if (s < SB) { + if (c < NCH) WG2_STSB(kbB + s * NSPLIT * B_E) + if (c + KSPLIT < NCH) WG2_LOADB(c + KSPLIT) + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + const int NIT = (NCH - kgrp + KSPLIT - 1) / KSPLIT; + for (int i = 0; i < NIT; ++i) { + if constexpr (!ALDG) + asm volatile("cp.async.wait_group %0;\n" ::"n"(SA - 2) : "memory"); + WG2_BAR(); + WG2_MMA(kaB + (i % SA) * A_E, kbB + (i % SB) * NSPLIT * B_E) + asm volatile("wgmma.wait_group.sync.aligned 1;\n" ::: "memory"); + if constexpr (NTS > 128) WG2_BAR(); + if (i >= 1) { + const int ca = (i + SA - 1) * KSPLIT + kgrp; + if (ca < NCH) WG2_PUTA(ca, kaB + ((i - 1) % SA) * A_E) + if (ALDG && ca + KSPLIT < NCH) WG2_LOADA(ca + KSPLIT) + const int cb = (i + SB - 1) * KSPLIT + kgrp; + if (cb < NCH) { + WG2_STSB(kbB + ((i - 1) % SB) * NSPLIT * B_E) + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if (cb + KSPLIT < NCH) WG2_LOADB(cb + KSPLIT) + } else if constexpr (ALDG) { + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + } + } + } + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); +#undef WG2_LOADB +#undef WG2_LOADA +#undef WG2_PUTA +#undef WG2_STSB +#undef WG2_MMA +#undef WG2_BAR + __syncthreads(); + float *Ck = Cs + kgrp * WG_BM * WG_CLD; + const int cb = (lane & 3) * 2; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Ck + (er0 + h * 8) * WG_CLD + + ngid * WG_BN + j * 8 + cb) = v; + } + } + __syncthreads(); + constexpr int NC = EPN / 4; + constexpr int MSTEP = WG_NT / NC; + constexpr int MITER = WG_BM / MSTEP; + constexpr int UNR = (MITER % 4 == 0) ? 4 : ((MITER % 2 == 0) ? 2 : 1); + static_assert(MITER % UNR == 0, "row batch must divide the row count"); + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + const int n = nb + lc; +#pragma unroll 1 + for (int t0 = 0; t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + v[u] = *reinterpret_cast(Cs + m * WG_CLD + lc); + if (resid != nullptr) { + const int co = m0 + m; + const int cc = co < Cout ? co : 0; + const __half2 *rp = reinterpret_cast( + resid + (size_t)cc * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + const int co = m0 + m; + if (co >= Cout) continue; + float4 vv = v[u]; +#pragma unroll + for (int g = 1; g < KSPLIT; ++g) { + float4 p = *reinterpret_cast( + Cs + g * WG_BM * WG_CLD + m * WG_CLD + lc); + vv.x += p.x; vv.y += p.y; vv.z += p.z; vv.w += p.w; + } + if (resid != nullptr) { + vv.x += __half2float(r0[u].x); vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (do_snake) { + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + } + __half *op = out + (size_t)co * SCout + n + H; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B * H; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int total = Cout * B * pad_len; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} diff --git a/deploy/nemotron_native/kernels/codec/base/wg3_kernel.inc b/deploy/nemotron_native/kernels/codec/base/wg3_kernel.inc new file mode 100644 index 0000000..fa55b32 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/base/wg3_kernel.inc @@ -0,0 +1,177 @@ +template +__device__ __noinline__ void wg3_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int WG_BNT = WG_BN * NSG; + constexpr int WG_CLD = WG_BNT + 8; + constexpr int WG_NT = 128 * NWG * NSG; + constexpr int WG_BM = 64 * NWG; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int SBO_AE = KG * WG_SLOT; + constexpr int LBO_BE = (WG_BN / 8) * WG_SLOT; + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + constexpr int ST_E = A_E + NSG * B_E; + __half *As0 = reinterpret_cast<__half *>(arena); + __half *Bs0 = As0 + A_E; + __half *As1 = (STAGES == 2) ? As0 + ST_E : As0; + __half *Bs1 = (STAGES == 2) ? Bs0 + ST_E : Bs0; + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int wgid = tid >> 7; + const int mgid = wgid % NWG, ngid = wgid / NWG; + const int warp = (tid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * WG_BM; + constexpr int WGROWS = 8; + const int nb = blockIdx.y * WG_BNT; + const int n0 = nb + ngid * WG_BN; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + if (tid < WG_BM) { + const int co = m0 + tid; + if (do_snake && co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + if (tid < WG_BNT) { + const int n = nb + tid; + colok[tid] = (n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mgid * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + constexpr int ATOT = WG_BM * KG; + constexpr int AREG = (ATOT + WG_NT - 1) / WG_NT; + constexpr int CIB = BK / KT; + constexpr int BTOT = NSG * CIB * (WG_BN / 8); + constexpr int BREG = (BTOT + WG_NT - 1) / WG_NT; + int4 ar[AREG]; + int4 br0[BREG], br1[BREG]; + uint32_t edge[BREG]; +#define WG_LOADG(kg0) { const int _ci0 = (kg0) / KT; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; ar[r] = make_int4(0, 0, 0, 0); if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; const int co = m0 + m, kg = (kg0) + v * 8; if (co < Cout && kg < Kg) ar[r] = *reinterpret_cast( w + (size_t)co * Kg + kg); } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; br0[r] = make_int4(0, 0, 0, 0); if constexpr (KT == 3) edge[r] = 0; else br1[r] = make_int4(0, 0, 0, 0); if (BTOT % WG_NT == 0 || idx < BTOT) { const int _g = idx / (CIB * (WG_BN / 8)); const int _r = idx - _g * (CIB * (WG_BN / 8)); const int ci = _ci0 + _r / (WG_BN / 8); const int nc = _r % (WG_BN / 8); if (ci < Cin) { const int4 *src = reinterpret_cast( in + (size_t)ci * SCin + nb + _g * WG_BN + nc * 8); br0[r] = src[0]; if constexpr (KT == 3) { if (nc == WG_BN / 8 - 1) edge[r] = src[1].x; } else { br1[r] = src[1]; } } } } } + WG_LOADG(0) +#define WG_STS(AD,BD) { _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; *reinterpret_cast( (AD) + (m >> 3) * SBO_AE + v * WG_SLOT + (m & 7) * 8) = ar[r]; } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; if (BTOT % WG_NT == 0 || idx < BTOT) { const int _g = idx / (CIB * (WG_BN / 8)); const int _r = idx - _g * (CIB * (WG_BN / 8)); const int cil = _r / (WG_BN / 8), nc = _r % (WG_BN / 8); unsigned int u[8]; u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; if constexpr (KT == 3) { u[4] = __shfl_down_sync(0xffffffffu, u[0], 1, WG_BN / 8); if (nc == WG_BN / 8 - 1) u[4] = edge[r]; u[5] = u[6] = u[7] = 0; } else { u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; } _Pragma("unroll") for (int k = 0; k < KT; ++k) { const int kh = k >> 1; int4 o; if (k & 1) { o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); } else { o.x = u[kh + 0]; o.y = u[kh + 1]; o.z = u[kh + 2]; o.w = u[kh + 3]; } const int row = cil * KT + k; *reinterpret_cast((BD) + _g * B_E + nc * WG_SLOT + (row >> 3) * LBO_BE + (row & 7) * 8) = o; } } } if constexpr (WG_BN == 64) { asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); } } +#define WG_MMA(AD,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t da0 = wg_desc((AD) + mgid * WGROWS * SBO_AE, WG_SLOT * 2, SBO_AE * 2); const uint64_t db0 = wg_desc((BD) + ngid * B_E, LBO_BE * 2, WG_SLOT * 2); _Pragma("unroll") for (int kk = 0; kk < BK; kk += 16) { const uint64_t da = da0 + (kk >> 3) * (WG_SLOT / 8); const uint64_t db = db0 + (kk >> 3) * (LBO_BE / 8); if constexpr (WG_BN == 64) { WGMMA_M64N64K16(acc, da, db); } else { WGMMA_M64N128K16(acc, da, db); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + WG_STS(As0, Bs0) + if (BK < Kg) WG_LOADG(BK) + __syncthreads(); + int cur = 0; + for (int kg0 = 0; kg0 < Kg; kg0 += BK) { + if (cur == 0) { WG_MMA(As0, Bs0) } else { WG_MMA(As1, Bs1) } + if (kg0 + BK < Kg) { + if (STAGES == 2) + asm volatile("wgmma.wait_group.sync.aligned 1;\n" ::: "memory"); + else + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + __syncthreads(); + if (cur == 0) { WG_STS(As1, Bs1) } else { WG_STS(As0, Bs0) } + if (kg0 + 2 * BK < Kg) WG_LOADG(kg0 + 2 * BK) + __syncthreads(); + } else { + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + } + if (STAGES == 2) cur ^= 1; + } +#undef WG_LOADG +#undef WG_STS +#undef WG_MMA + __syncthreads(); + const int r0 = er0; + const int cb = (lane & 3) * 2; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Cs + (r0 + h * 8) * WG_CLD + + ngid * WG_BN + j * 8 + cb) = v; + } + } + __syncthreads(); + constexpr int NC = WG_BNT / 4; + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + for (int m = lr; m < WG_BM; m += WG_NT / NC) { + const int co = m0 + m; + if (co >= Cout) continue; + const int n = nb + lc; + float4 v = *reinterpret_cast(Cs + m * WG_CLD + lc); + __half *op = out + (size_t)co * SCout + n + H; + if (resid != nullptr) { + const __half2 *rp = reinterpret_cast( + resid + (size_t)co * SCout + n + H); + __half2 r0h = rp[0], r1h = rp[1]; + v.x += __half2float(r0h.x); v.y += __half2float(r0h.y); + v.z += __half2float(r1h.x); v.w += __half2float(r1h.y); + } + if (do_snake) { + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + } + if (ok4) { + __half2 o0 = __floats2half2_rn(v.x, v.y); + __half2 o1 = __floats2half2_rn(v.z, v.w); + reinterpret_cast<__half2 *>(op)[0] = o0; + reinterpret_cast<__half2 *>(op)[1] = o1; + } else { + const float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B * H; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int total = Cout * B * pad_len; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} diff --git a/deploy/nemotron_native/kernels/codec/base/wg4_step_widths.inc b/deploy/nemotron_native/kernels/codec/base/wg4_step_widths.inc new file mode 100644 index 0000000..9fe1878 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/base/wg4_step_widths.inc @@ -0,0 +1,459 @@ +template +__device__ __noinline__ void wg4_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int CoutP, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + if constexpr (PDL && !LATE_PDL) PDL_ENTER(); + constexpr int NT = 128 * MSPLIT * NSPLIT; + constexpr int BM = 64 * MSPLIT; + constexpr int BNT = WG_BN * NSPLIT; + constexpr int NCRAW = BNT / 4; + constexpr int EPN = (NT % NCRAW == 0) ? BNT : (((BNT + 127) / 128) * 128); + constexpr int WG_CLD = EPN + 8; + constexpr int CG = CIB / 8; + constexpr int NBR = BNT + KT - 1; + constexpr int NBRP = BC8 ? ((NBR + 7) & ~7) : ((NBR + 31) & ~31); + constexpr int UCOL = NBRP / 32; + constexpr int NUNIT = CG * UCOL; + constexpr int NWARP = NT / 32; + constexpr int BITER = (NUNIT + NWARP - 1) / NWARP; + constexpr int B_E = CIB * NBRP; + constexpr int ST_E = B_E; + constexpr int LBO_B = NBRP * 8; + constexpr int ABUF = ABUF_; + static_assert(!OC8 || PREACT, + "K=7 C8 output is only used for residual-free pair heads"); + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *smem = reinterpret_cast<__half *>(arena); + __half *Hs = smem; + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int wgid = tid >> 7; + const int mg = wgid % MSPLIT, ng = wgid / MSPLIT; + const int warp = (tid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * BM; + const int nb = blockIdx.y * BNT; + const int H = KT - 1; + const int half = Cout >> 1; + const int NCH = (Cin + CIB - 1) / CIB; + for (int i = tid; i < BM; i += NT) { + const int co = m0 + i; + if (co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[i] = a; + row_ri[i] = __frcp_rn(a + 1e-9f); + } + } + for (int i = tid; i < EPN; i += NT) { + const int n = nb + i; + colok[i] = (i < BNT && n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mg * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + const int mtile = blockIdx.x * MSPLIT + mg; + const __half *const apt = w + ((size_t)mtile * NCH * KT) * 1024 + + (size_t)(tid & 127) * 8; + uint32_t areg[ABUF][KT][4]; + const int wi = tid >> 5; + int bdst[BITER], boff[BITER]; + bool blive[BITER]; + const __half *const bbase = in + (size_t)(lane >> 2) * SCin + nb + + 8 * (lane & 3); +#pragma unroll + for (int t = 0; t < BITER; ++t) { + const int u = wi + t * NWARP; + const int gc = (u < NUNIT) ? u / UCOL : 0; + const int cu = (u < NUNIT) ? u - gc * UCOL : 0; + blive[t] = u < NUNIT; + bdst[t] = gc * LBO_B + (cu * 32 + lane) * 8; + boff[t] = gc * 8 * SCin + cu * 32; + } + const size_t SC8i = (size_t)SCin * 8; + constexpr int BQ = B_E / 8; + constexpr int BCREG = (BQ + NT - 1) / NT; + int cdst[BCREG]; + const __half *csrc[BCREG]; + bool clive[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) { + const int q = tid + r * NT; + const int gc = q / NBRP, rem = q - gc * NBRP; + clive[r] = ((BQ % NT == 0) || q < BQ) && gc < CG; + cdst[r] = clive[r] ? gc * LBO_B + rem * 8 : 0; + csrc[r] = in + (clive[r] ? (size_t)gc * SC8i + + (size_t)nb * 8 + rem * 8 : 0); + } + uint32_t br[BITER][4]; + constexpr int CPA_STAGES = 4; + constexpr int BSTAGES = BC8 ? CPA_STAGES : STAGES; + static_assert(STAGES >= CPA_STAGES, "C8 B-panel pipeline allocation"); + static_assert(ABUF == 2 && BSTAGES == 4, "wg4 unrolled pipeline shape"); + uint64_t dslot[BSTAGES]; +#pragma unroll + for (int s = 0; s < BSTAGES; ++s) + dslot[s] = wg_desc(smem + s * ST_E + ng * WG_BN * 8, LBO_B * 2u, 128u); + uint32_t csm0[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) + csm0[r] = static_cast(__cvta_generic_to_shared(smem + cdst[r])); + uint32_t bsm0[BITER]; +#pragma unroll + for (int t = 0; t < BITER; ++t) + bsm0[t] = static_cast(__cvta_generic_to_shared(smem + bdst[t])); + const size_t CSTEP = (size_t)CG * SC8i; + const __half *crol[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) crol[r] = csrc[r] + CPA_STAGES * CSTEP; + const __half *arol = apt + (size_t)ABUF * (KT * 1024); +#define WG4_CPB(c,BD) { const size_t _o = (size_t)(c) * CG * SC8i; _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16((BD) + cdst[r], csrc[r] + _o, true); } +#define WG4_CPB_R(SLOT) { _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16o<(SLOT) * ST_E * 2>(csm0[r], crol[r], true); _Pragma("unroll") for (int r = 0; r < BCREG; ++r) crol[r] += CSTEP; } +#define WG4_LOADB(c) { const size_t _o = (size_t)(c) * CIB * SCin; _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { const int4 _v = *reinterpret_cast( bbase + _o + boff[t]); br[t][0] = _v.x; br[t][1] = _v.y; br[t][2] = _v.z; br[t][3] = _v.w; } } } +#define WG4_XPOSE(u) { const int _q = lane & 3; uint32_t _x, _y; _x = (_q & 1) ? u[0] : u[1]; _y = (_q & 1) ? u[2] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 1); _y = __shfl_xor_sync(0xffffffffu, _y, 1); if (_q & 1) { u[0] = _x; u[2] = _y; } else { u[1] = _x; u[3] = _y; } _x = (_q & 2) ? u[0] : u[2]; _y = (_q & 2) ? u[1] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 2); _y = __shfl_xor_sync(0xffffffffu, _y, 2); if (_q & 2) { u[0] = _x; u[1] = _y; } else { u[2] = _x; u[3] = _y; } } +#define WG4_STOREB(BD) { _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { uint32_t _u[4]; _u[0] = br[t][0]; _u[1] = br[t][1]; _u[2] = br[t][2]; _u[3] = br[t][3]; WG4_XPOSE(_u) const uint32_t _a = static_cast( __cvta_generic_to_shared((BD) + bdst[t])); asm volatile( "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 " "[%0], {%1,%2,%3,%4};\n" ::"r"(_a), "r"(_u[0]), "r"(_u[1]), "r"(_u[2]), "r"(_u[3])); } } } +#define WG4_STOREB_R(SLOT) { _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { uint32_t _u[4]; _u[0] = br[t][0]; _u[1] = br[t][1]; _u[2] = br[t][2]; _u[3] = br[t][3]; WG4_XPOSE(_u) asm volatile( "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 " "[%0+%5], {%1,%2,%3,%4};\n" ::"r"(bsm0[t]), "r"(_u[0]), "r"(_u[1]), "r"(_u[2]), "r"(_u[3]), "n"((SLOT) * ST_E * 2)); } } } +#define WG4_LOADA_R(BUF) { _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { const int4 _v = *reinterpret_cast(arol + kk * 1024); areg[BUF][kk][0] = _v.x; areg[BUF][kk][1] = _v.y; areg[BUF][kk][2] = _v.z; areg[BUF][kk][3] = _v.w; } arol += KT * 1024; } +#define WG4_MMA_R(BUF,SLOT) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { WGMMA_RS_ANY_T(WG_BN, acc, areg[BUF][kk], dslot[SLOT] + kk, 0); } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } +#define WG4_LOADA(BUF,c) { const __half *_p = apt + (size_t)(c) * (KT * 1024); _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { const int4 _v = *reinterpret_cast(_p + kk * 1024); areg[BUF][kk][0] = _v.x; areg[BUF][kk][1] = _v.y; areg[BUF][kk][2] = _v.z; areg[BUF][kk][3] = _v.w; } } +#define WG4_KEEPA(BUF) { _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) asm volatile("" : "+r"(areg[BUF][kk][0]), "+r"(areg[BUF][kk][1]), "+r"(areg[BUF][kk][2]), "+r"(areg[BUF][kk][3])); } +#define WG4_MMA(BUF,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t db0 = wg_desc((BD) + ng * WG_BN * 8, LBO_B * 2u, 128u); _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { WGMMA_RS_ANY_T(WG_BN, acc, areg[BUF][kk], db0 + kk, 0); } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + if constexpr (PDL && LATE_PDL) PDL_ENTER(); + constexpr int WGD = 1; + if constexpr (BC8) { +#pragma unroll 1 + for (int st = 0; st < CPA_STAGES; ++st) { + if (st < NCH) WG4_CPB(st, smem + st * ST_E) + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } + WG4_LOADA(0, 0) + if (ABUF > 1 && 1 < NCH) WG4_LOADA(1, 1) + if (ABUF > 2 && 2 < NCH) WG4_LOADA(2, 2) + } else { + WG4_LOADB(0) + WG4_LOADA(0, 0) + if (ABUF > 1 && 1 < NCH) WG4_LOADA(1, 1) + if (ABUF > 2 && 2 < NCH) WG4_LOADA(2, 2) + WG4_STOREB(smem) + if (1 < NCH) WG4_LOADB(1) + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + } +#define WG4_BODY(BUF,i,SLOT) { if constexpr (BC8) { asm volatile("cp.async.wait_group %0;\n" : : "n"(CPA_STAGES - 3) : "memory"); __syncthreads(); if ((i) >= 2) { if ((i) + 2 < NCH) { WG4_CPB_R(((SLOT) + 2) % CPA_STAGES) } asm volatile("cp.async.commit_group;\n" ::: "memory"); } } else { if (i) __syncthreads(); } WG4_MMA_R(BUF, SLOT) asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(WGD) : "memory"); if constexpr (NSPLIT > 1) __syncthreads(); if constexpr (!BC8) { const int cb = (i) + 1; if (cb < NCH) { WG4_STOREB_R(((SLOT) + 1) % BSTAGES) asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); if (cb + 1 < NCH) WG4_LOADB(cb + 1) } } if ((i) >= 1) { const int nb2 = ((BUF) + ABUF - 1) % ABUF; if ((i) + ABUF - 1 < NCH) { WG4_KEEPA(nb2) WG4_LOADA_R(nb2) } } } + for (int i = 0; i < NCH; i += 4) { + WG4_BODY(0, i, 0) + if (i + 1 < NCH) WG4_BODY(1, i + 1, 1) + if (i + 2 < NCH) WG4_BODY(0, i + 2, 2) + if (i + 3 < NCH) WG4_BODY(1, i + 3, 3) + } + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + if constexpr (BC8) + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); +#undef WG4_CPB +#undef WG4_CPB_R +#undef WG4_STOREB_R +#undef WG4_LOADA_R +#undef WG4_MMA_R +#undef WG4_LOADB +#undef WG4_XPOSE +#undef WG4_STOREB +#undef WG4_LOADA +#undef WG4_KEEPA +#undef WG4_MMA +#undef WG4_BODY +#define TL4_SW(c,r) ((c) * BM + (((r) & ~7) ^ (((c) & 7) << 3)) + ((r) & 7)) + __syncthreads(); + const int cb2 = (lane & 3) * 2; + const bool snake = PREACT && (m0 + er0) < half; + const float sn_a0 = snake ? row_alpha[er0] : 0.f; + const float sn_r0 = snake ? row_ri[er0] : 0.f; + const float sn_a1 = snake ? row_alpha[er0 + 8] : 0.f; + const float sn_r1 = snake ? row_ri[er0 + 8] : 0.f; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + const int row = er0 + h * 8; + if constexpr (PREACT) { + if (snake) { + const float a = h ? sn_a1 : sn_a0; + const float ri = h ? sn_r1 : sn_r0; + const float s0 = __sinf(a * v.x); + const float s1 = __sinf(a * v.y); + v.x = __fmaf_rn(s0 * s0, ri, v.x); + v.y = __fmaf_rn(s1 * s1, ri, v.y); + } else { + v.x = v.x > 0.f ? v.x : 0.01f * v.x; + v.y = v.y > 0.f ? v.y : 0.01f * v.y; + } + if constexpr (OC8) { + const int c0 = ng * WG_BN + j * 8 + cb2; + Hs[TL4_SW(c0, row)] = __float2half(v.x); + Hs[TL4_SW(c0 + 1, row)] = __float2half(v.y); + } else { + *reinterpret_cast<__half2 *>( + Hs + row * WG_CLD + ng * WG_BN + j * 8 + cb2) = + __floats2half2_rn(v.x, v.y); + } + } else { + *reinterpret_cast( + Cs + row * WG_CLD + ng * WG_BN + j * 8 + cb2) = v; + } + } + } + __syncthreads(); + if constexpr (PREACT) { + if constexpr (OC8) { + constexpr int GRP = BM / 8; + constexpr int TOTG = GRP * BNT; + constexpr int GITER = (TOTG + NT - 1) / NT; + const size_t SC8o = (size_t)SCout * 8; +#pragma unroll + for (int t = 0; t < GITER; ++t) { + const int idx = tid + t * NT; + if (TOTG % NT != 0 && idx >= TOTG) continue; + const int g = idx / BNT; + const int j = idx - g * BNT; + if (!colok[j]) continue; + const int m0g = g * 8; + if (m0 + m0g >= Cout) continue; + *reinterpret_cast( + out + (size_t)((m0 + m0g) >> 3) * SC8o + + (size_t)(nb + j + H) * 8) = + *reinterpret_cast(Hs + TL4_SW(j, m0g)); + } + } else { + static_assert((BNT & 3) == 0, "four-column preactivated epilogue"); + constexpr int NC = BNT / 4; + constexpr int TOT = BM * NC; + for (int q = tid; q < TOT; q += NT) { + const int m = q / NC; + const int lc = (q - m * NC) * 4; + const int co = m0 + m; + const bool any4 = (colok[lc] | colok[lc + 1] | + colok[lc + 2] | colok[lc + 3]) != 0; + if (co >= Cout || !any4) continue; + const __half2 *sp = reinterpret_cast( + Hs + m * WG_CLD + lc); + const __half2 v0 = sp[0], v1 = sp[1]; + __half *op = out + (size_t)co * SCout + nb + lc + H; + const bool ok4 = (colok[lc] & colok[lc + 1] & + colok[lc + 2] & colok[lc + 3]) != 0; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = v0; + reinterpret_cast<__half2 *>(op)[1] = v1; + } else { + if (colok[lc + 0]) op[0] = v0.x; + if (colok[lc + 1]) op[1] = v0.y; + if (colok[lc + 2]) op[2] = v1.x; + if (colok[lc + 3]) op[3] = v1.y; + } + } + } + } else { + constexpr bool COMPACT_EPI = + (WG_BN == 136 || WG_BN == 160 || WG_BN == 208); + if constexpr (COMPACT_EPI) { + static_assert((BNT & 3) == 0, "4-column epilogue groups"); + constexpr int NC = BNT / 4; + constexpr int TOT = BM * NC; + constexpr int MITER = (TOT + NT - 1) / NT; + constexpr int UNR = (MITER >= 8) ? 8 : ((MITER >= 4) ? 4 : + ((MITER >= 2) ? 2 : 1)); +#pragma unroll 1 + for (int t0 = 0; t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int q = tid + (t0 + u) * NT; + if (q >= TOT) continue; + const int m = q / NC; + const int lc = (q - m * NC) * 4; + const int co = m0 + m; + const bool any4 = (colok[lc] | colok[lc + 1] | + colok[lc + 2] | colok[lc + 3]) != 0; + if (co >= Cout || !any4) continue; + v[u] = *reinterpret_cast(Cs + m * WG_CLD + lc); + if constexpr (!PREACT) { + const int n = nb + lc; + const __half2 *rp = reinterpret_cast( + resid + (size_t)co * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int q = tid + (t0 + u) * NT; + if (q >= TOT) continue; + const int m = q / NC; + const int lc = (q - m * NC) * 4; + const int co = m0 + m; + const bool any4 = (colok[lc] | colok[lc + 1] | + colok[lc + 2] | colok[lc + 3]) != 0; + if (co >= Cout || !any4) continue; + const int n = nb + lc; + float4 vv = v[u]; + if constexpr (!PREACT) { + vv.x += __half2float(r0[u].x); vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (co < half) { + const float a = row_alpha[m], ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + __half *op = out + (size_t)co * SCout + n + H; + const bool ok4 = (colok[lc] & colok[lc + 1] & + colok[lc + 2] & colok[lc + 3]) != 0; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + } else { + constexpr int NC = EPN / 4; + constexpr int MSTEP = NT / NC; + constexpr int MITER = BM / MSTEP; + constexpr int UNR = (MITER % 8 == 0) ? 8 + : ((MITER % 4 == 0) ? 4 + : ((MITER % 2 == 0) ? 2 : 1)); + static_assert(MITER % UNR == 0, "row batch must divide row count"); + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && + colok[lc + 2] && colok[lc + 3]; + const bool live4 = colok[lc] || colok[lc + 1] || + colok[lc + 2] || colok[lc + 3]; + const int n = nb + lc; +#pragma unroll 1 + for (int t0 = 0; live4 && t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + v[u] = *reinterpret_cast( + Cs + m * WG_CLD + lc); + if constexpr (!PREACT) { + const int co = m0 + m; + const int cc = co < Cout ? co : 0; + const __half2 *rp = reinterpret_cast( + resid + (size_t)cc * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + const int co = m0 + m; + if (co >= Cout) continue; + float4 vv = v[u]; + if constexpr (!PREACT) { + vv.x += __half2float(r0[u].x); + vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); + vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (co < half) { + const float a = row_alpha[m], ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + __half *op = out + (size_t)co * SCout + n + H; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B; + for (int idx = bxy * NT + tid; idx < total; idx += nbx * NT) { + const int b = idx / Cin, ci = idx - b * Cin; + const int x = b * L + native_true_tokens(b,T); + __half *const dst = ns + ((size_t)b * Cin + ci) * H; + if constexpr (BC8) { + const __half *const src = + in + (size_t)(ci >> 3) * ((size_t)SCin * 8) + + (size_t)x * 8 + (ci & 7); +#pragma unroll + for (int j = 0; j < H; ++j) dst[j] = src[j * 8]; + } else { + const __half *const src = in + (size_t)ci * SCin + x; +#pragma unroll + for (int j = 0; j < H; j += 2) + *reinterpret_cast(dst + j) = + *reinterpret_cast(src + j); + } + } + } + if (pad_src != nullptr) { + const int total = Cout * B; + for (int idx = bxy * NT + tid; idx < total; idx += nbx * NT) { + const int b = idx / Cout, co = idx - b * Cout; + const int x = b * L + pad_off; + const __half *const src = + pad_src + ((size_t)b * Cout + co) * pad_len; + if constexpr (OC8) { + __half *const dst = + out + (size_t)(co >> 3) * ((size_t)SCout * 8) + + (size_t)x * 8 + (co & 7); + for (int j = 0; j < pad_len; ++j) dst[j * 8] = src[j]; + } else { + __half *const dst = out + (size_t)co * SCout + x; + for (int j = 0; j < pad_len; ++j) dst[j] = src[j]; + } + } + } + } +} diff --git a/deploy/nemotron_native/kernels/codec/base/wg5_step_widths.inc b/deploy/nemotron_native/kernels/codec/base/wg5_step_widths.inc new file mode 100644 index 0000000..467481f --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/base/wg5_step_widths.inc @@ -0,0 +1,400 @@ +template +__device__ __noinline__ +void wg5_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int KT = 3; + constexpr int BM = 64 * MSPLIT; + constexpr int BNT = WG_BN; + constexpr int NTS = 128 * MSPLIT; + constexpr int WG_NT = NTS * KSPLIT; + constexpr int NCRAW = BNT / 4; + constexpr int EPN = + (WG_NT % NCRAW == 0) ? BNT : (((BNT + 127) / 128) * 128); + constexpr int WG_CLD_R = (OC8 ? WG_BN : EPN) + 8; + constexpr int WG_CLD = (WG_CLD_R % 16 == 0) ? WG_CLD_R + 4 : WG_CLD_R; + constexpr int CG = CIB / 8; + constexpr int CSUB = CIB / 16; + constexpr int AK = KT * CSUB; + constexpr int NBRP = BC8 ? (((BNT + KT - 1) + 7) & ~7) + : (((BNT + KT - 1) + 31) & ~31); + constexpr int UCOL = NBRP / 32; + constexpr int NUNIT = CG * UCOL; + constexpr int NWARP = NTS / 32; + constexpr int BITER = (NUNIT + NWARP - 1) / NWARP; + constexpr int B_E = CIB * NBRP; + constexpr int ST_E = SB * B_E; + constexpr int LBO_B = NBRP * 8; + constexpr int ABUF = 3; + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *smem = reinterpret_cast<__half *>(arena); + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int kgrp = tid / NTS; + const int stid = tid - kgrp * NTS; + const int mg = stid >> 7; + const int warp = (stid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * BM; + const int nb = blockIdx.y * BNT; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + constexpr int NCH = (CIN + CIB - 1) / CIB; + static_assert(NCH % KSPLIT == 0, "wg5 K-split must divide the chunk count"); + if (tid < BM) { + const int co = m0 + tid; + if (co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + for (int i = tid; i < EPN; i += WG_NT) { + const int n = nb + i; + colok[i] = (i < BNT && n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mg * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (kgrp == 0 && eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (kgrp == 0 && eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + const int mtile = blockIdx.x * MSPLIT + mg; + const __half *const apt = w + ((size_t)mtile * NCH * AK) * 1024 + + (size_t)(stid & 127) * 8; + uint32_t areg[ABUF][AK][4]; + const int wi = stid >> 5; + int bdst[BITER], boff[BITER]; + bool blive[BITER]; + const __half *const bbase = in + (size_t)(lane >> 2) * SCin + nb + + 8 * (lane & 3); +#pragma unroll + for (int t = 0; t < BITER; ++t) { + const int u = wi + t * NWARP; + const int gc = (u < NUNIT) ? u / UCOL : 0; + const int cu = (u < NUNIT) ? u - gc * UCOL : 0; + blive[t] = u < NUNIT; + bdst[t] = gc * LBO_B + (cu * 32 + lane) * 8; + boff[t] = gc * 8 * SCin + cu * 32; + } + const size_t SC8i = (size_t)SCin * 8; + constexpr int BQ = B_E / 8; + constexpr int BCREG = (BQ + NTS - 1) / NTS; + int cdst[BCREG]; + const __half *csrc[BCREG]; + bool clive[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) { + const int q = stid + r * NTS; + const int gc = q / NBRP, rem = q - gc * NBRP; + clive[r] = ((BQ % NTS == 0) || q < BQ) && gc < CG; + cdst[r] = clive[r] ? gc * LBO_B + rem * 8 : 0; + csrc[r] = in + (clive[r] ? (size_t)gc * SC8i + + (size_t)nb * 8 + rem * 8 : 0); + } + uint32_t br[BITER][4]; + __half *const kbB = smem + kgrp * ST_E; + const unsigned barid = 1u + kgrp + (unsigned)(KSPLIT * native_subid()); + uint64_t dbuf[ABUF]; +#pragma unroll + for (int s = 0; s < ABUF; ++s) + dbuf[s] = wg_desc(kbB + s * B_E, LBO_B * 2u, 128u); + uint32_t csm0[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) + csm0[r] = static_cast( + __cvta_generic_to_shared(kbB + cdst[r])); + const size_t CSTEP = (size_t)KSPLIT * CG * SC8i; + constexpr int ASTEP = KSPLIT * AK * 1024; + const __half *crol[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) + crol[r] = csrc[r] + (size_t)(ABUF * KSPLIT + kgrp) * CG * SC8i; + const __half *arol = apt + (size_t)(ABUF * KSPLIT + kgrp) * (AK * 1024); + int crolc = ABUF * KSPLIT + kgrp; +#define WG5_BAR() asm volatile("bar.sync %0, %1;\n" ::"r"(barid), "r"(NTS)) +#define WG5_CPB(c,BD) { const size_t _o = (size_t)(c) * CG * SC8i; _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16((BD) + cdst[r], csrc[r] + _o, true); } +#define WG5_LOADA(BUF,c) { const __half *_p = apt + (size_t)(c) * (AK * 1024); _Pragma("unroll") for (int kk = 0; kk < AK; ++kk) { const int4 _v = *reinterpret_cast(_p + kk * 1024); areg[BUF][kk][0] = _v.x; areg[BUF][kk][1] = _v.y; areg[BUF][kk][2] = _v.z; areg[BUF][kk][3] = _v.w; } } +#define WG5_KEEPA(BUF) { _Pragma("unroll") for (int kk = 0; kk < AK; ++kk) asm volatile("" : "+r"(areg[BUF][kk][0]), "+r"(areg[BUF][kk][1]), "+r"(areg[BUF][kk][2]), "+r"(areg[BUF][kk][3])); } +#define WG5_LOADB(c) { const size_t _o = (size_t)(c) * CIB * SCin; _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { const int4 _v = *reinterpret_cast( bbase + _o + boff[t]); br[t][0] = _v.x; br[t][1] = _v.y; br[t][2] = _v.z; br[t][3] = _v.w; } } } +#define WG5_BXPOSE(u) { const int _q = lane & 3; uint32_t _x, _y; _x = (_q & 1) ? u[0] : u[1]; _y = (_q & 1) ? u[2] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 1); _y = __shfl_xor_sync(0xffffffffu, _y, 1); if (_q & 1) { u[0] = _x; u[2] = _y; } else { u[1] = _x; u[3] = _y; } _x = (_q & 2) ? u[0] : u[2]; _y = (_q & 2) ? u[1] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 2); _y = __shfl_xor_sync(0xffffffffu, _y, 2); if (_q & 2) { u[0] = _x; u[1] = _y; } else { u[2] = _x; u[3] = _y; } } +#define WG5_STOREB(BD) { _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { uint32_t _u[4]; _u[0] = br[t][0]; _u[1] = br[t][1]; _u[2] = br[t][2]; _u[3] = br[t][3]; WG5_BXPOSE(_u) const uint32_t _a = static_cast( __cvta_generic_to_shared((BD) + bdst[t])); asm volatile( "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 " "[%0], {%1,%2,%3,%4};\n" ::"r"(_a), "r"(_u[0]), "r"(_u[1]), "r"(_u[2]), "r"(_u[3])); } } } +#define WG5_MMA(BUF) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); _Pragma("unroll") for (int h = 0; h < CSUB; ++h) { _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { WGMMA_RS_ANY_T(WG_BN, acc, areg[BUF][h * KT + kk], dbuf[BUF] + kk + h * (LBO_B / 4), 0); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } +#define WG5_CPB_R(SLOT) { _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16o<(SLOT) * B_E * 2>(csm0[r], crol[r], true); } +#define WG5_LOADA_R(BUF) { _Pragma("unroll") for (int kk = 0; kk < AK; ++kk) { const int4 _v = *reinterpret_cast(arol + kk * 1024); areg[BUF][kk][0] = _v.x; areg[BUF][kk][1] = _v.y; areg[BUF][kk][2] = _v.z; areg[BUF][kk][3] = _v.w; } } +#define WG5_ROLL() { crolc += KSPLIT; arol += ASTEP; if constexpr (BC8) { _Pragma("unroll") for (int r = 0; r < BCREG; ++r) crol[r] += CSTEP; } } + if constexpr (PDL) { + (void)0; + if (tid == 0) + (void)0; + } + static_assert(SB == 3 && WGD == 1 && ABUF == 3, "wg5 pipeline shape"); + if constexpr (BC8) { +#pragma unroll 1 + for (int s = 0; s < SB; ++s) { + const int c = s * KSPLIT + kgrp; + if (c < NCH) WG5_CPB(c, kbB + s * B_E) + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } + } else { + WG5_LOADB(kgrp) +#pragma unroll 1 + for (int s = 0; s < SB; ++s) { + const int c = s * KSPLIT + kgrp; + if (c < NCH) WG5_STOREB(kbB + s * B_E) + if (c + KSPLIT < NCH) WG5_LOADB(c + KSPLIT) + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + } + WG5_LOADA(0, kgrp) + if (KSPLIT + kgrp < NCH) WG5_LOADA(1, KSPLIT + kgrp) + if (2 * KSPLIT + kgrp < NCH) WG5_LOADA(2, 2 * KSPLIT + kgrp) + constexpr int NIT = NCH / KSPLIT; + static_assert(NIT % ABUF == 0, "wg5 mainloop unroll shape"); +#define WG5_BODY(BUF,i) { if constexpr (BC8) asm volatile("cp.async.wait_group %0;\n" : : "n"(SB - 2) : "memory"); WG5_BAR(); WG5_MMA(BUF) asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(WGD) : "memory"); if constexpr (NTS > 128) WG5_BAR(); if ((i) >= 1) { const int nb2 = ((BUF) + 2) % ABUF; if (crolc < NCH) { if constexpr (BC8) { WG5_CPB_R(nb2) } else { WG5_STOREB(kbB + nb2 * B_E) asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); if (crolc + KSPLIT < NCH) WG5_LOADB(crolc + KSPLIT) } WG5_KEEPA(nb2) WG5_LOADA_R(nb2) } WG5_ROLL() if constexpr (BC8) asm volatile("cp.async.commit_group;\n" ::: "memory"); } } +#pragma unroll + for (int i = 0; i < NIT; i += 3) { + WG5_BODY(0, i) + WG5_BODY(1, i + 1) + WG5_BODY(2, i + 2) + } + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + if constexpr (BC8) + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); +#undef WG5_BXPOSE +#undef WG5_CPB +#undef WG5_CPB_R +#undef WG5_LOADA_R +#undef WG5_ROLL +#undef WG5_LOADB +#undef WG5_STOREB +#undef WG5_MMA +#undef WG5_LOADA +#undef WG5_KEEPA +#undef WG5_BODY +#undef WG5_BAR + __syncthreads(); + const int cb2 = (lane & 3) * 2; + if constexpr (OC8) { + __half *Ck = smem + kgrp * BM * WG_CLD; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + *reinterpret_cast<__half2 *>(Ck + (er0 + h * 8) * WG_CLD + + j * 8 + cb2) = + __floats2half2_rn(acc[j * 4 + h * 2], + acc[j * 4 + h * 2 + 1]); + } + } + } else { + float *Ck = Cs + kgrp * BM * WG_CLD; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], + acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Ck + (er0 + h * 8) * WG_CLD + + j * 8 + cb2) = v; + } + } + } + __syncthreads(); + if constexpr (OC8) { + constexpr int GRP = BM / 8; + constexpr int TOTG = GRP * WG_BN; + constexpr int GITER = (TOTG + WG_NT - 1) / WG_NT; + const size_t SC8o = (size_t)SCout * 8; +#pragma unroll + for (int t = 0; t < GITER; ++t) { + const int idx = tid + t * WG_NT; + if (TOTG % WG_NT != 0 && idx >= TOTG) continue; + const int j = idx - (idx / WG_BN) * WG_BN, g = idx / WG_BN; + if (!colok[j]) continue; + const int m0g = g * 8; + if (m0 + m0g >= Cout) continue; + const int n = nb + j; + __half rr[8]; + if constexpr (CIN == 1536) { + if constexpr (RC8) { + *reinterpret_cast(rr) = + *reinterpret_cast( + resid + (size_t)((m0 + m0g) >> 3) * SC8o + + (size_t)(n + H) * 8); + } else { +#pragma unroll + for (int h = 0; h < 8; ++h) + rr[h] = resid[(size_t)(m0 + m0g + h) * SCout + n + H]; + } + } + uint32_t pk[4]; + const bool snake = (m0 + m0g) < half; +#pragma unroll + for (int h = 0; h < 8; h += 2) { + float f[2]; +#pragma unroll + for (int e = 0; e < 2; ++e) { + const int m = m0g + h + e; + float a = __half2float(smem[m * WG_CLD + j]); +#pragma unroll + for (int kg = 1; kg < KSPLIT; ++kg) + a += __half2float( + smem[kg * BM * WG_CLD + m * WG_CLD + j]); + if constexpr (CIN == 1536) + a += __half2float(rr[h + e]); + if (snake) { + const float sn = __sinf(row_alpha[m] * a); + a = __fmaf_rn(sn * sn, row_ri[m], a); + } else { + a = a > 0.f ? a : 0.01f * a; + } + f[e] = a; + } + const __half2 p = __floats2half2_rn(f[0], f[1]); + pk[h >> 1] = *reinterpret_cast(&p); + } + *reinterpret_cast( + out + (size_t)((m0 + m0g) >> 3) * SC8o + + (size_t)(n + H) * 8) = make_int4(pk[0], pk[1], pk[2], pk[3]); + } + } else { + constexpr int NC = EPN / 4; + constexpr int MSTEP = WG_NT / NC; + constexpr int MITER = BM / MSTEP; + constexpr int UNR = (MITER % 4 == 0) ? 4 : ((MITER % 2 == 0) ? 2 : 1); + static_assert(MITER % UNR == 0, "row batch must divide the row count"); + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + const bool live4 = colok[lc] || colok[lc + 1] || colok[lc + 2] || + colok[lc + 3]; + const int n = nb + lc; +#pragma unroll 1 + for (int t0 = 0; live4 && t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + v[u] = *reinterpret_cast(Cs + m * WG_CLD + lc); + if constexpr (CIN == 1536) { + const int co = m0 + m; + const int cc = co < Cout ? co : 0; + if constexpr (RC8) { + const __half *rp = + resid + (size_t)(cc >> 3) * ((size_t)SCout * 8) + + (size_t)(n + H) * 8 + (cc & 7); + r0[u].x = rp[0]; r0[u].y = rp[8]; + r1[u].x = rp[16]; r1[u].y = rp[24]; + } else { + const __half2 *rp = reinterpret_cast( + resid + (size_t)cc * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + const int co = m0 + m; + if (co >= Cout) continue; + float4 vv = v[u]; +#pragma unroll + for (int g = 1; g < KSPLIT; ++g) { + float4 p = *reinterpret_cast( + Cs + g * BM * WG_CLD + m * WG_CLD + lc); + vv.x += p.x; vv.y += p.y; vv.z += p.z; vv.w += p.w; + } + if constexpr (CIN == 1536) { + vv.x += __half2float(r0[u].x); vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + __half *op = out + (size_t)co * SCout + n + H; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + static_assert(CIN == 768 || CIN == 1536); + constexpr int COUT = (CIN == 768) ? 1536 : 768; + if (ns != nullptr) { + const int total = CIN * B; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int b = idx / CIN, ci = idx - b * CIN; + const int x = b * L + native_true_tokens(b,T); + __half *const dst = ns + ((size_t)b * CIN + ci) * H; + if constexpr (BC8) { + const __half *const src = + in + (size_t)(ci >> 3) * ((size_t)SCin * 8) + + (size_t)x * 8 + (ci & 7); +#pragma unroll + for (int j = 0; j < H; ++j) dst[j] = src[j * 8]; + } else { + const __half *const src = in + (size_t)ci * SCin + x; +#pragma unroll + for (int j = 0; j < H; j += 2) + *reinterpret_cast(dst + j) = + *reinterpret_cast(src + j); + } + } + } + if (pad_src != nullptr) { + const int total = COUT * B; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int b = idx / COUT, co = idx - b * COUT; + const int x = b * L + pad_off; + const __half *const src = + pad_src + ((size_t)b * COUT + co) * pad_len; + if constexpr (OC8) { + __half *const dst = + out + (size_t)(co >> 3) * ((size_t)SCout * 8) + + (size_t)x * 8 + (co & 7); + for (int j = 0; j < pad_len; ++j) dst[j * 8] = src[j]; + } else { + __half *const dst = out + (size_t)co * SCout + x; + for (int j = 0; j < pad_len; ++j) dst[j] = src[j]; + } + } + } + } +} diff --git a/deploy/nemotron_native/kernels/codec/base/wg_kernel.inc b/deploy/nemotron_native/kernels/codec/base/wg_kernel.inc new file mode 100644 index 0000000..2d5c766 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/base/wg_kernel.inc @@ -0,0 +1,182 @@ +__device__ __forceinline__ uint64_t wg_desc(const void *p, uint32_t lbo_bytes, + uint32_t sbo_bytes) { + uint32_t a = static_cast(__cvta_generic_to_shared(p)); + uint64_t d = static_cast((a & 0x3FFFFu) >> 4); + d |= static_cast((lbo_bytes & 0x3FFFFu) >> 4) << 16; + d |= static_cast((sbo_bytes & 0x3FFFFu) >> 4) << 32; + return d; +} +#include "wgmma_macro.inc" +constexpr int WG_SLOT = 72; +template +__device__ __noinline__ void wg_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int WG_CLD = WG_BN + 8; + constexpr int WG_NT = 128 * NWG; + constexpr int WG_BM = 64 * NWG; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int SBO_AE = KG * WG_SLOT; + constexpr int LBO_BE = (WG_BN / 8) * WG_SLOT; + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *As0 = reinterpret_cast<__half *>(arena); + __half *Bs0 = As0 + A_E; + __half *As1 = (STAGES == 2) ? Bs0 + B_E : As0; + __half *Bs1 = (STAGES == 2) ? Bs0 + B_E + A_E : Bs0; + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int wgid = tid >> 7; + const int warp = (tid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * WG_BM; + constexpr int WGROWS = 8; + const int n0 = blockIdx.y * WG_BN; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + if (tid < WG_BM) { + const int co = m0 + tid; + if (do_snake && co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + if (tid < WG_BN) { + const int n = n0 + tid; + colok[tid] = (n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = wgid * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + constexpr int ATOT = WG_BM * KG; + constexpr int AREG = (ATOT + WG_NT - 1) / WG_NT; + constexpr int CIB = BK / KT; + constexpr int BTOT = CIB * (WG_BN / 8); + constexpr int BREG = (BTOT + WG_NT - 1) / WG_NT; + int4 ar[AREG]; + int4 br0[BREG], br1[BREG]; + uint32_t edge[BREG]; +#define WG_LOADG(kg0) { const int _ci0 = (kg0) / KT; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; ar[r] = make_int4(0, 0, 0, 0); if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; const int co = m0 + m, kg = (kg0) + v * 8; if (co < Cout && kg < Kg) ar[r] = *reinterpret_cast( w + (size_t)co * Kg + kg); } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; br0[r] = make_int4(0, 0, 0, 0); if constexpr (KT == 3) edge[r] = 0; else br1[r] = make_int4(0, 0, 0, 0); if (BTOT % WG_NT == 0 || idx < BTOT) { const int ci = _ci0 + idx / (WG_BN / 8); const int nc = idx % (WG_BN / 8); if (ci < Cin) { const int4 *src = reinterpret_cast( in + (size_t)ci * SCin + n0 + nc * 8); br0[r] = src[0]; if constexpr (KT == 3) { if (nc == WG_BN / 8 - 1) edge[r] = src[1].x; } else { br1[r] = src[1]; } } } } } + WG_LOADG(0) +#define WG_STS(AD,BD) { _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; *reinterpret_cast( (AD) + (m >> 3) * SBO_AE + v * WG_SLOT + (m & 7) * 8) = ar[r]; } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; if (BTOT % WG_NT == 0 || idx < BTOT) { const int cil = idx / (WG_BN / 8), nc = idx % (WG_BN / 8); unsigned int u[8]; u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; if constexpr (KT == 3) { u[4] = __shfl_down_sync(0xffffffffu, u[0], 1, WG_BN / 8); if (nc == WG_BN / 8 - 1) u[4] = edge[r]; u[5] = u[6] = u[7] = 0; } else { u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; } _Pragma("unroll") for (int k = 0; k < KT; ++k) { const int kh = k >> 1; int4 o; if (k & 1) { o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); } else { o.x = u[kh + 0]; o.y = u[kh + 1]; o.z = u[kh + 2]; o.w = u[kh + 3]; } const int row = cil * KT + k; *reinterpret_cast((BD) + nc * WG_SLOT + (row >> 3) * LBO_BE + (row & 7) * 8) = o; } } } if constexpr (WG_BN == 64) { asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); } } +#define WG_MMA(AD,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t da0 = wg_desc((AD) + wgid * WGROWS * SBO_AE, WG_SLOT * 2, SBO_AE * 2); const uint64_t db0 = wg_desc((BD), LBO_BE * 2, WG_SLOT * 2); _Pragma("unroll") for (int kk = 0; kk < BK; kk += 16) { const uint64_t da = da0 + (kk >> 3) * (WG_SLOT / 8); const uint64_t db = db0 + (kk >> 3) * (LBO_BE / 8); if constexpr (WG_BN == 64) { WGMMA_M64N64K16(acc, da, db); } else { WGMMA_M64N128K16(acc, da, db); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + WG_STS(As0, Bs0) + if (BK < Kg) WG_LOADG(BK) + __syncthreads(); + int cur = 0; + for (int kg0 = 0; kg0 < Kg; kg0 += BK) { + if (cur == 0) { WG_MMA(As0, Bs0) } else { WG_MMA(As1, Bs1) } + if (kg0 + BK < Kg) { + if (STAGES == 2) + asm volatile("wgmma.wait_group.sync.aligned 1;\n" ::: "memory"); + else + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + __syncthreads(); + if (cur == 0) { WG_STS(As1, Bs1) } else { WG_STS(As0, Bs0) } + if (kg0 + 2 * BK < Kg) WG_LOADG(kg0 + 2 * BK) + __syncthreads(); + } else { + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + } + cur ^= 1; + } +#undef WG_LOADG +#undef WG_STS +#undef WG_MMA + __syncthreads(); + const int r0 = er0; + const int cb = (lane & 3) * 2; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Cs + (r0 + h * 8) * WG_CLD + j * 8 + cb) = v; + } + } + __syncthreads(); + constexpr int NC = WG_BN / 4; + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + for (int m = lr; m < WG_BM; m += WG_NT / NC) { + const int co = m0 + m; + if (co >= Cout) continue; + const int n = n0 + lc; + float4 v = *reinterpret_cast(Cs + m * WG_CLD + lc); + __half *op = out + (size_t)co * SCout + n + H; + if (resid != nullptr) { + const __half2 *rp = reinterpret_cast( + resid + (size_t)co * SCout + n + H); + __half2 r0h = rp[0], r1h = rp[1]; + v.x += __half2float(r0h.x); v.y += __half2float(r0h.y); + v.z += __half2float(r1h.x); v.w += __half2float(r1h.y); + } + if (do_snake) { + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + } + if (ok4) { + __half2 o0 = __floats2half2_rn(v.x, v.y); + __half2 o1 = __floats2half2_rn(v.z, v.w); + reinterpret_cast<__half2 *>(op)[0] = o0; + reinterpret_cast<__half2 *>(op)[1] = o1; + } else { + const float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B * H; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int total = Cout * B * pad_len; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} diff --git a/deploy/nemotron_native/kernels/codec/base/wgmma_macro.inc b/deploy/nemotron_native/kernels/codec/base/wgmma_macro.inc new file mode 100644 index 0000000..db624e9 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/base/wgmma_macro.inc @@ -0,0 +1,102 @@ +#define WGMMA_M64N8K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " "{" "%0,%1,%2,%3}" ", %4, %5, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N8K16(D,DA,DB) WGMMA_M64N8K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N8K16_ZERO_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " "{" "%0,%1,%2,%3}" ", %4, %5, p, 1, 1, " #TA ", " #TB ";\n}\n" : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) : "l"(DA), "l"(DB), "r"(0)) +#define WGMMA_M64N16K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %10, 0;\n" "wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7}" ", %8, %9, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N16K16(D,DA,DB) WGMMA_M64N16K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N24K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %14, 0;\n" "wgmma.mma_async.sync.aligned.m64n24k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11}" ", %12, %13, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N24K16(D,DA,DB) WGMMA_M64N24K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N32K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %18, 0;\n" "wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15}" ", %16, %17, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N32K16(D,DA,DB) WGMMA_M64N32K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N40K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %22, 0;\n" "wgmma.mma_async.sync.aligned.m64n40k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19}" ", %20, %21, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N40K16(D,DA,DB) WGMMA_M64N40K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N48K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %26, 0;\n" "wgmma.mma_async.sync.aligned.m64n48k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N48K16(D,DA,DB) WGMMA_M64N48K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N56K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %30, 0;\n" "wgmma.mma_async.sync.aligned.m64n56k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N56K16(D,DA,DB) WGMMA_M64N56K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N64K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %34, 0;\n" "wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31}" ", %32, %33, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N64K16(D,DA,DB) WGMMA_M64N64K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N72K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %38, 0;\n" "wgmma.mma_async.sync.aligned.m64n72k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35}" ", %36, %37, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N72K16(D,DA,DB) WGMMA_M64N72K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N80K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %42, 0;\n" "wgmma.mma_async.sync.aligned.m64n80k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39}" ", %40, %41, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N80K16(D,DA,DB) WGMMA_M64N80K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N88K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %46, 0;\n" "wgmma.mma_async.sync.aligned.m64n88k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43}" ", %44, %45, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N88K16(D,DA,DB) WGMMA_M64N88K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N96K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %50, 0;\n" "wgmma.mma_async.sync.aligned.m64n96k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47}" ", %48, %49, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N96K16(D,DA,DB) WGMMA_M64N96K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N104K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %54, 0;\n" "wgmma.mma_async.sync.aligned.m64n104k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51}" ", %52, %53, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N104K16(D,DA,DB) WGMMA_M64N104K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N112K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %58, 0;\n" "wgmma.mma_async.sync.aligned.m64n112k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55}" ", %56, %57, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N112K16(D,DA,DB) WGMMA_M64N112K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N120K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %62, 0;\n" "wgmma.mma_async.sync.aligned.m64n120k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59}" ", %60, %61, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N120K16(D,DA,DB) WGMMA_M64N120K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N128K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %66, 0;\n" "wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63}" ", %64, %65, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N128K16(D,DA,DB) WGMMA_M64N128K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N136K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %70, 0;\n" "wgmma.mma_async.sync.aligned.m64n136k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67}" ", %68, %69, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N136K16(D,DA,DB) WGMMA_M64N136K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N144K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %74, 0;\n" "wgmma.mma_async.sync.aligned.m64n144k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71}" ", %72, %73, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N144K16(D,DA,DB) WGMMA_M64N144K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N152K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %78, 0;\n" "wgmma.mma_async.sync.aligned.m64n152k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75}" ", %76, %77, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N152K16(D,DA,DB) WGMMA_M64N152K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N160K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %82, 0;\n" "wgmma.mma_async.sync.aligned.m64n160k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79}" ", %80, %81, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N160K16(D,DA,DB) WGMMA_M64N160K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N168K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %86, 0;\n" "wgmma.mma_async.sync.aligned.m64n168k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83}" ", %84, %85, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N168K16(D,DA,DB) WGMMA_M64N168K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N176K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %90, 0;\n" "wgmma.mma_async.sync.aligned.m64n176k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87}" ", %88, %89, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N176K16(D,DA,DB) WGMMA_M64N176K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N184K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %94, 0;\n" "wgmma.mma_async.sync.aligned.m64n184k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91}" ", %92, %93, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N184K16(D,DA,DB) WGMMA_M64N184K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N192K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %98, 0;\n" "wgmma.mma_async.sync.aligned.m64n192k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95}" ", %96, %97, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N192K16(D,DA,DB) WGMMA_M64N192K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N200K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %102, 0;\n" "wgmma.mma_async.sync.aligned.m64n200k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99}" ", %100, %101, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N200K16(D,DA,DB) WGMMA_M64N200K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N208K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %106, 0;\n" "wgmma.mma_async.sync.aligned.m64n208k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103}" ", %104, %105, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N208K16(D,DA,DB) WGMMA_M64N208K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N216K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %110, 0;\n" "wgmma.mma_async.sync.aligned.m64n216k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107}" ", %108, %109, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N216K16(D,DA,DB) WGMMA_M64N216K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N224K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %114, 0;\n" "wgmma.mma_async.sync.aligned.m64n224k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111}" ", %112, %113, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N224K16(D,DA,DB) WGMMA_M64N224K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N232K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %118, 0;\n" "wgmma.mma_async.sync.aligned.m64n232k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115}" ", %116, %117, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N232K16(D,DA,DB) WGMMA_M64N232K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N240K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %122, 0;\n" "wgmma.mma_async.sync.aligned.m64n240k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119}" ", %120, %121, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N240K16(D,DA,DB) WGMMA_M64N240K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N248K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %126, 0;\n" "wgmma.mma_async.sync.aligned.m64n248k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123}" ", %124, %125, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N248K16(D,DA,DB) WGMMA_M64N248K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N256K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %130, 0;\n" "wgmma.mma_async.sync.aligned.m64n256k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123,%124,%125,%126,%127}" ", %128, %129, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]), "+f"(D[124]), "+f"(D[125]), "+f"(D[126]), "+f"(D[127]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N256K16(D,DA,DB) WGMMA_M64N256K16_T(D, DA, DB, 0, 1) +#define WGMMA_ANY_T(NW,D,DA,DB,TA,TB) do { if constexpr (NW == 8) { WGMMA_M64N8K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 16) { WGMMA_M64N16K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 24) { WGMMA_M64N24K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 32) { WGMMA_M64N32K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 40) { WGMMA_M64N40K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 48) { WGMMA_M64N48K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 56) { WGMMA_M64N56K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 64) { WGMMA_M64N64K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 72) { WGMMA_M64N72K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 80) { WGMMA_M64N80K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 88) { WGMMA_M64N88K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 96) { WGMMA_M64N96K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 104) { WGMMA_M64N104K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 112) { WGMMA_M64N112K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 120) { WGMMA_M64N120K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 128) { WGMMA_M64N128K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 136) { WGMMA_M64N136K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 144) { WGMMA_M64N144K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 152) { WGMMA_M64N152K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 160) { WGMMA_M64N160K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 168) { WGMMA_M64N168K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 176) { WGMMA_M64N176K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 184) { WGMMA_M64N184K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 192) { WGMMA_M64N192K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 200) { WGMMA_M64N200K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 208) { WGMMA_M64N208K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 216) { WGMMA_M64N216K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 224) { WGMMA_M64N224K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 232) { WGMMA_M64N232K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 240) { WGMMA_M64N240K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 248) { WGMMA_M64N248K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 256) { WGMMA_M64N256K16_T(D, DA, DB, TA, TB); } } while (0) +#define WGMMA_ANY(NW,D,DA,DB) WGMMA_ANY_T(NW, D, DA, DB, 0, 1) +#define WGMMA_RS_M64N96K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %53, 0;\n" "wgmma.mma_async.sync.aligned.m64n96k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47}" ", {%48,%49,%50,%51}, %52, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N104K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %57, 0;\n" "wgmma.mma_async.sync.aligned.m64n104k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51}" ", {%52,%53,%54,%55}, %56, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N112K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %61, 0;\n" "wgmma.mma_async.sync.aligned.m64n112k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55}" ", {%56,%57,%58,%59}, %60, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N120K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %65, 0;\n" "wgmma.mma_async.sync.aligned.m64n120k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59}" ", {%60,%61,%62,%63}, %64, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N128K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %69, 0;\n" "wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63}" ", {%64,%65,%66,%67}, %68, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_ANY_T(NW,D,A,DB,TB) do { if constexpr (NW == 96) { WGMMA_RS_M64N96K16_T(D, A, DB, TB); } else if constexpr (NW == 104) { WGMMA_RS_M64N104K16_T(D, A, DB, TB); } else if constexpr (NW == 112) { WGMMA_RS_M64N112K16_T(D, A, DB, TB); } else if constexpr (NW == 120) { WGMMA_RS_M64N120K16_T(D, A, DB, TB); } else if constexpr (NW == 128) { WGMMA_RS_M64N128K16_T(D, A, DB, TB); } } while (0) +#define WGMMA_RS_M64N8K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %9, 0;\n" "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " "{" "%0,%1,%2,%3}" ", {%4,%5,%6,%7}, %8, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N16K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %13, 0;\n" "wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7}" ", {%8,%9,%10,%11}, %12, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N24K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %17, 0;\n" "wgmma.mma_async.sync.aligned.m64n24k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11}" ", {%12,%13,%14,%15}, %16, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N32K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %21, 0;\n" "wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15}" ", {%16,%17,%18,%19}, %20, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N40K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %25, 0;\n" "wgmma.mma_async.sync.aligned.m64n40k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N48K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %29, 0;\n" "wgmma.mma_async.sync.aligned.m64n48k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N56K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %33, 0;\n" "wgmma.mma_async.sync.aligned.m64n56k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N64K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %37, 0;\n" "wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31}" ", {%32,%33,%34,%35}, %36, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N72K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %41, 0;\n" "wgmma.mma_async.sync.aligned.m64n72k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35}" ", {%36,%37,%38,%39}, %40, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N80K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %45, 0;\n" "wgmma.mma_async.sync.aligned.m64n80k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39}" ", {%40,%41,%42,%43}, %44, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N88K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %49, 0;\n" "wgmma.mma_async.sync.aligned.m64n88k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43}" ", {%44,%45,%46,%47}, %48, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N136K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %73, 0;\n" "wgmma.mma_async.sync.aligned.m64n136k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67}" ", {%68,%69,%70,%71}, %72, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N144K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %77, 0;\n" "wgmma.mma_async.sync.aligned.m64n144k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71}" ", {%72,%73,%74,%75}, %76, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N152K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %81, 0;\n" "wgmma.mma_async.sync.aligned.m64n152k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75}" ", {%76,%77,%78,%79}, %80, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N160K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %85, 0;\n" "wgmma.mma_async.sync.aligned.m64n160k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79}" ", {%80,%81,%82,%83}, %84, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N168K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %89, 0;\n" "wgmma.mma_async.sync.aligned.m64n168k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83}" ", {%84,%85,%86,%87}, %88, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N176K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %93, 0;\n" "wgmma.mma_async.sync.aligned.m64n176k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87}" ", {%88,%89,%90,%91}, %92, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N184K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %97, 0;\n" "wgmma.mma_async.sync.aligned.m64n184k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91}" ", {%92,%93,%94,%95}, %96, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N192K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %101, 0;\n" "wgmma.mma_async.sync.aligned.m64n192k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95}" ", {%96,%97,%98,%99}, %100, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N200K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %105, 0;\n" "wgmma.mma_async.sync.aligned.m64n200k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99}" ", {%100,%101,%102,%103}, %104, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N208K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %109, 0;\n" "wgmma.mma_async.sync.aligned.m64n208k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103}" ", {%104,%105,%106,%107}, %108, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N216K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %113, 0;\n" "wgmma.mma_async.sync.aligned.m64n216k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107}" ", {%108,%109,%110,%111}, %112, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N224K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %117, 0;\n" "wgmma.mma_async.sync.aligned.m64n224k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111}" ", {%112,%113,%114,%115}, %116, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N232K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %121, 0;\n" "wgmma.mma_async.sync.aligned.m64n232k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115}" ", {%116,%117,%118,%119}, %120, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N240K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %125, 0;\n" "wgmma.mma_async.sync.aligned.m64n240k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119}" ", {%120,%121,%122,%123}, %124, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N248K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %129, 0;\n" "wgmma.mma_async.sync.aligned.m64n248k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123}" ", {%124,%125,%126,%127}, %128, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N256K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %133, 0;\n" "wgmma.mma_async.sync.aligned.m64n256k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123,%124,%125,%126,%127}" ", {%128,%129,%130,%131}, %132, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]), "+f"(D[124]), "+f"(D[125]), "+f"(D[126]), "+f"(D[127]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#undef WGMMA_RS_ANY_T +#define WGMMA_RS_ANY_T(NW,D,A,DB,TB) do { if constexpr (NW == 8) { WGMMA_RS_M64N8K16_T(D, A, DB, TB); } else if constexpr (NW == 16) { WGMMA_RS_M64N16K16_T(D, A, DB, TB); } else if constexpr (NW == 24) { WGMMA_RS_M64N24K16_T(D, A, DB, TB); } else if constexpr (NW == 32) { WGMMA_RS_M64N32K16_T(D, A, DB, TB); } else if constexpr (NW == 40) { WGMMA_RS_M64N40K16_T(D, A, DB, TB); } else if constexpr (NW == 48) { WGMMA_RS_M64N48K16_T(D, A, DB, TB); } else if constexpr (NW == 56) { WGMMA_RS_M64N56K16_T(D, A, DB, TB); } else if constexpr (NW == 64) { WGMMA_RS_M64N64K16_T(D, A, DB, TB); } else if constexpr (NW == 72) { WGMMA_RS_M64N72K16_T(D, A, DB, TB); } else if constexpr (NW == 80) { WGMMA_RS_M64N80K16_T(D, A, DB, TB); } else if constexpr (NW == 88) { WGMMA_RS_M64N88K16_T(D, A, DB, TB); } else if constexpr (NW == 96) { WGMMA_RS_M64N96K16_T(D, A, DB, TB); } else if constexpr (NW == 104) { WGMMA_RS_M64N104K16_T(D, A, DB, TB); } else if constexpr (NW == 112) { WGMMA_RS_M64N112K16_T(D, A, DB, TB); } else if constexpr (NW == 120) { WGMMA_RS_M64N120K16_T(D, A, DB, TB); } else if constexpr (NW == 128) { WGMMA_RS_M64N128K16_T(D, A, DB, TB); } else if constexpr (NW == 136) { WGMMA_RS_M64N136K16_T(D, A, DB, TB); } else if constexpr (NW == 144) { WGMMA_RS_M64N144K16_T(D, A, DB, TB); } else if constexpr (NW == 152) { WGMMA_RS_M64N152K16_T(D, A, DB, TB); } else if constexpr (NW == 160) { WGMMA_RS_M64N160K16_T(D, A, DB, TB); } else if constexpr (NW == 168) { WGMMA_RS_M64N168K16_T(D, A, DB, TB); } else if constexpr (NW == 176) { WGMMA_RS_M64N176K16_T(D, A, DB, TB); } else if constexpr (NW == 184) { WGMMA_RS_M64N184K16_T(D, A, DB, TB); } else if constexpr (NW == 192) { WGMMA_RS_M64N192K16_T(D, A, DB, TB); } else if constexpr (NW == 200) { WGMMA_RS_M64N200K16_T(D, A, DB, TB); } else if constexpr (NW == 208) { WGMMA_RS_M64N208K16_T(D, A, DB, TB); } else if constexpr (NW == 216) { WGMMA_RS_M64N216K16_T(D, A, DB, TB); } else if constexpr (NW == 224) { WGMMA_RS_M64N224K16_T(D, A, DB, TB); } else if constexpr (NW == 232) { WGMMA_RS_M64N232K16_T(D, A, DB, TB); } else if constexpr (NW == 240) { WGMMA_RS_M64N240K16_T(D, A, DB, TB); } else if constexpr (NW == 248) { WGMMA_RS_M64N248K16_T(D, A, DB, TB); } else if constexpr (NW == 256) { WGMMA_RS_M64N256K16_T(D, A, DB, TB); } } while (0) diff --git a/deploy/nemotron_native/kernels/codec/first_packet/codec.cu b/deploy/nemotron_native/kernels/codec/first_packet/codec.cu new file mode 100644 index 0000000..9e198ed --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/first_packet/codec.cu @@ -0,0 +1 @@ +#include "regions/stage1/codec.cu" diff --git a/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/codec.cu b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/codec.cu new file mode 100644 index 0000000..7cf7265 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/codec.cu @@ -0,0 +1,4 @@ +#define NATIVE_CODEC_PREFETCH 1 +#define NATIVE_CODEC_PACK_SUBCTAS 1 +#define D7W 96 +#include "native/codec/native_stage1.cu" diff --git a/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/kernel_kf_codec.h b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/kernel_kf_codec.h new file mode 100644 index 0000000..9a8b4d5 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/kernel_kf_codec.h @@ -0,0 +1,13 @@ +#pragma once +#include +struct CodecArgs { + const void *inp[41]; + const void *st[26]; + void *ns[26]; + void *waveform; + void *scratch; + int B; + int F; +}; +void codec_launch(const CodecArgs &a, void *stream); +size_t codec_scratch_bytes(int B, int F); diff --git a/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/legacy_codec.cuh b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/legacy_codec.cuh new file mode 100644 index 0000000..650a428 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/legacy_codec.cuh @@ -0,0 +1,1389 @@ +#include "kernel_kf_codec.h" +#include +#include +#include +#include +using namespace nvcuda; +template __device__ void prepare_attribute(F,int) {} +#ifndef TC2_QPT +#define TC2_QPT 4 +#endif +#define DIVUP(a,b) (((a) + (b) - 1) / (b)) +__device__ __forceinline__ float half_snake(float v, int c, int half, + const float *__restrict__ alpha) { + if (c < half) { + const float a = alpha[c]; + const float s = __sinf(a * v); + return __fmaf_rn(s, s / (a + 1e-9f), v); + } + return v > 0.f ? v : 0.01f * v; +} +#ifndef WG_ST3 +#define WG_ST3 1 +#define WG_MB3 3 +#define WG_ST7 1 +#define WG_MB7 3 +#define WG_BN3 128 +#define WG_BN7 128 +#endif +#define PDL_ENTER() do { (void)0; if (threadIdx.x == 0) (void)0; } while (0) +#define PDL_TRIGGER() do { if (threadIdx.x == 0) (void)0; } while (0) +template +__device__ __forceinline__ void pdl_launch(KFn kern, dim3 grid, dim3 blk, size_t shmem, + cudaStream_t s, KArgs... args) { + native_invoke(kern, grid, blk, shmem, args...); +} +template +__device__ __forceinline__ void pdl_launch_pf(KFn kern, dim3 grid, dim3 blk, size_t shmem, + cudaStream_t s, const void *pf, size_t pfb, KArgs... args) { + native_invoke_pf(kern, grid, blk, shmem, pf, pfb, args...); +} +#include "wg_kernel.inc" +#include "wg2_kernel.inc" +#include "wg3_kernel.inc" +#include "wg4_step_widths.inc" +#include "wg5_step_widths.inc" +#include "repack3.inc" +#include "tail_kernel.inc" +template +__device__ __noinline__ void dense_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + PDL_TRIGGER(); + constexpr int BK = 16 * KT; + constexpr int NT = NWARPS * 32; + constexpr int WTM = BM / WM; + constexpr int WTN = BN / WN; + constexpr int FM = WTM / 16, FN = WTN / 16; + constexpr int LDA = BK + 8, LDB = BN + 8; + constexpr int GEMM_SM = BM * LDA + BK * LDB; + constexpr int EPI_SM = NWARPS * WTM * WTN * 2; + constexpr int SM_H = GEMM_SM > EPI_SM ? GEMM_SM : EPI_SM; + __half *smem=reinterpret_cast<__half*>(native_arena()); + char *colok=reinterpret_cast(native_aux()); + __half *As = smem; + __half *Bs = smem + BM * LDA; + float *Cw = reinterpret_cast(smem); + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int wm = warp / WN, wn = warp % WN; + const int m0 = blockIdx.x * BM; + const int n0 = blockIdx.y * BN; + const int H = KT - 1; + const int Kg = Cin * KT; + if (tid < BN) { + const int n = n0 + tid; + colok[tid] = (n < N && (n % L) < T) ? 1 : 0; + } + wmma::fragment acc[FM][FN]; +#pragma unroll + for (int i = 0; i < FM; ++i) +#pragma unroll + for (int j = 0; j < FN; ++j) wmma::fill_fragment(acc[i][j], 0.f); + constexpr int ATOT = BM * (BK / 8); + constexpr int AREG = (ATOT + NT - 1) / NT; + constexpr int BTOT = 16 * (BN / 8); + constexpr int BREG = (BTOT + NT - 1) / NT; + int4 ar[AREG]; + int4 br0[BREG], br1[BREG]; +#define CODEC_LOADG(kg0) { const int _ci0 = (kg0) / KT; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * NT; ar[r] = make_int4(0, 0, 0, 0); if (ATOT % NT == 0 || idx < ATOT) { const int m = idx / (BK / 8), v = idx % (BK / 8); const int co = m0 + m, kg = (kg0) + v * 8; if (co < Cout && kg < Kg) ar[r] = *reinterpret_cast(w + (size_t)co * Kg + kg); } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * NT; br0[r] = make_int4(0, 0, 0, 0); br1[r] = make_int4(0, 0, 0, 0); if (BTOT % NT == 0 || idx < BTOT) { const int cil = idx / (BN / 8), nc = idx % (BN / 8); const int ci = _ci0 + cil; if (ci < Cin) { const int4 *src = reinterpret_cast( in + (size_t)ci * SCin + n0 + nc * 8); br0[r] = src[0]; br1[r] = src[1]; } } } } + CODEC_LOADG(0) + for (int kg0 = 0; kg0 < Kg; kg0 += BK) { + __syncthreads(); +#pragma unroll + for (int r = 0; r < AREG; ++r) { + const int idx = tid + r * NT; + if (ATOT % NT == 0 || idx < ATOT) + *reinterpret_cast(As + (idx / (BK / 8)) * LDA + + (idx % (BK / 8)) * 8) = ar[r]; + } +#pragma unroll + for (int r = 0; r < BREG; ++r) { + const int idx = tid + r * NT; + if (BTOT % NT != 0 && idx >= BTOT) continue; + const int cil = idx / (BN / 8), nc = idx % (BN / 8); + unsigned int u[8]; + u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; + u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; +#pragma unroll + for (int k = 0; k < KT; ++k) { + const int kh = k >> 1; + int4 o; + if (k & 1) { + o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); + o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); + o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); + o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); + } else { + o.x = u[kh + 0]; o.y = u[kh + 1]; + o.z = u[kh + 2]; o.w = u[kh + 3]; + } + *reinterpret_cast(Bs + (cil * KT + k) * LDB + nc * 8) = o; + } + } + __syncthreads(); + if (kg0 + BK < Kg) CODEC_LOADG(kg0 + BK) +#pragma unroll + for (int kf = 0; kf < BK; kf += 16) { + wmma::fragment af[FM]; + wmma::fragment bf[FN]; +#pragma unroll + for (int i = 0; i < FM; ++i) + wmma::load_matrix_sync(af[i], As + (wm * WTM + i * 16) * LDA + kf, LDA); +#pragma unroll + for (int j = 0; j < FN; ++j) + wmma::load_matrix_sync(bf[j], Bs + kf * LDB + wn * WTN + j * 16, LDB); +#pragma unroll + for (int i = 0; i < FM; ++i) +#pragma unroll + for (int j = 0; j < FN; ++j) + wmma::mma_sync(acc[i][j], af[i], bf[j], acc[i][j]); + } + } +#undef CODEC_LOADG + __syncthreads(); + float *Cme = Cw + warp * WTM * WTN; +#pragma unroll + for (int i = 0; i < FM; ++i) +#pragma unroll + for (int j = 0; j < FN; ++j) + wmma::store_matrix_sync(Cme + (i * 16) * WTN + j * 16, acc[i][j], WTN, + wmma::mem_row_major); + __syncwarp(); + const int half = Cout >> 1; + const int lane = tid & 31; + for (int e = lane; e < WTM * WTN; e += 32) { + const int mm = e / WTN, nn = e % WTN; + const int co = m0 + wm * WTM + mm; + const int nl = wn * WTN + nn; + if (co >= Cout || !colok[nl]) continue; + const int n = n0 + nl; + float v = Cme[e] + bias[co]; + if (resid != nullptr) + v += __half2float(resid[(size_t)co * SCout + n + H]); + if (do_snake) v = half_snake(v, co, half, alpha); + out[(size_t)co * SCout + n + H] = __float2half(v); + } + if (blockIdx.y == 0) { + const int nbx = gridDim.x; + if (ns != nullptr) { + const int B = N / L; + const int total = Cin * B * H; + for (int idx = blockIdx.x * NT + tid; idx < total; idx += nbx * NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int B = N / L; + const int total = Cout * B * pad_len; + for (int idx = blockIdx.x * NT + tid; idx < total; idx += nbx * NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} +template +__device__ __noinline__ void transpose_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + __half *__restrict__ out, const __half *__restrict__ pad_src, + __half *__restrict__ ns, int Cin, int Cout, int Tin, + int Lin, int SCin, int jbase, int Lout, int SCout, int Hout, int pad_off, + int pad_len, int B) { + if constexpr (PDL) PDL_ENTER(); + constexpr int Ktaps = 2 * STRIDE; + constexpr int QPT = QPT_; + constexpr int VPT = QPT * STRIDE; + static_assert(QPT % 2 == 0, "half2 input loads need an even frame count"); + static_assert(VPT % 2 == 0, "half2 staging needs an even slot"); + float *ws=reinterpret_cast(native_aux()); + __half *sm=reinterpret_cast<__half*>(native_arena()); + const int co = blockIdx.y; + const int b = blockIdx.z; + const int tid = threadIdx.x, BLK = blockDim.x; + if (G != 4 || STRIDE != 7 || native_meta()->lengths[b] > 0) { + for (int i = tid; i < G * Ktaps; i += BLK) + ws[i] = __half2float(w[(size_t)(co * G + i / Ktaps) * Ktaps + (i % Ktaps)]); + __syncthreads(); + const float bs = bias[0 + co]; + const int half = Cout >> 1; + const bool snake = co < half; + const float av = snake ? alpha[co] : 0.f; + const float ri = snake ? __frcp_rn(av + 1e-9f) : 0.f; + const __half *ic = in + (size_t)(co * G) * SCin + b * Lin + jbase; + __half *oc = out + (size_t)co * SCout + b * Lout + Hout; + constexpr bool VEC16 = (VPT % 8) == 0; + const int ophase = + VEC16 ? 0 : (int)((((uintptr_t)oc) >> 1) & 7u); + __half *const sbase = sm + ophase; + const int hd = (8 - ophase) & 7; + const int lane = tid & 31; + for (int qb = 1 + blockIdx.x * BLK * QPT; qb <= Tin; + qb += gridDim.x * BLK * QPT) { + const int q = qb + QPT * tid; + const bool a0 = (q <= Tin); + float jq[QPT][G], jp0[G]; + #pragma unroll + for (int g = 0; g < G; ++g) { + #pragma unroll + for (int e2 = 0; e2 < QPT / 2; ++e2) { + const __half2 v = a0 ? *reinterpret_cast( + ic + (size_t)g * SCin + q + 2 * e2) + : __floats2half2_rn(0.f, 0.f); + jq[2 * e2][g] = __half2float(v.x); + jq[2 * e2 + 1][g] = __half2float(v.y); + } + } + #pragma unroll + for (int g = 0; g < G; ++g) { + const float up = __shfl_up_sync(0xffffffffu, jq[QPT - 1][g], 1); + jp0[g] = (lane == 0) + ? (a0 ? __half2float(ic[(size_t)g * SCin + q - 1]) : 0.f) + : up; + } + __half *const sp = sbase + tid * VPT; + __align__(16) __half vreg[VPT]; + #pragma unroll + for (int r = 0; r < STRIDE; ++r) { + float wc[G], wp[G]; + #pragma unroll + for (int g = 0; g < G; ++g) { + wc[g] = ws[g * Ktaps + r]; + wp[g] = ws[g * Ktaps + r + STRIDE]; + } + #pragma unroll + for (int e = 0; e < QPT; ++e) { + float v = bs; + #pragma unroll + for (int g = 0; g < G; ++g) { + const float cur = jq[e][g]; + const float prv = (e == 0) ? jp0[g] : jq[e - 1][g]; + v = __fmaf_rn(wc[g], cur, __fmaf_rn(wp[g], prv, v)); + } + if (snake) { + const float sn = __sinf(av * v); + v = __fmaf_rn(sn * sn, ri, v); + } else { + v = v > 0.f ? v : 0.01f * v; + } + vreg[e * STRIDE + r] = __float2half(v); + } + } + if constexpr (VEC16) { + #pragma unroll + for (int m = 0; m < VPT / 8; ++m) + *reinterpret_cast(sp + m * 8) = + *reinterpret_cast(vreg + m * 8); + } else { + #pragma unroll + for (int i = 0; i < VPT; ++i) sp[i] = vreg[i]; + } + __syncthreads(); + const int nval = min(BLK * QPT, Tin - qb + 1); + const int tot = nval * STRIDE; + __half *dst = oc + (size_t)(qb - 1) * STRIDE; + if constexpr (VEC16) { + const int tot2 = tot >> 1; + for (int i = tid; i < tot2; i += BLK) + reinterpret_cast<__half2 *>(dst)[i] = + reinterpret_cast(sbase)[i]; + if ((tot & 1) && tid == 0) dst[tot - 1] = sbase[tot - 1]; + } else { + const int head = hd < tot ? hd : tot; + for (int i = tid; i < head; i += BLK) dst[i] = sbase[i]; + const int nb4 = (tot - head) >> 3; + const int4 *s4 = reinterpret_cast(sbase + head); + int4 *d4 = reinterpret_cast(dst + head); + for (int i = tid; i < nb4; i += BLK) d4[i] = s4[i]; + for (int i = head + (nb4 << 3) + tid; i < tot; i += BLK) + dst[i] = sbase[i]; + } + __syncthreads(); + } + } + if (blockIdx.x == 0) { + if (ns != nullptr) { + for (int ci = co + tid * Cout; ci < Cin; ci += BLK * Cout) + ns[(size_t)b * Cin + ci] = + in[(size_t)ci * SCin + b * Lin + jbase + native_true_tokens(b,Tin)]; + } + if (pad_src != nullptr) { + for (int j = tid; j < pad_len; j += BLK) + out[(size_t)co * SCout + b * Lout + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } +} +template +__device__ __noinline__ void transpose_conv_early_bpack_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + __half *__restrict__ out, const __half *__restrict__ pad_src, + __half *__restrict__ ns, int Cin, int Cout, int Tin, + int Lin, int SCin, int jbase, int Lout, int SCout, int Hout, int pad_off, + int pad_len, int B) { + if constexpr (PDL) PDL_ENTER(); + static_assert(BPACK <= 32); + static_assert(LPB == 8 || LPB == 16 || LPB == 32); + constexpr int Ktaps = 2 * STRIDE; + constexpr int WARP = 32; + constexpr int NT = WARP * BPACK; + constexpr int BPB = WARP / LPB; + float *ws=reinterpret_cast(native_aux()); + __half *sm=reinterpret_cast<__half*>(native_arena()); + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int lane = tid & 31; + const int sub = lane % LPB; + const int bsub = lane / LPB; + const int co = blockIdx.y; + const int b = blockIdx.z * (BPACK * BPB) + warp * BPB + bsub; + const bool active = b < B; + for (int i = tid; i < G * Ktaps; i += NT) + ws[i] = __half2float(w[(size_t)(co * G + i / Ktaps) * Ktaps + + (i % Ktaps)]); + __syncthreads(); + const int qb = 1 + blockIdx.x * LPB; + const int q = qb + sub; + __half *sp = sm + (warp * WARP + lane) * STRIDE; + const bool live = active && q <= Tin; + { + const float bs = bias[co]; + const int half = Cout >> 1; + const bool snake = co < half; + const float av = snake ? alpha[co] : 0.f; + const float ri = snake ? __frcp_rn(av + 1e-9f) : 0.f; + const __half *ic = in + (size_t)(co * G) * SCin + b * Lin + jbase; + float jq[G], jp[G]; +#pragma unroll + for (int g = 0; g < G; ++g) + jq[g] = live ? __half2float(ic[(size_t)g * SCin + q]) : 0.f; +#pragma unroll + for (int g = 0; g < G; ++g) { + const float up = __shfl_up_sync(0xffffffffu, jq[g], 1, LPB); + jp[g] = (sub == 0) + ? (live ? __half2float(ic[(size_t)g * SCin + q - 1]) : 0.f) + : up; + } + if (live) +#pragma unroll + for (int r = 0; r < STRIDE; ++r) { + float v = bs; +#pragma unroll + for (int g = 0; g < G; ++g) + v = __fmaf_rn(ws[g * Ktaps + r], jq[g], + __fmaf_rn(ws[g * Ktaps + r + STRIDE], jp[g], v)); + if (snake) { + const float sn = __sinf(av * v); + v = __fmaf_rn(sn * sn, ri, v); + } else { + v = v > 0.f ? v : 0.01f * v; + } + sp[r] = __float2half(v); + } + } + __syncwarp(); + if (active) { + const int nval = min(LPB, Tin - (qb - 1)); + const int tot2 = (nval * STRIDE) >> 1; + const __half *src = sm + (warp * WARP + bsub * LPB) * STRIDE; + __half *dst = out + (size_t)co * SCout + b * Lout + Hout + + (size_t)(qb - 1) * STRIDE; + for (int i = sub; i < tot2; i += LPB) + reinterpret_cast<__half2 *>(dst)[i] = + reinterpret_cast(src)[i]; + if ((nval * STRIDE) & 1) + if (sub == 0) dst[nval * STRIDE - 1] = src[nval * STRIDE - 1]; + } + if (blockIdx.x == 0 && active) { + if (ns != nullptr) { + for (int ci = co + sub * Cout; ci < Cin; ci += LPB * Cout) + ns[(size_t)b * Cin + ci] = + in[(size_t)ci * SCin + b * Lin + jbase + native_true_tokens(b,Tin)]; + } + if (pad_src != nullptr) { + for (int j = sub; j < pad_len; j += LPB) + out[(size_t)co * SCout + b * Lout + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } +} +template +__device__ __noinline__ void transpose_conv_bpack2_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + __half *__restrict__ out, const __half *__restrict__ pad_src, + __half *__restrict__ ns, int Cin, int Cout, int Tin, + int Lin, int SCin, int jbase, int Lout, int SCout, int Hout, int pad_off, + int pad_len, int B) { + if constexpr (PDL) PDL_ENTER(); + static_assert(BPACK <= 32); + static_assert(LPB == 8 || LPB == 16 || LPB == 32); + constexpr int Ktaps = 2 * STRIDE; + constexpr int WARP = 32; + constexpr int NT = WARP * BPACK; + constexpr int BPB = WARP / LPB; + constexpr int QPT = 2; + constexpr int VPT = QPT * STRIDE; + float *ws=reinterpret_cast(native_aux()); + __half *sm=reinterpret_cast<__half*>(native_arena()); + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int lane = tid & 31; + const int sub = lane % LPB; + const int bsub = lane / LPB; + const int co = blockIdx.y; + const int b = blockIdx.z * (BPACK * BPB) + warp * BPB + bsub; + const bool active = b < B; + for (int i = tid; i < G * Ktaps; i += NT) + ws[i] = __half2float(w[(size_t)(co * G + i / Ktaps) * Ktaps + + (i % Ktaps)]); + __syncthreads(); + const int qb = 1 + blockIdx.x * (LPB * QPT); + const int q = qb + sub * QPT; + __half *sp = sm + (warp * WARP + lane) * VPT; + const bool a0 = active && (q <= Tin); + { + const float bs = bias[co]; + const int half = Cout >> 1; + const bool snake = co < half; + const float av = snake ? alpha[co] : 0.f; + const float ri = snake ? __frcp_rn(av + 1e-9f) : 0.f; + const __half *ic = in + (size_t)(co * G) * SCin + b * Lin + jbase; + float jq0[G], jq1[G], jp0[G]; +#pragma unroll + for (int g = 0; g < G; ++g) { + const __half2 v = a0 ? *reinterpret_cast( + ic + (size_t)g * SCin + q) + : __floats2half2_rn(0.f, 0.f); + jq0[g] = __half2float(v.x); + jq1[g] = __half2float(v.y); + } +#pragma unroll + for (int g = 0; g < G; ++g) { + const float up = __shfl_up_sync(0xffffffffu, jq1[g], 1, LPB); + jp0[g] = (sub == 0) + ? (a0 ? __half2float(ic[(size_t)g * SCin + q - 1]) + : 0.f) + : up; + } + if (a0) { + float vv[VPT]; +#pragma unroll + for (int r = 0; r < STRIDE; ++r) { + float wc[G], wp[G]; +#pragma unroll + for (int g = 0; g < G; ++g) { + wc[g] = ws[g * Ktaps + r]; + wp[g] = ws[g * Ktaps + r + STRIDE]; + } +#pragma unroll + for (int e = 0; e < QPT; ++e) { + float v = bs; +#pragma unroll + for (int g = 0; g < G; ++g) { + const float cur = (e == 0) ? jq0[g] : jq1[g]; + const float prv = (e == 0) ? jp0[g] : jq0[g]; + v = __fmaf_rn(wc[g], cur, __fmaf_rn(wp[g], prv, v)); + } + if (snake) { + const float sn = __sinf(av * v); + v = __fmaf_rn(sn * sn, ri, v); + } else { + v = v > 0.f ? v : 0.01f * v; + } + vv[e * STRIDE + r] = v; + } + } + __half2 *sp2 = reinterpret_cast<__half2 *>(sp); +#pragma unroll + for (int i = 0; i < VPT / 2; ++i) + sp2[i] = __floats2half2_rn(vv[2 * i], vv[2 * i + 1]); + } + } + __syncwarp(); + if (active) { + const int nval = min(LPB * QPT, Tin - (qb - 1)); + const int tot = nval * STRIDE; + const __half *src = sm + (warp * WARP + bsub * LPB) * VPT; + __half *dst = out + (size_t)co * SCout + b * Lout + Hout + + (size_t)(qb - 1) * STRIDE; + for (int i = sub; i < (tot >> 1); i += LPB) + reinterpret_cast<__half2 *>(dst)[i] = + reinterpret_cast(src)[i]; + if (tot & 1) + if (sub == 0) dst[tot - 1] = src[tot - 1]; + } + if (blockIdx.x == 0 && active) { + if (ns != nullptr) { + for (int ci = co + sub * Cout; ci < Cin; ci += LPB * Cout) + ns[(size_t)b * Cin + ci] = + in[(size_t)ci * SCin + b * Lin + jbase + native_true_tokens(b,Tin)]; + } + if (pad_src != nullptr) { + for (int j = sub; j < pad_len; j += LPB) + out[(size_t)co * SCout + b * Lout + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } +} +#define TCONV_BPACK(G,STRIDE,BPACK,LPB,TIN,COUT,BB,...) pdl_launch(transpose_conv_early_bpack_kernel, dim3(DIVUP(TIN, LPB), COUT, DIVUP(BB, BPACK * (32 / LPB))), dim3(32 * BPACK), (size_t)(32 * (BPACK)) * (STRIDE) * 2 + 64, s, __VA_ARGS__) +#define TCONV_BP2(G,STRIDE,LPB,TIN,COUT,BB,...) pdl_launch(transpose_conv_bpack2_kernel, dim3(DIVUP(TIN, (LPB) * 2), COUT, DIVUP(BB, 32)), dim3(32 * (LPB)), (size_t)(32 * (LPB)) * (4 * (STRIDE)) + 64, s, __VA_ARGS__) +#define TCONV_BPACK_DISPATCH1(G,STRIDE,BPACK,TIN,COUT,BB,...) do { if ((TIN) >= 32) TCONV_BP2(G, STRIDE, 16, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) >= 16) TCONV_BP2(G, STRIDE, 8, TIN, COUT, BB, __VA_ARGS__); else TCONV_BPACK(G, STRIDE, BPACK, 16, TIN, COUT, BB, __VA_ARGS__); } while (0) +#define TCONV_BPACK_DISPATCH(G,STRIDE,BPACK,TIN,COUT,BB,...) do { if ((TIN) >= 32) TCONV_BP2(G, STRIDE, 16, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) >= 16) TCONV_BP2(G, STRIDE, 8, TIN, COUT, BB, __VA_ARGS__); else TCONV_BPACK(G, STRIDE, BPACK, 16, TIN, COUT, BB, __VA_ARGS__); } while (0) +#define TCONV_BPACK_DISPATCH8(G,STRIDE,BPACK,TIN,COUT,BB,...) do { if ((TIN) >= 32) TCONV_BP2(G, STRIDE, 16, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) >= 16) TCONV_BP2(G, STRIDE, 8, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) <= 8) TCONV_BPACK(G, STRIDE, BPACK, 8, TIN, COUT, BB, __VA_ARGS__); else TCONV_BPACK(G, STRIDE, BPACK, 16, TIN, COUT, BB, __VA_ARGS__); } while (0) +__device__ __noinline__ void pack_latent_kernel(const __half *__restrict__ latent, + const __half *__restrict__ state, + __half *__restrict__ out, int Cin, int T, + int L, int SC, int H, int B) { + const int total = Cin * B * L; + for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; + idx += gridDim.x * blockDim.x) { + const int j = idx % L, r = idx / L; + const int b = r % B, ci = r / B; + out[(size_t)ci * SC + b * L + j] = + (j < H) ? state[((size_t)b * Cin + ci) * H + j] + : latent[((size_t)b * Cin + ci) * T + (j - H)]; + } +} +#define PC_VEC 8 +__device__ __noinline__ void post_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, __half *__restrict__ y, + __half *__restrict__ ns, int Cin, int T, int L, int SC, int KT, int B, + int N) { + float *pw=reinterpret_cast(native_arena()); + for (int i = threadIdx.x; i < Cin * KT; i += 256) pw[i] = __half2float(w[i]); + __syncthreads(); + const int H = KT - 1; + const float bs = bias[0]; + for (int n = (blockIdx.x * 256 + threadIdx.x) * PC_VEC; n < N; + n += gridDim.x * 256 * PC_VEC) { + float v[PC_VEC]; +#pragma unroll + for (int i = 0; i < PC_VEC; ++i) v[i] = bs; + for (int ci = 0; ci < Cin; ++ci) { + const __half *xr = in + (size_t)ci * SC + n; + const float *wr = pw + ci * KT; + const int4 lo = *reinterpret_cast(xr); + const int4 hi = *reinterpret_cast(xr + 8); + unsigned u[8]; + u[0] = lo.x; u[1] = lo.y; u[2] = lo.z; u[3] = lo.w; + u[4] = hi.x; u[5] = hi.y; u[6] = hi.z; u[7] = hi.w; + float x[PC_VEC + 6]; +#pragma unroll + for (int j = 0; j < (PC_VEC + 6) / 2; ++j) { + const float2 f = __half22float2( + *reinterpret_cast(&u[j])); + x[2 * j] = f.x; + x[2 * j + 1] = f.y; + } +#pragma unroll + for (int i = 0; i < PC_VEC; ++i) + for (int k = 0; k < 7; ++k) + if (k < KT) v[i] = __fmaf_rn(wr[k], x[i + k], v[i]); + } + const int b0 = n / L, t0 = n - (n / L) * L; +#pragma unroll + for (int i = 0; i < PC_VEC; ++i) { + int t = t0 + i, b = b0; + if (t >= L) { t -= L; b = b0 + 1; } + if (t < T && b < B) + y[(size_t)b * T + t] = + __float2half(fminf(fmaxf(v[i], -1.f), 1.f)); + } + } + if (ns != nullptr) { + const int total = B * Cin * H; + for (int idx = blockIdx.x * 256 + threadIdx.x; idx < total; + idx += gridDim.x * 256) { + const int j = idx % H, r = idx / H; + const int ci = r % Cin, b = r / Cin; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SC + b * L + native_true_tokens(b,T) + j]; + } + } +} +#define RP_JOBS 6 +struct RepackSet { + const __half *src[RP_JOBS]; + __half *dst[RP_JOBS]; + int Cout[RP_JOBS], CoutP[RP_JOBS], Cin[RP_JOBS]; + int base[RP_JOBS], nch[RP_JOBS], rs[RP_JOBS]; + int k3base; + const __half *k3src[2]; + __half *k3dst[2]; +}; +__host__ __device__ inline int native_RP_COUT(int j) {const int v[6]={768, 384, 256, 128, 64, 32};return v[j];} +__host__ __device__ inline int native_RP_CIN(int j) {const int v[6]={384, 768, 128, 256, 32, 64};return v[j];} +__host__ __device__ inline int native_RP_COUTP(int j) {const int v[6]={768, 384, 256, 128, 128, 128};return v[j];} +template +__device__ __noinline__ void repack_kernel(RepackSet js) { + const int bx = blockIdx.x; + int job, job_bx; + if (bx < 1152) { + job = 0; job_bx = 0; + } else if (bx < 2304) { + job = 1; job_bx = 1152; + } else if (bx < 2432) { + job = 2; job_bx = 2304; + } else if (bx < 2560) { + job = 3; job_bx = 2432; + } else if (bx < 2576) { + job = 4; job_bx = 2560; + } else { + job = 5; job_bx = 2576; + } + const int Cout = js.Cout[job], CoutP = js.CoutP[job], Cin = js.Cin[job]; + const int idx = (bx - job_bx) * blockDim.x + threadIdx.x; + const int co = idx / Cin, ci = idx - co * Cin; + const int c = ci / CIB, cl = ci - c * CIB; + __half *d = js.dst[job] + (size_t)c * (KT * CoutP * CIB) + + (size_t)(co >> 3) * (CIB * 8) + (cl >> 3) * 64 + (co & 7) * 8 + + (cl & 7); + const __half *sp = js.src[job] + (size_t)(co * Cin + ci) * KT; + const bool live = co < Cout; +#pragma unroll + for (int k = 0; k < KT; ++k) + d[(size_t)k * CoutP * CIB] = live ? sp[k] : __float2half(0.f); +} +template +__device__ __noinline__ void repack_rs_kernel( + const RepackSet &js, const __half *__restrict__ latent, const __half *__restrict__ latent_state, + __half *__restrict__ latent_out, int latent_C, int latent_T, int latent_L, + int latent_SC, int latent_H, int latent_B) { + PDL_TRIGGER(); + constexpr int CO_TILE = 16; + constexpr int CI_TILE = 64; + constexpr int ROW_H = CI_TILE * KT + 8; + constexpr int ROW4 = CI_TILE * KT / 8; + __half *tile=reinterpret_cast<__half*>(native_arena()); + const int bx = blockIdx.x; + { + const int idx = bx + threadIdx.x * gridDim.x; + const int total = latent_C * latent_B * latent_L; + if (idx < total) { + const int j = idx % latent_L, r = idx / latent_L; + const int b = r % latent_B, ci = r / latent_B; + latent_out[(size_t)ci * latent_SC + b * latent_L + j] = + (j < latent_H) + ? latent_state[((size_t)b * latent_C + ci) * latent_H + j] + : latent[((size_t)b * latent_C + ci) * latent_T + + (j - latent_H)]; + } + } + if (bx >= js.k3base) { + constexpr int K3 = 3; + constexpr int R3_H = CI_TILE * K3 + 8; + constexpr int R3_4 = CI_TILE * K3 / 8; + const int r = bx - js.k3base; + const int conv = r / 1152; + const int loc = r - conv * 1152; + const int Cout3 = conv ? 768 : 1536; + const int Cin3 = conv ? 1536 : 768; + const int nci3 = Cin3 / CI_TILE; + const int cot3 = loc / nci3; + if (cot3 * CO_TILE >= Cout3) return; + const int co03 = cot3 * CO_TILE; + const int ci03 = (loc - cot3 * nci3) * CI_TILE; + const int NCH3 = Cin3 >> 4; + const __half *const s3 = js.k3src[conv]; + __half *const d3 = js.k3dst[conv]; + const int t3 = threadIdx.x; +#pragma unroll + for (int n = 0; n < CO_TILE * R3_4 / 128; ++n) { + const int i = t3 + n * 128; + const int co = i / R3_4, q = i - co * R3_4; + reinterpret_cast(tile + co * R3_H)[q] = + *reinterpret_cast( + s3 + ((size_t)(co03 + co) * Cin3 + ci03) * K3 + q * 8); + } + __syncthreads(); + const int w3 = t3 >> 5, l3 = t3 & 31; + const int rA3 = l3 >> 2, q43 = l3 & 3; + const int cb3 = co03 >> 4; + __half *const db3 = d3 + ((size_t)(cb3 >> 2) * NCH3 * K3) * 1024 + + ((cb3 & 3) * 32 + l3) * 8; +#pragma unroll + for (int u = w3; u < (CI_TILE / 16) * K3; u += 4) { + const int cl = u / K3, tap = u - cl * K3; + const int cgi = (ci03 >> 4) + cl; + const __half *sa = tile + rA3 * R3_H + (cl * 16 + 2 * q43) * K3 + tap; + const __half *sb = sa + 8 * R3_H; + uint32_t o[4]; + o[0] = __half_as_ushort(sa[0]) | + ((uint32_t)__half_as_ushort(sa[K3]) << 16); + o[1] = __half_as_ushort(sb[0]) | + ((uint32_t)__half_as_ushort(sb[K3]) << 16); + o[2] = __half_as_ushort(sa[8 * K3]) | + ((uint32_t)__half_as_ushort(sa[9 * K3]) << 16); + o[3] = __half_as_ushort(sb[8 * K3]) | + ((uint32_t)__half_as_ushort(sb[9 * K3]) << 16); + *reinterpret_cast(db3 + (size_t)(cgi * K3 + tap) * 1024) = + make_int4(o[0], o[1], o[2], o[3]); + } + return; + } + int job = 0; +#pragma unroll + for (int q = 1; q < RP_JOBS; ++q) if (bx >= js.base[q]) job = q; + const int local = bx - js.base[job]; + const int Cout = js.Cout[job], Cin = js.Cin[job]; + const int nci = (Cin + CI_TILE - 1) / CI_TILE; + const int cot = local / nci; + const int co0 = cot * CO_TILE; + const int ci0 = (local - cot * nci) * CI_TILE; + const int NCH = js.nch[job]; + const int t = threadIdx.x; + const int valid4 = (min(Cin - ci0, CI_TILE) * KT) >> 3; + for (int i = t; i < CO_TILE * ROW4; i += 128) { + const int co = i / ROW4, q = i - co * ROW4; + int4 v = make_int4(0, 0, 0, 0); + if (co0 + co < Cout && q < valid4) + v = *reinterpret_cast( + js.src[job] + ((size_t)(co0 + co) * Cin + ci0) * KT + q * 8); + reinterpret_cast(tile + co * ROW_H)[q] = v; + } + __syncthreads(); + const int warp = t >> 5, lane = t & 31; + const int rA = lane >> 2, q4 = lane & 3; + const int cb = co0 >> 4; + const int mtile = cb >> 2, wsub = cb & 3; + __half *const dbase = js.dst[job] + + ((size_t)mtile * NCH * KT) * 1024 + (wsub * 32 + lane) * 8; + if (!js.rs[job]) { + const int CoutP = js.CoutP[job]; + constexpr int UN2 = 16 * (CI_TILE / 8) * KT; +#pragma unroll + for (int n = 0; n < UN2 / 128; ++n) { + const int u = t + n * 128; + const int col = u & 15, cig = (u >> 4) & (CI_TILE / 8 - 1); + const int tap = u / (16 * (CI_TILE / 8)); + const int gci = ci0 + cig * 8; + if (gci >= Cin) continue; + const int gco = co0 + col; + uint32_t o[4]; + const __half *sp = tile + col * ROW_H; +#pragma unroll + for (int e = 0; e < 4; ++e) + o[e] = __half_as_ushort(sp[(cig * 8 + 2 * e) * KT + tap]) | + ((uint32_t)__half_as_ushort( + sp[(cig * 8 + 2 * e + 1) * KT + tap]) << 16); + __half *d = js.dst[job] + + (size_t)(gci >> 4) * ((size_t)KT * CoutP * 16) + + (size_t)tap * (CoutP * 16) + (size_t)(gco >> 3) * 128 + + ((gci >> 3) & 1) * 64 + (gco & 7) * 8; + *reinterpret_cast(d) = make_int4(o[0], o[1], o[2], o[3]); + } + return; + } + constexpr int UNITS = (CI_TILE / 16) * KT; +#pragma unroll + for (int u = warp; u < UNITS; u += 4) { + const int cl = u / KT, tap = u - cl * KT; + const int cgi = (ci0 >> 4) + cl; + if (cgi >= NCH) continue; + const __half *sa = tile + rA * ROW_H + (cl * 16 + 2 * q4) * KT + tap; + const __half *sb = sa + 8 * ROW_H; + uint32_t o[4]; + o[0] = __half_as_ushort(sa[0]) | + ((uint32_t)__half_as_ushort(sa[KT]) << 16); + o[1] = __half_as_ushort(sb[0]) | + ((uint32_t)__half_as_ushort(sb[KT]) << 16); + o[2] = __half_as_ushort(sa[8 * KT]) | + ((uint32_t)__half_as_ushort(sa[9 * KT]) << 16); + o[3] = __half_as_ushort(sb[8 * KT]) | + ((uint32_t)__half_as_ushort(sb[9 * KT]) << 16); + *reinterpret_cast(dbase + (size_t)(cgi * KT + tap) * 1024) = + make_int4(o[0], o[1], o[2], o[3]); + } +} +__device__ __noinline__ void repack3_rs_kernel( + const __half *__restrict__ win, const __half *__restrict__ wsk, + __half *__restrict__ din, __half *__restrict__ dsk, + const __half *__restrict__ asrc, __half *__restrict__ adst, int SCA, + int ncol, int ntile) { + PDL_TRIGGER(); + if (blockIdx.y == 12) { + if (adst == nullptr) return; + for (int u = blockIdx.x; u < 96 * ntile; u += gridDim.x) { + const int g = u / ntile; + const int col = (u - g * ntile) * 128 + (int)threadIdx.x; + if (col >= ncol) continue; + __half v[8]; +#pragma unroll + for (int h = 0; h < 8; ++h) + v[h] = asrc[(size_t)(g * 8 + h) * SCA + col]; + *reinterpret_cast(adst + (size_t)g * ((size_t)SCA * 8) + + (size_t)col * 8) = + *reinterpret_cast(v); + } + return; + } + constexpr int KT = 3; + constexpr int CO_TILE = 16; + constexpr int CI_TILE = 64; + constexpr int ROW_H = CI_TILE * KT + 8; + constexpr int ROW4 = CI_TILE * KT / 8; + __half *tile=reinterpret_cast<__half*>(native_arena()); + const int jy = blockIdx.y; + const int conv = jy & 1, b = 5 - (jy >> 1); + const int Cout = conv ? 768 : 1536; + const int Cin = conv ? 1536 : 768; + const __half *src = (conv ? wsk : win) + (size_t)(b + 1) * Cout * Cin * KT; + __half *dst = (conv ? dsk : din) + (size_t)b * ((size_t)Cout * Cin * KT); + const int nci = Cin / CI_TILE; + const int cot = blockIdx.x / nci; + const int co0 = cot * CO_TILE; + const int ci0 = (blockIdx.x - cot * nci) * CI_TILE; + const int NCH = Cin >> 4; + const int t = threadIdx.x; +#pragma unroll + for (int n = 0; n < CO_TILE * ROW4 / 128; ++n) { + const int i = t + n * 128; + const int co = i / ROW4, q = i - co * ROW4; + reinterpret_cast(tile + co * ROW_H)[q] = + *reinterpret_cast( + src + ((size_t)(co0 + co) * Cin + ci0) * KT + q * 8); + } + __syncthreads(); + const int warp = t >> 5, lane = t & 31; + const int rA = lane >> 2, q4 = lane & 3; + const int cb = co0 >> 4; + __half *const dbase = dst + ((size_t)(cb >> 2) * NCH * KT) * 1024 + + ((cb & 3) * 32 + lane) * 8; +#pragma unroll + for (int u = warp; u < (CI_TILE / 16) * KT; u += 4) { + const int cl = u / KT, tap = u - cl * KT; + const int cgi = (ci0 >> 4) + cl; + const __half *sa = tile + rA * ROW_H + (cl * 16 + 2 * q4) * KT + tap; + const __half *sb = sa + 8 * ROW_H; + uint32_t o[4]; + o[0] = __half_as_ushort(sa[0]) | + ((uint32_t)__half_as_ushort(sa[KT]) << 16); + o[1] = __half_as_ushort(sb[0]) | + ((uint32_t)__half_as_ushort(sb[KT]) << 16); + o[2] = __half_as_ushort(sa[8 * KT]) | + ((uint32_t)__half_as_ushort(sa[9 * KT]) << 16); + o[3] = __half_as_ushort(sb[8 * KT]) | + ((uint32_t)__half_as_ushort(sb[9 * KT]) << 16); + *reinterpret_cast(dbase + (size_t)(cgi * KT + tap) * 1024) = + make_int4(o[0], o[1], o[2], o[3]); + } +} +__device__ __noinline__ void repack_activation_only(const __half *asrc,__half *adst,int SCA,int ncol,int ntile){ + if (adst == nullptr) return; + for (int u = blockIdx.x; u < 96 * ntile; u += gridDim.x) { + const int g = u / ntile; + const int col = (u - g * ntile) * 128 + (int)threadIdx.x; + if (col >= ncol) continue; + __half v[8]; +#pragma unroll + for (int h = 0; h < 8; ++h) + v[h] = asrc[(size_t)(g * 8 + h) * SCA + col]; + *reinterpret_cast(adst + (size_t)g * ((size_t)SCA * 8) + + (size_t)col * 8) = + *reinterpret_cast(v); + } + return; +} +namespace { +struct Stack { + int T, L, N, SC, H; +}; +inline __host__ __device__ int tc_blk(int Tin, int qpt = 2) { + int best = 256, bw = 1 << 30; + for (int b = 256; b >= 32; b >>= 1) { + const int per = qpt * b; + const int waste = (Tin + per - 1) / per * per - Tin; + if (waste < bw) { bw = waste; best = b; } + } + return best; +} +inline __host__ __device__ int roundup(int a, int b) { return (a + b - 1) / b * b; } +__host__ __device__ Stack mk(int B, int T, int K) { + Stack s; + s.T = T; + s.H = K - 1; + s.L = s.H + T; + s.N = B * s.L; + s.SC = roundup(s.N, 128) + 128; + return s; +} +struct Plan { + Stack s0, s1, s2, s3, s4; + __half *LB, *X0, *M0, *Y0; + __half *A, *Ac, *Bb, *M1; + __half *P, *Q, *M2; + __half *R, *S, *M3; + __half *U; + __half *W7[RP_JOBS]; + __half *W3i, *W3s; + size_t total; +}; +static constexpr int kTuneB = 32; +inline __host__ __device__ int nref(const Stack &st) { return kTuneB * st.L; } +inline __host__ __device__ bool use_rp3i(int, int F) { return F == 2 || (size_t)kTuneB * (4 * F + 2) > 400; } +inline __host__ __device__ bool use_rp3s(int, int F) { return F == 2 || (size_t)kTuneB * (4 * F + 2) > 400; } +__host__ __device__ Plan build_plan(int B, int F, __half *base) { + Plan p; + p.s0 = mk(B, 2 * F, 3); + p.s1 = mk(B, 4 * F, 3); + p.s2 = mk(B, 36 * F, 7); + p.s3 = mk(B, 252 * F, 7); + p.s4 = mk(B, 1764 * F, 7); + size_t off = 0; + auto alloc = [&](int ch, const Stack &s) -> __half * { + __half *r = base ? base + off : nullptr; + off += (size_t)ch * s.SC; + return r; + }; + p.LB = alloc(40, p.s0); + p.X0 = alloc(768, p.s0); + p.M0 = alloc(1536, p.s0); + p.Y0 = alloc(768, p.s0); + p.A = alloc(768, p.s1); + p.Ac = use_rp3i(B, F) ? alloc(768, p.s1) : nullptr; + p.Bb = alloc(768, p.s1); + p.M1 = alloc(1536, p.s1); + p.P = alloc(384, p.s2); + p.Q = alloc(384, p.s2); + p.M2 = alloc(768, p.s2); + p.R = alloc(128, p.s3); + p.S = alloc(128, p.s3); + p.M3 = alloc(256, p.s3); + p.U = alloc(32, p.s4); + for (int j = 0; j < RP_JOBS; ++j) { + p.W7[j] = base ? base + off : nullptr; + off += (size_t)native_RP_COUTP(j) * native_RP_CIN(j) * 7; + } + p.W3i = nullptr; + if (use_rp3i(B, F)) { + p.W3i = base ? base + off + (size_t)3 * 1536 * 768 : nullptr; + off += (size_t)7 * 3 * 1536 * 768; + } + p.W3s = nullptr; + if (use_rp3s(B, F)) { + p.W3s = base ? base + off + (size_t)3 * 768 * 1536 : nullptr; + off += (size_t)7 * 3 * 768 * 1536; + } + p.total = off; + return p; +} +__host__ __device__ const __half *H16(const void *p) { return static_cast(p); } +__host__ __device__ const float *F32(const void *p) { return static_cast(p); } +} +size_t codec_scratch_bytes(int B, int F) { + return build_plan(B, F, nullptr).total * sizeof(__half); +} +template +__device__ __forceinline__ void wg_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, bool prep = false) { + constexpr int WG_BM = 64 * NWG, WG_NT = 128 * NWG, WG_CLD = WG_BN + 8; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int GEMM_B = STAGES * (A_E + B_E) * 2; + constexpr int EPI_B = WG_BM * WG_CLD * 4; + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute(wg_conv_kernel, SH); + return; + } + native_invoke(wg_conv_kernel,dim3(dim3(DIVUP(Cout, WG_BM), DIVUP(st.N, WG_BN))),dim3(WG_NT),SH, + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg2_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, + bool prep = false) { + constexpr int BK = CIB * KT; + constexpr int WG_BM = 64 * MSPLIT; + constexpr int WG_NT = 128 * MSPLIT * KSPLIT * NSPLIT; + constexpr int WG_BNT = WG_BN * NSPLIT; + constexpr int NCRAW = WG_BNT / 4; + constexpr int EPN = + (WG_NT % NCRAW == 0) ? WG_BNT : (((WG_BNT + 127) / 128) * 128); + constexpr int WG_CLD = EPN + 8; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int GEMM_B = KSPLIT * (SA * A_E + SB * NSPLIT * B_E) * 2; + constexpr int EPI_B = KSPLIT * WG_BM * WG_CLD * 4; + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute( + wg2_conv_kernel, + SH); + return; + } + native_invoke_pf(wg2_conv_kernel,dim3(dim3(DIVUP(Cout, WG_BM), DIVUP(st.N, WG_BNT))),dim3(WG_NT),SH, + w, (size_t)KT * Cin * Cout * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg3_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, + bool prep = false) { + constexpr int WG_BM = 64 * NWG, WG_NT = 128 * NWG * NSG; + constexpr int WG_BNT = WG_BN * NSG, WG_CLD = WG_BNT + 8; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int GEMM_B = STAGES * (A_E + NSG * B_E) * 2; + constexpr int EPI_B = WG_BM * WG_CLD * 4; + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute(wg3_conv_kernel, SH); + return; + } + native_invoke(wg3_conv_kernel,dim3(dim3(DIVUP(Cout, WG_BM), DIVUP(st.N, WG_BNT))),dim3(WG_NT),SH, + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg4_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, int CoutP, + const Stack &st, int SCin, int do_snake, int pad_off, + int pad_len, bool prep = false) { + constexpr int NT = 128 * MSPLIT * NSPLIT; + constexpr int BM = 64 * MSPLIT, BNT = WG_BN * NSPLIT; + constexpr int NBR = BNT + KT - 1; + constexpr int NBRP = BC8 ? ((NBR + 7) & ~7) : ((NBR + 31) & ~31); + constexpr int NCRAW = BNT / 4; + constexpr int EPN = + (NT % NCRAW == 0) ? BNT : (((BNT + 127) / 128) * 128); + constexpr int B_E = CIB * NBRP; + constexpr int GEMM_B = STAGES * B_E * 2; + constexpr int EPI_B = BM * (EPN + 8) * (PREACT ? 2 : 4); + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute( + wg4_conv_kernel, + SH); + return; + } + pdl_launch_pf(wg4_conv_kernel, + dim3(DIVUP(Cout, BM), DIVUP(st.N, BNT)), dim3(NT), SH, s, + w, (size_t)KT * CoutP * Cin * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, + st.T, st.L, SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg5_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, + bool prep = false) { + constexpr int BM = 64 * MSPLIT; + constexpr int NT = 128 * MSPLIT * KSPLIT; + constexpr int NCRAW = WG_BN / 4; + constexpr int EPN = (NT % NCRAW == 0) ? WG_BN : (((WG_BN + 127) / 128) * 128); + constexpr int NBRP = BC8 ? (((WG_BN + 2) + 7) & ~7) + : (((WG_BN + 2) + 31) & ~31); + constexpr int B_E = CIB * NBRP; + constexpr int GEMM_B = KSPLIT * SB * B_E * 2; + constexpr int EPI_COLS = OC8 ? WG_BN : EPN; + constexpr int EPI_CLD_R = EPI_COLS + 8; + constexpr int EPI_CLD = (EPI_CLD_R % 16 == 0) ? EPI_CLD_R + 4 : EPI_CLD_R; + constexpr int EPI_B = KSPLIT * BM * EPI_CLD * (OC8 ? 2 : 4); + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + static_assert(SH <= 227 * 1024, "wg5 shared over budget"); + if (prep) { + prepare_attribute( + wg5_conv_kernel, + SH); + return; + } + if constexpr (PDL) { + cudaLaunchConfig_t cfg = {}; + cudaLaunchAttribute attr[1] = {}; + native_config_grid(cfg) = dim3(DIVUP(Cout, BM), DIVUP(st.N, WG_BN)); + native_config_threads(cfg) = dim3(NT); + cfg.dynamicSmemBytes = SH; + cfg.stream = s; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = 1; + cfg.attrs = attr; + cfg.numAttrs = 1; + native_cudaLaunchKernelEx_pf( + &cfg, + wg5_conv_kernel, + w, (size_t)3 * Cin * Cout * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); + } else { + native_invoke_pf(wg5_conv_kernel,dim3(dim3(DIVUP(Cout, BM), DIVUP(st.N, WG_BN))),dim3(NT),SH, + w, (size_t)3 * Cin * Cout * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, + st.L, SCin, st.SC, st.N, do_snake, pad_off, pad_len); + } +} +#ifndef D7W +#define D7W 96 +#endif +#define D7_M2_C8I(BN,LATE) wg4_launch<7, 16, 2, 1, D7W, 4, 1, false, 2, true, false, LATE>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_M2P_C8O(BN,LATE) wg4_launch<7, 16, 2, 1, D7W, 4, 1, true, 2, false, true, LATE>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_M2N2(BN) wg4_launch<7, 16, 2, 2, BN, 4, 1, false, 2>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_M1(BN) wg4_launch<7, 16, 1, 1, BN, 4, 2, false, 2>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_C8I(W) \ + wg4_launch<7, 16, 2, 1, W, 4, 1, false, 2, true, false, false>( \ + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, \ + SCin, 1, pad_off, pad_len, prep) +#define D7_C8O(W) \ + wg4_launch<7, 16, 2, 1, W, 4, 1, true, 2, false, true, false>( \ + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, \ + SCin, 1, pad_off, pad_len, prep) +template +__device__ __forceinline__ int d7_band_cost(int N, int mtiles, int resident) { + const int v = mtiles * DIVUP(N, BN); + const int sub = (DIVUP(v, resident) >= 2) ? 2 : 1; + return DIVUP(DIVUP(v, sub), resident) * sub * (128 + BN + 6); +} +__device__ __forceinline__ int d7_band(int N, int mtiles, int resident) { + int best = 64, bc = d7_band_cost<64>(N, mtiles, resident); + const int c96 = d7_band_cost<96>(N, mtiles, resident); + if (c96 < bc) { bc = c96; best = 96; } + if (d7_band_cost<112>(N, mtiles, resident) < bc) best = 112; + return best; +} +__device__ __forceinline__ void dense7(cudaStream_t s, int job, const __half *in, const __half *w, + const float *bias, const float *alpha, const __half *resid, + __half *out, const __half *pad_src, __half *ns, + const Stack &st, int SCin, int pad_off, int pad_len, + bool prep = false) { + const int Cin = native_RP_CIN(job), Cout = native_RP_COUT(job), CoutP = native_RP_COUTP(job); + const int b_ = d7_band(st.N, DIVUP(Cout, 128), native_resident()); + if ((job & 1) == 0) { + if (b_ == 64) D7_C8O(64); + else if (b_ == 96) D7_C8O(96); + else D7_C8O(112); + } else { + if (b_ == 64) D7_C8I(64); + else if (b_ == 96) D7_C8I(96); + else D7_C8I(112); + } +} +__device__ __forceinline__ int wg5_band_skip(int N, int mtiles, int resident) { + const int wmin = DIVUP(mtiles * DIVUP(N, 104), resident); + if (DIVUP(mtiles * DIVUP(N, 24), resident) == wmin) return 24; + if (DIVUP(mtiles * DIVUP(N, 40), resident) == wmin) return 40; + if (DIVUP(mtiles * DIVUP(N, 48), resident) == wmin) return 48; + if (DIVUP(mtiles * DIVUP(N, 64), resident) == wmin) return 64; + return 104; +} +__device__ __forceinline__ int wg5_band_main_long(int N, int mtiles, + int resident) { + const int wmin = DIVUP(mtiles * DIVUP(N, 104), resident); + if (DIVUP(mtiles * DIVUP(N, 64), resident) == wmin) return 64; + return 104; +} +__device__ __forceinline__ int wg5_band_main_short(int N, int mtiles, + int resident) { + const int wmin = DIVUP(mtiles * DIVUP(N, 96), resident); + if (DIVUP(mtiles * DIVUP(N, 56), resident) == wmin) return 56; + if (DIVUP(mtiles * DIVUP(N, 80), resident) == wmin) return 80; + return 96; +} +#define DENSE_LAUNCH(KT,BM,BN,NW,WM,WN) native_invoke(dense_conv_kernel,dim3(dim3(DIVUP(Cout, BM), DIVUP(st.N, BN))),dim3(NW * 32),0, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, SCin, st.SC, st.N, do_snake, pad_off, pad_len) +__device__ __forceinline__ void dense(cudaStream_t s, int KT, const __half *in, const __half *w, + const float *bias, const float *alpha, const __half *resid, + __half *out, const __half *pad_src, __half *ns, int Cin, + int Cout, const Stack &st, int SCin, int do_snake, + int pad_off, int pad_len, const __half *w3 = nullptr, + int layout_mode = 0, bool prep = false) { + if (KT == 3 && w3 != nullptr) { + if (Cin == 1536) { +#define WG5S(BN,O,R,PDL) wg5_launch<16, 1, 4, BN, 4, 3, 1, 1, true, O, R, 1536, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5SW(BN,O,R,PDL) wg5_launch<32, 1, 2, BN, 4, 3, 1, 1, true, O, R, 1536, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5S3(O,R) do { const int n_ = nref(st); if (n_ > 500 && n_ <= 640) { WG5SW(56, O, R, false); } else { const int b_ = wg5_band_skip(st.N, DIVUP(Cout, 64), native_resident()); if (b_ == 24) WG5S(24, O, R, true); else if (b_ == 40) WG5S(40, O, R, true); else if (b_ == 48) WG5S(48, O, R, true); else if (b_ == 64) WG5S(64, O, R, true); else WG5S(104, O, R, true); } } while (0) + if (layout_mode == 1) WG5S3(true, false); + else if (layout_mode == 2) WG5S3(true, true ); + else if (layout_mode == 3) WG5S3(false, true ); + else WG5S3(false, false); +#undef WG5S3 +#undef WG5SW +#undef WG5S + } else { +#define WG5I(BN,BC,MS,PDL) wg5_launch<16, MS, 2, BN, 4, 3, 1, 1, BC, true, false, 768, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5I4(BN,BC,PDL) wg5_launch<16, 1, 4, BN, 4, 3, 1, 1, BC, true, false, 768, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5I2W(BN,BC,PDL) wg5_launch<32, 1, 2, BN, 4, 3, 1, 1, BC, true, false, 768, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5I3(BC) do { const int n_ = nref(st); if (n_ > 640) { const int b_ = wg5_band_main_long(st.N, DIVUP(Cout, 128), native_resident()); if (b_ == 64) WG5I(64, BC, 2, true); else WG5I(104, BC, 2, true); } else if (n_ > 500) { WG5I2W(64, BC, false); } else { const int b_ = wg5_band_main_short(st.N, DIVUP(Cout, 64), native_resident()); if (b_ == 56) WG5I4(56, BC, true); else if (b_ == 80) WG5I4(80, BC, true); else WG5I4(96, BC, true); } } while (0) + if (layout_mode == 1) WG5I3(true); + else WG5I3(false); +#undef WG5I3 +#undef WG5I2W +#undef WG5I4 +#undef WG5I + } + } else if (KT == 3) { + const int n_ = nref(st); + if (Cin == 1536 && Cout == 768) { + if (n_ <= 704) + wg2_launch<3, 16, 1, 4, 1, 64, 2, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else + wg2_launch<3, 16, 1, 4, 1, 128, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + } else if (Cin == 768 && Cout == 1536) { + if (n_ > 1000) + wg2_launch<3, 16, 1, 2, 1, 128, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ > 640 && n_ <= 768) + wg2_launch<3, 16, 1, 2, 1, 64, 2, 2, 2>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ <= 320) + wg2_launch<3, 16, 1, 4, 1, 64, 2, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ > 384 && n_ <= 512) + wg2_launch<3, 16, 1, 4, 1, 112, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ > 512 && n_ <= 640) + wg2_launch<3, 16, 1, 4, 1, 120, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else + wg2_launch<3, 16, 1, 4, 1, 128, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + } else { + wg2_launch<3, 16, 1, 1, 1, 64, 3, 2, 2>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, + SCin, do_snake, pad_off, pad_len, prep); + } + } else if (Cout >= 128 && + ((Cin >= 384 && st.T == 288) || + (Cin < 384 && st.T == 2016))) { + wg3_launch<7, 112, 2, 2, 1, 128, 2>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, + SCin, do_snake, pad_off, pad_len); + } else { + wg3_launch<7, 112, 1, 4, 1, 128, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, + SCin, do_snake, pad_off, pad_len); + } +} +__device__ void codec_program(const CodecArgs &a, const Plan &p, const RepackSet &js, + int repack_grid, + const __half *const *state_override = nullptr) { + cudaStream_t s = nullptr; + const int B = a.B; + const void *const *W = a.inp; + const __half *const *ST = state_override ? state_override : + reinterpret_cast(a.st); + __half *const *NS = reinterpret_cast<__half *const *>(a.ns); + dense(s, 3, p.LB, H16(W[1]), F32(W[2]), nullptr, nullptr, p.X0, ST[1], NS[0], + 40, 768, p.s0, p.s0.SC, 0, 0, 2); + const __half *k3iw = H16(W[3]); + const float *k3ib = F32(W[4]); + const float *k3ia = F32(W[5]); + const __half *k3sw = H16(W[6]); + const float *k3sb = F32(W[7]); + const float *k3sa = F32(W[8]); + const size_t IW = (size_t)1536 * 768 * 3, SW = (size_t)768 * 1536 * 3; + dense(s, 3, p.X0, k3iw, k3ib, k3ia, nullptr, p.M0, ST[2], NS[1], 768, 1536, + p.s0, p.s0.SC, 1, 0, 2, p.W3i ? p.W3i - IW : nullptr, 0); + dense(s, 3, p.M0, k3sw, k3sb, k3sa, p.X0, p.Y0, ST[3], NS[2], 1536, 768, + p.s0, p.s0.SC, 1, 1, 1, p.W3s ? p.W3s - SW : nullptr, 0); + TCONV_BPACK_DISPATCH8(1, 2, 8, p.s0.T, 768, B, + p.Y0, H16(W[9]), F32(W[10]), F32(W[11]), p.A, ST[4], NS[3], 768, 768, + p.s0.T, p.s0.L, p.s0.SC, 1, p.s1.L, p.s1.SC, 2, 0, 2, B); + if (p.W3i != nullptr) + native_invoke(repack_activation_only,dim3(min(1152,96*DIVUP(p.s1.SC,128))),dim3(128),0, + p.A, p.Ac, p.s1.SC, p.s1.SC, DIVUP(p.s1.SC,128)); + __half *cur = p.A, *nxt = p.Bb; + for (int i = 0; i < 6; ++i) { + const int bi = i + 1; + const __half *ain = (i == 0 && p.Ac) ? p.Ac : cur; + dense(s, 3, ain, k3iw + bi * IW, k3ib + bi * 1536, k3ia + bi * 768, + nullptr, p.M1, ST[5 + 2 * i], NS[4 + 2 * i], 768, 1536, p.s1, + p.s1.SC, 1, 0, 2, p.W3i ? p.W3i + (bi - 1) * IW : nullptr, + (i > 0 || p.Ac) ? 1 : 0); + const __half *ps = (i < 5) ? ST[6 + 2 * i] : ST[16]; + const int po = (i < 5) ? 0 : 1, pl = 1 + (i < 5); + dense(s, 3, p.M1, k3sw + bi * SW, k3sb + bi * 768, k3sa + bi * 384, cur, + nxt, ps, NS[5 + 2 * i], 1536, 768, p.s1, p.s1.SC, 1, po, pl, + p.W3s ? p.W3s + (bi - 1) * SW : nullptr, + i == 0 ? 1 : (i < 5 ? 2 : 3)); + __half *t = cur; cur = nxt; nxt = t; + } + TCONV_BPACK_DISPATCH1(2, 9, 16, p.s1.T, 384, B, + cur, H16(W[12]), F32(W[13]), F32(W[14]), p.P, ST[17], NS[16], 768, 384, + p.s1.T, p.s1.L, p.s1.SC, 1, p.s2.L, p.s2.SC, 6, 0, 6, B); + dense7(s, 0, p.P, p.W7[0], F32(W[22]), F32(W[23]), nullptr, p.M2, ST[18], + NS[17], p.s2, p.s2.SC, 0, 6); + dense7(s, 1, p.M2, p.W7[1], F32(W[25]), F32(W[26]), p.P, p.Q, ST[19], + NS[18], p.s2, p.s2.SC, 5, 1); + TCONV_BPACK_DISPATCH(3, 7, 16, p.s2.T, 128, B, + p.Q, H16(W[15]), F32(W[16]), F32(W[17]), p.R, ST[20], NS[19], 384, 128, + p.s2.T, p.s2.L, p.s2.SC, 5, p.s3.L, p.s3.SC, 6, 0, 6, B); + dense7(s, 2, p.R, p.W7[2], F32(W[28]), F32(W[29]), nullptr, p.M3, ST[21], + NS[20], p.s3, p.s3.SC, 0, 6); + dense7(s, 3, p.M3, p.W7[3], F32(W[31]), F32(W[32]), p.R, p.S, ST[22], + NS[21], p.s3, p.s3.SC, 5, 1); + if (p.s3.T >= 1500) + pdl_launch(transpose_conv_kernel<4, 7, 8, true>, dim3(1, 32, B), + dim3(tc_blk(p.s3.T, 8)), + (size_t)tc_blk(p.s3.T, 8) * (8 * 7) * 2 + 64, s, + p.S, H16(W[18]), F32(W[19]), F32(W[20]), p.U, ST[23], + NS[22], 128, 32, p.s3.T, p.s3.L, p.s3.SC, 5, p.s4.L, + p.s4.SC, 6, 0, 6, B); + else + pdl_launch(transpose_conv_kernel<4, 7, 4, true>, dim3(1, 32, B), + dim3(tc_blk(p.s3.T, 4)), + (size_t)tc_blk(p.s3.T, 4) * (4 * 7) * 2 + 64, s, + p.S, H16(W[18]), F32(W[19]), F32(W[20]), p.U, ST[23], + NS[22], 128, 32, p.s3.T, p.s3.L, p.s3.SC, 5, p.s4.L, + p.s4.SC, 6, 0, 6, B); + { + const int ntile = DIVUP(p.s4.N, TL_BN); + const int grid = ntile < 132 ? ntile : 132; + pdl_launch_pf(tail_fused_kernel, dim3(grid), dim3(TL_NT), TL_SHARED, s, + p.U, (size_t)32 * p.s4.SC * sizeof(__half), + p.U, p.W7[4], p.W7[5], H16(W[39]), F32(W[34]), F32(W[35]), + F32(W[37]), F32(W[38]), F32(W[40]), ST[24], ST[25], + static_cast<__half *>(a.waveform), NS[23], NS[24], NS[25], B, + p.s4.T, p.s4.L, p.s4.SC, p.s4.N, ntile); + } +} diff --git a/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/native_stage1.cu b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/native_stage1.cu new file mode 100644 index 0000000..36bac91 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/native_stage1.cu @@ -0,0 +1,357 @@ +#include "native_stage1.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#ifndef NATIVE_CODEC_PREFETCH +#define NATIVE_CODEC_PREFETCH 1 +#endif +#ifndef NATIVE_CODEC_PACK_SUBCTAS +#define NATIVE_CODEC_PACK_SUBCTAS 1 +#endif +namespace cg=cooperative_groups; +constexpr int PHYSICAL_THREADS=512,MAIN_ARENA=227840,AUX_BYTES=4096,META_OFFSET=231936,TOTAL_SHARED=232192; +constexpr int SLICE_BYTES=MAIN_ARENA+AUX_BYTES; +constexpr int MAX_SUBCTA=8; +struct NativeMeta { + dim3 grid,threads; + uint3 block[MAX_SUBCTA]; + NativeStage1Status *status;const int32_t *lengths;int padded_frames; + int submask,tshift,tmask,slice; + int resident; + int identity_slots; +}; +static_assert(sizeof(NativeMeta)<=TOTAL_SHARED-META_OFFSET); +__device__ __forceinline__ char *native_base(){extern __shared__ __align__(16) char arena[];return arena;} +__device__ __forceinline__ NativeMeta *native_meta(){return reinterpret_cast(native_base()+META_OFFSET);} +__device__ __forceinline__ int native_subid(){const NativeMeta *m=native_meta();return (int)((threadIdx.x>>m->tshift)&(unsigned)m->submask);} +__device__ __forceinline__ char *native_arena(){const NativeMeta *m=native_meta();return native_base()+native_subid()*m->slice;} +__device__ __forceinline__ char *native_aux(){const NativeMeta *m=native_meta();return native_base()+native_subid()*m->slice+(m->slice-AUX_BYTES);} +__device__ __forceinline__ uint3 native_tid(){uint3 r;r.x=threadIdx.x&(unsigned)native_meta()->tmask;r.y=0u;r.z=0u;return r;} +__device__ __forceinline__ dim3 native_block(){const uint3 b=native_meta()->block[native_subid()];return dim3(b.x,b.y,b.z);} +__device__ __forceinline__ dim3 native_grid(){return native_meta()->grid;} +__device__ __forceinline__ dim3 native_threads(){return native_meta()->threads;} +__device__ __forceinline__ int native_resident(){return native_meta()->resident;} +__device__ __forceinline__ void native_local_sync(){ + const NativeMeta *m=native_meta();unsigned n=m->threads.x,id=15u-(unsigned)native_subid(); + asm volatile("barrier.sync %0, %1;"::"r"(id),"r"(n):"memory"); +} +__device__ __forceinline__ dim3 &native_config_grid(cudaLaunchConfig_t &c){return c.gridDim;} +__device__ __forceinline__ dim3 &native_config_threads(cudaLaunchConfig_t &c){return c.blockDim;} +__device__ __forceinline__ int native_true_tokens(int row,int padded_tokens){ + const NativeMeta *m=native_meta(); + return (padded_tokens/m->padded_frames)*m->lengths[row]; +} +__device__ __forceinline__ void native_warm(const char *p,size_t n,int k,int nk,int me){ + if(p==nullptr||n<16||nk<=0)return; + const size_t a=(size_t)p&15u;p-=a;n=(n+a)&~(size_t)15; + const uint4 *q=reinterpret_cast(p); + const size_t cnt=n>>4,step=(size_t)nk*PHYSICAL_THREADS; + uint32_t acc=0u; + for(size_t i=(size_t)k*PHYSICAL_THREADS+(size_t)me;i __device__ __forceinline__ void native_invoke_pf(F fn,dim3 g,dim3 t,size_t shared,const void *pf,size_t pfb,A&&... args){ + const int me=threadIdx.x,owner=blockIdx.x,owners=gridDim.x; + const cg::grid_group all=cg::this_grid();NativeMeta *meta=native_meta(); + if(t.x>512||t.y!=1||t.z!=1||t.x%32||shared>MAIN_ARENA){ + if(me==0&&owner==0)meta->status->code=10; + all.sync();return; + } + int sub=1,tshift=9; + const unsigned tx=t.x; + const int vtotal=(int)(g.x*g.y*g.z); + if(NATIVE_CODEC_PACK_SUBCTAS && (tx&(tx-1))==0u){ + int s=0;while((1u<MAX_SUBCTA)cand=MAX_SUBCTA; + const int want=(vtotal+owners-1)/owners; + while(cand>1&&cand>want)cand>>=1; + while(cand>1&&shared>(size_t)(SLICE_BYTES/cand-AUX_BYTES))cand>>=1; + sub=cand; + } + const int slice=(SLICE_BYTES/sub)&~127; + if(me==0){meta->grid=g;meta->threads=t;meta->submask=sub-1;meta->tshift=tshift; + meta->tmask=(int)((1u<slice=slice;} + const int total=vtotal,grp=me>>tshift; + const bool holder=me=busy&&pf!=nullptr)native_warm((const char*)pf,pfb,owner-busy,owners-busy,me); + for(int base=owner*sub;baseblock[grp]=(g.y==1u&&g.z==1u)?make_uint3((unsigned)i,0u,0u) + :make_uint3(i%g.x,(i/g.x)%g.y,i/(g.x*g.y)); + __syncthreads(); + if(holder&&i(args)...); + asm volatile("cp.async.wait_all;":::"memory"); + if(t.x>=128)asm volatile("wgmma.wait_group.sync.aligned 0;":::"memory"); + } + __syncthreads(); + } + all.sync(); +} +template __device__ __forceinline__ void native_invoke(F fn,dim3 g,dim3 t,size_t shared,A&&... args){ + native_invoke_pf(fn,g,t,shared,nullptr,(size_t)0,static_cast(args)...); +} +template __device__ __forceinline__ void native_cudaLaunchKernelEx(cudaLaunchConfig_t *c,F fn,A... args){ + native_invoke(fn,c->gridDim,c->blockDim,c->dynamicSmemBytes,args...); +} +template __device__ __forceinline__ void native_cudaLaunchKernelEx_pf(cudaLaunchConfig_t *c,F fn,const void *pf,size_t pfb,A... args){ + native_invoke_pf(fn,c->gridDim,c->blockDim,c->dynamicSmemBytes,pf,pfb,args...); +} +#define blockIdx native_block() +#define gridDim native_grid() +#define blockDim native_threads() +#define threadIdx native_tid() +#define __syncthreads() native_local_sync() +#include "legacy_codec.cuh" +#undef __syncthreads +#undef threadIdx +#undef blockDim +#undef gridDim +#undef blockIdx +namespace { +struct FastArgs { + CodecArgs codec;Plan plan;RepackSet repack;int repack_grid;int input_frames;int blocks; + const int64_t *codes;const int32_t *lengths,*slots; + const __half *state[26];__half *next[26];float *pcm;int32_t *samples; + NativeStage1Status *status;int nslots; + int channels[26],histories[26]; +}; +static_assert(sizeof(FastArgs)<=4096,"precomputed kernel argument must fit the conservative parameter budget"); +struct Handle {FastArgs args{};Plan plans[17]{};RepackSet repacks[17]{};int repack_grids[17]{};void *scratch=nullptr;size_t bytes=0;void *immutable=nullptr;size_t immutable_bytes=0;int setup_launches=0;std::vector packed;int max_frames=0,blocks=0,device=-1;}; +thread_local std::string error; +bool check(cudaError_t e,const char *where){if(e==cudaSuccess)return true;error=std::string(where)+": "+cudaGetErrorString(e);return false;} +__device__ void reject(NativeStage1Status *s,int code,int row,int frame,int detail){ + if(atomicCAS(&s->code,0,code)==0){s->row=row;s->frame=frame;s->detail=detail;} +} +__global__ __launch_bounds__(512,1) void stage1_fast_persistent(const __grid_constant__ FastArgs a){ + { + const auto all=cg::this_grid(); + const size_t me=(size_t)blockIdx.x*512+threadIdx.x,step=(size_t)gridDim.x*512; + if(threadIdx.x==0){native_meta()->status=a.status;native_meta()->lengths=a.lengths;native_meta()->padded_frames=a.codec.F;native_meta()->resident=a.blocks;native_meta()->identity_slots=(a.nslots==a.codec.B);} + if(me==0)*a.status=NativeStage1Status{0,-1,-1,0}; + all.sync();const int B=a.codec.B,F=a.codec.F,inputF=a.input_frames; + for(int row=threadIdx.x;rowidentity_slots,0); + __syncthreads(); + for(int row=me;rowinputF)reject(a.status,8,row,-1,n); + if(slot<0||slot>=a.nslots)reject(a.status,2,row,-1,slot); + for(int j=0;j=0&&n<=inputF)?n*1764:0; + } + for(size_t i=me;i<(size_t)B*inputF;i+=step){ + int row=i/inputF,t=i%inputF;bool any=false; + for(int c=0;c<16;++c){int64_t v=a.codes[i*16+c];if(v<0||v>=1024)reject(a.status,4,row,t,c);any|=v!=0;} + if(t=a.lengths[row]&&any)reject(a.status,6,row,t,0); + } + if(a.nslots!=B) + for(int p=0;p<26;++p){ + const size_t vectors=(size_t)a.nslots*a.channels[p]*a.histories[p]/8; + const uint4 *src=reinterpret_cast(a.state[p]); + uint4 *dst=reinterpret_cast(a.next[p]); + for(size_t i=me;icode)return; + if(!native_meta()->identity_slots) for(int p=1;p<26;++p){ + const size_t stride=(size_t)a.channels[p]*a.histories[p]; + __half *dest=(__half*)a.codec.st[p]; + for(int row=blockIdx.x;row(a.state[p]+(size_t)slot*stride); + uint4 *dst=reinterpret_cast(dest+(size_t)row*stride); + for(size_t offset=threadIdx.x;offset>(2*(c%5)))&3)-2)*.5f; + } + v=__float2half_rn(x); + } + packed[(size_t)c*SC+(size_t)row*L+j]=v; + } + all.sync(); + } + codec_program(a.codec,a.plan,a.repack,a.repack_grid, + native_meta()->identity_slots?a.state:nullptr); + if(a.status->code)return; + const size_t me=(size_t)blockIdx.x*512+threadIdx.x,step=(size_t)gridDim.x*512; + const int B=a.codec.B,F=a.codec.F,inputF=a.input_frames; + for(int p=0;p<26;++p){ + const size_t stride=(size_t)a.channels[p]*a.histories[p];const __half *src=(const __half*)a.codec.ns[p]; + for(int row=blockIdx.x;row(src+(size_t)row*stride); + const uint4 *prior=reinterpret_cast(a.state[p]+(size_t)slot*stride); + uint4 *dest=reinterpret_cast(a.next[p]+(size_t)slot*stride); + for(size_t offset=threadIdx.x;offset(a.codec.waveform); + float2 *pcm=reinterpret_cast(a.pcm); + const int row_pairs=inputF*882; + for(size_t i=me;i<(size_t)B*row_pairs;i+=step){ + const int row=i/row_pairs,t=i%(size_t)row_pairs; + pcm[i]=tstatus=a.status;native_meta()->resident=a.blocks;} + if(blockIdx.x==0&&threadIdx.x==0)*a.status=NativeStage1Status{0,-1,-1,0}; + all.sync(); + native_invoke(repack_rs_kernel<7>,dim3(a.repack_grid),dim3(128),MAIN_ARENA, + a.repack,H16(a.codec.inp[0]),H16(a.codec.st[0]),a.plan.LB, + 40,a.plan.s0.T,a.plan.s0.L,a.plan.s0.SC,2,a.codec.B); + native_invoke(repack3_rs_kernel,dim3(1152,12),dim3(128),MAIN_ARENA, + H16(a.codec.inp[3]),H16(a.codec.inp[6]),a.plan.W3i,a.plan.W3s, + a.plan.A,a.plan.Ac,a.plan.s1.SC,a.plan.s1.SC,DIVUP(a.plan.s1.SC,128)); +} +bool stack(Handle *h,int slot,const NativeStage1Weights *w,int first,int field,size_t elements){ + const int ids[7]={first,first+3,first+5,first+7,first+9,first+11,first+13}; + const size_t stride=elements*(field==0?sizeof(__half):sizeof(float));void *ptr=nullptr; + if(!check(cudaMalloc(&ptr,7*stride),"packed-weight allocation"))return false; + h->packed.push_back(ptr);h->args.codec.inp[slot]=ptr; + for(int i=0;i<7;++i){const void *src=field==0?w[ids[i]].weight:field==1?(const void*)w[ids[i]].bias:(const void*)w[ids[i]].alpha; + if(!check(cudaMemcpy((char*)ptr+i*stride,src,stride,cudaMemcpyDeviceToDevice),"immutable weight stack"))return false;} + return true; +} +void release(Handle *h){if(h->immutable)cudaFree(h->immutable);if(h->scratch)cudaFree(h->scratch);for(void *p:h->packed)cudaFree(p);delete h;} +} +extern "C" void *native_stage1_create(int b,int maxf,int slots,int blocks,const NativeStage1Weights *w){ + error.clear();if((b!=1&&b!=16&&b!=32&&b!=64&&b!=128&&b!=256&&b!=512&&b!=1024&&b!=2048)||maxf<1||maxf>16||slots4096||blocks<0||!w){error="invalid create arguments";return nullptr;} + for(int p=0;p<26;++p)if(!w[p].weight||!w[p].bias||(p!=0&&p!=25&&!w[p].alpha)){error="missing weight pointer";return nullptr;} + Handle *h=new(std::nothrow)Handle;if(!h){error="host allocation";return nullptr;} + h->max_frames=maxf;h->args.nslots=slots;h->args.codec.B=b; + const int cs[26]={40,768,1536,768,768,1536,768,1536,768,1536,768,1536,768,1536,768,1536,768,384,768,384,128,256,128,32,64,32}; + const int hs[26]={2,2,2,1,2,2,2,2,2,2,2,2,2,2,2,2,1,6,6,1,6,6,1,6,6,6}; + cudaDeviceProp prop{};int occupancy=0,coop=0; + if(!check(cudaGetDevice(&h->device),"device")||!check(cudaGetDeviceProperties(&prop,h->device),"properties")|| + !check(cudaDeviceGetAttribute(&coop,cudaDevAttrCooperativeLaunch,h->device),"cooperative support")|| + !check(cudaFuncSetAttribute(stage1_fast_persistent,cudaFuncAttributeMaxDynamicSharedMemorySize,TOTAL_SHARED),"shared-memory optin")|| + !check(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&occupancy,stage1_fast_persistent,512,TOTAL_SHARED),"occupancy")) {release(h);return nullptr;} + if(prop.major!=9||prop.minor!=0||!coop||occupancy<1){error="requires resident SM90 cooperative kernel";release(h);return nullptr;} + h->blocks=blocks?blocks:occupancy*prop.multiProcessorCount; + if(h->blocks>occupancy*prop.multiProcessorCount){error="grid exceeds cooperative occupancy";release(h);return nullptr;} + h->args.blocks=h->blocks; + const int capacityF=maxf<=3?3:(maxf<=5?5:std::max(8,maxf)); + size_t decoder=codec_scratch_bytes(b,capacityF),halfs=(size_t)b*40*2*capacityF+(size_t)b*1764*capacityF; + for(int p=0;p<26;++p){h->args.channels[p]=cs[p];h->args.histories[p]=hs[p];halfs+=2*(size_t)b*cs[p]*hs[p];} + h->bytes=decoder+halfs*sizeof(__half); + if(!check(cudaMalloc(&h->scratch,h->bytes),"scratch allocation")){release(h);return nullptr;} + if(!check(cudaMemset(h->scratch,0,h->bytes),"initial scratch padding")){release(h);return nullptr;} + h->args.codec.scratch=h->scratch;__half *q=(__half*)((char*)h->scratch+decoder); + h->args.codec.inp[0]=q;q+=(size_t)b*40*2*capacityF;h->args.codec.waveform=q;q+=(size_t)b*1764*capacityF; + for(int p=0;p<26;++p){h->args.codec.st[p]=q;q+=(size_t)b*cs[p]*hs[p];h->args.codec.ns[p]=q;q+=(size_t)b*cs[p]*hs[p];} + const void **in=h->args.codec.inp; + in[1]=w[0].weight;in[2]=w[0].bias; + if(!stack(h,3,w,1,0,(size_t)1536*768*3)||!stack(h,4,w,1,1,1536)||!stack(h,5,w,1,2,768)|| + !stack(h,6,w,2,0,(size_t)768*1536*3)||!stack(h,7,w,2,1,768)||!stack(h,8,w,2,2,384)){release(h);return nullptr;} + const int upops[4]={3,16,19,22};for(int i=0;i<4;++i){in[9+3*i]=w[upops[i]].weight;in[10+3*i]=w[upops[i]].bias;in[11+3*i]=w[upops[i]].alpha;} + for(int i=0;i<3;++i){int p=17+3*i;in[21+6*i]=w[p].weight;in[22+6*i]=w[p].bias;in[23+6*i]=w[p].alpha; + in[24+6*i]=w[p+1].weight;in[25+6*i]=w[p+1].bias;in[26+6*i]=w[p+1].alpha;} + in[39]=w[25].weight;in[40]=w[25].bias; + __half *immutable_w7[6]{};__half *immutable_w3i=nullptr,*immutable_w3s=nullptr; + if(capacityF>=3){ + size_t wh=0;for(int j=0;j<6;++j)wh+=(size_t)native_RP_COUTP(j)*native_RP_CIN(j)*7; + const size_t IW=(size_t)3*1536*768;wh+=14*IW; + h->immutable_bytes=wh*sizeof(__half); + if(!check(cudaMalloc(&h->immutable,h->immutable_bytes+sizeof(NativeStage1Status)),"immutable weight allocation")){release(h);return nullptr;} + if(!check(cudaMemset(h->immutable,0,h->immutable_bytes+sizeof(NativeStage1Status)),"immutable padding initialization")){release(h);return nullptr;} + __half *q=(__half*)h->immutable; + for(int j=0;j<6;++j){immutable_w7[j]=q;q+=(size_t)native_RP_COUTP(j)*native_RP_CIN(j)*7;} + immutable_w3i=q+IW;q+=7*IW;immutable_w3s=q+IW;q+=7*IW; + if((char*)q!=(char*)h->immutable+h->immutable_bytes){error="immutable layout size mismatch";release(h);return nullptr;} + } + for(int f:{2,3,5,8}) if(f<=capacityF){ + Plan p=build_plan(b,f,static_cast<__half*>(h->scratch)); + for(int j=0;j<6;++j)p.W7[j]=immutable_w7[j];p.W3i=immutable_w3i;p.W3s=immutable_w3s; + const void *const *W=h->args.codec.inp; + RepackSet js{}; + const void *SRC[RP_JOBS] = {W[21], W[24], W[27], W[30], W[33], W[36]}; + for (int j = 0; j < RP_JOBS; ++j) { + js.src[j] = H16(SRC[j]); + js.dst[j] = p.W7[j]; + js.Cout[j] = native_RP_COUT(j); + js.CoutP[j] = native_RP_COUTP(j); + js.Cin[j] = native_RP_CIN(j); + } + int acc = 0; + for (int j = 0; j < RP_JOBS; ++j) { + js.base[j] = acc; + js.nch[j] = native_RP_CIN(j) / 16; + js.rs[j] = (j < 4) ? 1 : 0; + acc += (native_RP_COUTP(j) / 16) * DIVUP(native_RP_CIN(j), 64); + } + js.k3base = acc; + js.k3src[0] = H16(W[3]); + js.k3src[1] = H16(W[6]); + js.k3dst[0] = p.W3i ? p.W3i - (size_t)3 * 1536 * 768 : nullptr; + js.k3dst[1] = p.W3s ? p.W3s - (size_t)3 * 768 * 1536 : nullptr; + const int k3n = p.W3i ? 2304 : 0; + h->plans[f]=p;h->repacks[f]=js;h->repack_grids[f]=acc+k3n; + } + if(capacityF>=3){ + FastArgs a=h->args;a.codec.F=3;a.plan=h->plans[3];a.repack=h->repacks[3];a.repack_grid=h->repack_grids[3]; + a.status=(NativeStage1Status*)((char*)h->immutable+h->immutable_bytes); + int setup_occupancy=0; + if(!check(cudaFuncSetAttribute(stage1_prepare_immutable,cudaFuncAttributeMaxDynamicSharedMemorySize,TOTAL_SHARED),"setup shared-memory optin")|| + !check(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&setup_occupancy,stage1_prepare_immutable,512,TOTAL_SHARED),"setup occupancy")|| + setup_occupancy<1||h->blocks>setup_occupancy*prop.multiProcessorCount){ + if(error.empty())error="setup grid exceeds cooperative occupancy";release(h);return nullptr;} + void *params[]={&a}; + if(!check(cudaLaunchCooperativeKernel((void*)stage1_prepare_immutable,dim3(h->blocks),dim3(512),params,TOTAL_SHARED,nullptr),"immutable setup launch")|| + !check(cudaStreamSynchronize(nullptr),"immutable setup completion")){release(h);return nullptr;} + NativeStage1Status status{}; + if(!check(cudaMemcpy(&status,a.status,sizeof(status),cudaMemcpyDeviceToHost),"immutable setup status")){release(h);return nullptr;} + if(status.code){error="immutable setup device status "+std::to_string(status.code);release(h);return nullptr;} + h->setup_launches=1; + } + return h; +} +extern "C" int native_stage1_run(void *handle,const int64_t *codes,const int32_t *lengths,const int32_t *slots, + const void *const *states,void *const *next,float *pcm,int32_t *samples,NativeStage1Status *status,int f,void *stream){ + error.clear();Handle *h=(Handle*)handle; + if(!h||!codes||!lengths||!slots||!states||!next||!pcm||!samples||!status){error="null run argument";return 1;} + if(f<1||f>8||f>h->max_frames){error="true-state variant supports public F1..8";return 1;} + const int paddedF=f<=2?2:(f<=3?3:(f<=5?5:8)); + int device=-1;if(!check(cudaGetDevice(&device),"device")||device!=h->device){error="device mismatch";return 1;} + FastArgs a=h->args;a.input_frames=f;a.plan=h->plans[paddedF];a.repack=h->repacks[paddedF];a.repack_grid=h->repack_grids[paddedF];a.codec.F=paddedF;a.codes=codes;a.lengths=lengths;a.slots=slots;a.pcm=pcm;a.samples=samples;a.status=status; + for(int p=0;p<26;++p){if(!states[p]||!next[p]||states[p]==next[p]){error="invalid/aliased state";return 1;}a.state[p]=(const __half*)states[p];a.next[p]=(__half*)next[p];} + void *params[]={&a};if(!check(cudaLaunchCooperativeKernel((void*)stage1_fast_persistent,dim3(h->blocks),dim3(512),params,TOTAL_SHARED,(cudaStream_t)stream),"fast Stage1 launch"))return 1; + return check(cudaPeekAtLastError(),"immediate launch status")?0:1; +} +extern "C" int native_stage1_destroy(void *p){error.clear();if(!p)return 0;release((Handle*)p);return check(cudaPeekAtLastError(),"destroy status")?0:1;} +extern "C" const char *native_stage1_last_error(){return error.c_str();} +extern "C" size_t native_stage1_scratch_bytes(void *p){return p?((Handle*)p)->bytes:0;} +extern "C" int native_stage1_resident_blocks(void *p){return p?((Handle*)p)->blocks:0;} +extern "C" int native_stage1_setup_kernel_launches(void *p){return p?((Handle*)p)->setup_launches:0;} +extern "C" size_t native_stage1_immutable_weight_bytes(void *p){return p?((Handle*)p)->immutable_bytes:0;} diff --git a/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/native_stage1.h b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/native_stage1.h new file mode 100644 index 0000000..a95b5f0 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/native_stage1.h @@ -0,0 +1,26 @@ +#pragma once +#include +#include +struct NativeStage1Weights { const void *weight; const float *bias; const float *alpha; }; +struct NativeStage1Status { int32_t code, row, frame, detail; }; +#ifdef __cplusplus +extern "C" { +#endif +void *native_stage1_create(int batch, int max_frames, int state_slots, + int resident_blocks, const NativeStage1Weights *weights26); +int native_stage1_run(void *handle, const int64_t *codes, const int32_t *lengths, + const int32_t *slots, const void *const *states26, + void *const *next_states26, float *pcm, int32_t *samples, + NativeStage1Status *status, int frames, void *stream); +int native_stage1_run_profile(void *handle, const int64_t *codes, const int32_t *lengths, + const int32_t *slots, const void *const *states26, + void *const *next_states26, float *pcm, int32_t *samples, + NativeStage1Status *status, int frames, void *stream, + uint64_t *phase_timestamps30); +int native_stage1_destroy(void *handle); +const char *native_stage1_last_error(); +size_t native_stage1_scratch_bytes(void *handle); +int native_stage1_resident_blocks(void *handle); +#ifdef __cplusplus +} +#endif diff --git a/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/repack3.inc b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/repack3.inc new file mode 100644 index 0000000..2f61c93 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/repack3.inc @@ -0,0 +1,66 @@ +__device__ __noinline__ void repack3_kernel( + const __half *__restrict__ win, const __half *__restrict__ wsk, + __half *__restrict__ din, __half *__restrict__ dsk) { + int4 (*sm)[49]=reinterpret_cast(native_arena()); + const int job = blockIdx.y; + const int Cout = job ? 768 : 1536; + const int Cin = job ? 1536 : 768; + const __half *src = job ? wsk : win; + __half *dst = job ? dsk : din; + const int GPB = Cin >> 3; + const int Kg = Cin * 3; + const int GT = GPB >> 4; + const int RT = Cout >> 5; + const int t = threadIdx.x; + for (int tile = blockIdx.x; tile < 6 * RT * GT; tile += gridDim.x) { + const int gblk = tile % GT; + const int rest = tile / GT; + const int corow = rest % RT; + const int b = rest / RT; + const int co0 = corow * 32; + const int g0 = gblk * 16; + const __half *sbase = src + (size_t)(b + 1) * Cout * Kg + + (size_t)co0 * Kg + g0 * 24; +#pragma unroll + for (int j = 0; j < 6; ++j) { + const int idx = t + j * 256; + const int r = idx / 48, c = idx - (idx / 48) * 48; + sm[r][c] = *reinterpret_cast(sbase + (size_t)r * Kg + + c * 8); + } + __syncthreads(); + __half *dbase = dst + (size_t)b * (3 * (size_t)Cout * Cin); +#pragma unroll + for (int p = 0; p < 2; ++p) { + const int u = t + p * 256; + const int col = (u & 7) + 8 * (u >> 7); + const int g = g0 + ((u >> 3) & 15); + const int co = co0 + col; + const int c = g >> 1, v = g & 1; + const int4 a = sm[col][((u >> 3) & 15) * 3]; + const int4 e = sm[col][((u >> 3) & 15) * 3 + 1]; + const int4 f = sm[col][((u >> 3) & 15) * 3 + 2]; + uint32_t h[12]; + h[0] = a.x; h[1] = a.y; h[2] = a.z; h[3] = a.w; + h[4] = e.x; h[5] = e.y; h[6] = e.z; h[7] = e.w; + h[8] = f.x; h[9] = f.y; h[10] = f.z; h[11] = f.w; + __half *d = dbase + (size_t)c * (48 * Cout) + + (size_t)(co >> 3) * 128 + v * 64 + (co & 7) * 8; +#pragma unroll + for (int tp = 0; tp < 3; ++tp) { + uint32_t o[4]; +#pragma unroll + for (int jj = 0; jj < 4; ++jj) { + const uint32_t g0w = h[3 * jj], g1w = h[3 * jj + 1], + g2w = h[3 * jj + 2]; + o[jj] = (tp == 0) ? __byte_perm(g0w, g1w, 0x7610) + : (tp == 1) ? __byte_perm(g0w, g2w, 0x5432) + : __byte_perm(g1w, g2w, 0x7610); + } + *reinterpret_cast(d + (size_t)tp * (16 * Cout)) = + make_int4(o[0], o[1], o[2], o[3]); + } + } + __syncthreads(); + } +} diff --git a/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/tail_kernel.inc b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/tail_kernel.inc new file mode 100644 index 0000000..054ee0f --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/tail_kernel.inc @@ -0,0 +1,502 @@ +#define TL_NT 512 +#define TL_MW 512 +#define TL_MLD TL_MW +#define TL_VW 504 +#define TL_UW 520 +#define TL_BN 496 +#define TL_VLD 520 +#define TL_LBO_U (TL_UW * 8) +#define TL_LBO_M (TL_MLD * 8) +#define TL_LBO_V (TL_VLD * 8) +#define TL_SBO_A 128 +#define TL_PU_E (4 * TL_UW * 8) +#define TL_PM_E (8 * TL_MLD * 8) +#define TL_PV_E (32 * TL_VLD) +#define TL_AI_E (14 * 1024) +#define TL_AS_E (32 * 512 + 512) +#define TL_SHARED ((2 * TL_PU_E + TL_PM_E + TL_PV_E + TL_AI_E + TL_AS_E) * 2) +__device__ __forceinline__ void tl_stmatrix4(void *dst, const uint32_t r[4], + bool trans) { + const uint32_t a = static_cast(__cvta_generic_to_shared(dst)); + if (trans) + asm volatile( + "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1,%2,%3,%4};\n" ::"r"( + a), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3])); + else + asm volatile( + "stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1,%2,%3,%4};\n" ::"r"( + a), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3])); +} +__device__ __forceinline__ void tl_ldmatrix4t(uint32_t r[4], const void *src) { + const uint32_t a = static_cast(__cvta_generic_to_shared(src)); + asm volatile( + "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(a)); +} +__device__ __forceinline__ void tl_stmatrix1t(void *dst, uint32_t r) { + const uint32_t a = static_cast(__cvta_generic_to_shared(dst)); + asm volatile( + "stmatrix.sync.aligned.m8n8.x1.trans.shared.b16 [%0], {%1};\n" :: + "r"(a), "r"(r)); +} +template +__device__ __forceinline__ void tl_skip_band( + const __half *PU, __half *PM, const __half *AS, + __half *PV, int c0, int warp, int lane, int er0, const float *sh_ask, + const float *sh_ari_v, float b0v, float b1v) { + constexpr int NB8 = BW / 8; + float acc[BW / 2]; +#pragma unroll + for (int i = 0; i < BW / 2; ++i) acc[i] = ((i & 2) == 0) ? b0v : b1v; + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int c = 0; c < 4; ++c) { + const uint64_t da0 = + wg_desc(AS + c * 8 * 512, 128u, TL_SBO_A * 2u); + const uint64_t db0 = + wg_desc(PM + (2 * c) * TL_LBO_M + (c0 + 2) * 8, + TL_LBO_M * 2u, 128u); +#pragma unroll + for (int p = 0; p < 4; ++p) { + const uint64_t da = da0 + 64 * (2 * p); + const uint64_t db = db0 + 2 * p; + if constexpr (BW == 120) { + WGMMA_M64N120K16_T(acc, da, db, 0, 0); + } else { + WGMMA_M64N128K16_T(acc, da, db, 0, 0); + } + } + } + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + const int m_addr = lane >> 3, row = lane & 7; + const int cg = (warp & 1) * 2; + if (er0 >= 32) { +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = __floats2half2_rn( + acc[jj * 4 + h * 2 + 0], acc[jj * 4 + h * 2 + 1]); + r[m] = *reinterpret_cast(&hv); + } else { + r[m] = 0; + } + } + const int jj = j + (m_addr >> 1), h = m_addr & 1; + tl_stmatrix4(PV + (cg + h) * TL_LBO_V + (c0 + jj * 8 + row) * 8, + r, true); + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + if (er0 < 32) { +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; + const int jja = j + (m_addr >> 1), ha = m_addr & 1; + tl_ldmatrix4t(r, PV + (cg + ha) * TL_LBO_V + + (c0 + jja * 8 + row + 1) * 8); +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = *reinterpret_cast<__half2 *>(&r[m]); + acc[jj * 4 + h * 2 + 0] += __half2float(hv.x); + acc[jj * 4 + h * 2 + 1] += __half2float(hv.y); + } + } + } +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; + const int jja = j + (m_addr >> 1), ha = m_addr & 1; + tl_ldmatrix4t(r, PU + (cg + ha) * TL_LBO_U + + (c0 + jja * 8 + row + 16) * 8); +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = *reinterpret_cast<__half2 *>(&r[m]); + acc[jj * 4 + h * 2 + 0] += __half2float(hv.x); + acc[jj * 4 + h * 2 + 1] += __half2float(hv.y); + } + } + } + } + if (er0 >= 32) return; + if (er0 < 16) { + const float a0 = sh_ask[er0], r0 = sh_ari_v[er0]; + const float a1 = sh_ask[er0 + 8], r1 = sh_ari_v[er0 + 8]; +#pragma unroll + for (int i = 0; i < BW / 2; ++i) { + const int h = (i & 2) >> 1; + const float v = acc[i]; + const float sn = __sinf((h ? a1 : a0) * v); + acc[i] = __fmaf_rn(sn * sn, h ? r1 : r0, v); + } + } else { +#pragma unroll + for (int i = 0; i < BW / 2; ++i) { + const float v = acc[i]; + acc[i] = v > 0.f ? v : 0.01f * v; + } + } +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = __floats2half2_rn( + acc[jj * 4 + h * 2 + 0], acc[jj * 4 + h * 2 + 1]); + r[m] = *reinterpret_cast(&hv); + } else { + r[m] = 0; + } + } + const int jj = j + (m_addr >> 1), h = m_addr & 1; + tl_stmatrix4(PM + (cg + h) * TL_LBO_V + + (c0 + jj * 8 + row) * 8, + r, true); + } +} +__device__ __noinline__ void tail_fused_kernel( + const __half *__restrict__ U, const __half *__restrict__ win, + const __half *__restrict__ wsk, const __half *__restrict__ wpo, + const float *__restrict__ bin, const float *__restrict__ ain, + const float *__restrict__ bsk, const float *__restrict__ ask, + const float *__restrict__ bpo, const __half *__restrict__ st_m4, + const __half *__restrict__ st_v, __half *__restrict__ wave, + __half *__restrict__ ns_u, __half *__restrict__ ns_m4, + __half *__restrict__ ns_v, int B, int T4, int L4, int SC4, int N4, + int NTILE) { + char *tl_arena=reinterpret_cast(native_arena()); + __half *const PU = reinterpret_cast<__half *>(tl_arena); + __half *const PM = PU + 2 * TL_PU_E; + __half *const PV = PM + TL_PM_E; + __half *const AI = PV + TL_PV_E; + __half *const AS = AI + TL_AI_E; + __half *const PB = AS + 32 * 512; + float *sh_ain=(float*)native_aux(), *sh_ari=sh_ain+32, *sh_ask=sh_ari+32, *sh_ari_v=sh_ask+16; + const int tid = threadIdx.x; + const int wg = tid >> 7; + const int warp = (tid & 127) >> 5; + const int lane = tid & 31; + const int er0 = warp * 16 + (lane >> 2); + const int cb2 = (lane & 3) * 2; + for (int i = tid; i < 32; i += TL_NT) { + const float a = ain[i]; + sh_ain[i] = a; + sh_ari[i] = __frcp_rn(a + 1e-9f); + } + for (int i = tid; i < 16; i += TL_NT) { + const float a = ask[i]; + sh_ask[i] = a; + sh_ari_v[i] = __frcp_rn(a + 1e-9f); + } + for (int i = tid; i < TL_AI_E / 8; i += TL_NT) { + const int blk = i / 128, off = (i - blk * 128) * 8; + const int c = blk / 7, k = blk - c * 7; + *reinterpret_cast(AI + blk * 1024 + off) = + *reinterpret_cast(win + (size_t)c * 14336 + + (size_t)k * 2048 + off); + } + for (int i = tid; i < 32 * 64; i += TL_NT) { + const int blk = i / 64, off = (i - blk * 64) * 8; + const int c = blk >> 3, k = blk & 7; + *reinterpret_cast(AS + blk * 512 + off) = + (k < 7) ? *reinterpret_cast( + wsk + (size_t)c * 14336 + (size_t)k * 2048 + off) + : make_int4(0, 0, 0, 0); + } + for (int i = tid; i < 4 * (TL_VLD - TL_VW) * 8; i += TL_NT) { + const int g = i / ((TL_VLD - TL_VW) * 8); + PV[g * TL_LBO_V + TL_VW * 8 + (i - g * ((TL_VLD - TL_VW) * 8))] = + __float2half(0.f); + } + for (int i = tid; i < 512; i += TL_NT) PB[i] = __float2half(0.f); + for (int i = tid; i < 2 * 16 * 8; i += TL_NT) { + const int c = i / (16 * 8); + const int r = i - c * 16 * 8; + const int ki = r / 8, k = r & 7; + PB[c * 128 + (ki >> 3) * 64 + (ki & 7) * 8 + k] = + (k < 7) ? wpo[(c * 16 + ki) * 7 + k] : __float2half(0.f); + } + const float bias_i0 = bin[er0], bias_i1 = bin[er0 + 8]; + const float bias_s0 = (er0 < 32) ? bsk[er0] : 0.f; + const float bias_s1 = (er0 + 8 < 32) ? bsk[er0 + 8] : 0.f; + const float bpo0 = bpo[0]; + constexpr int UCOL = TL_UW / 32; + constexpr int UUNIT = 4 * UCOL; + constexpr int UWARP = TL_NT / 32; + constexpr int UITER = (UUNIT + UWARP - 1) / UWARP; + const int uwarp = tid >> 5; + const int uq = lane & 3; + int ugc[UITER], uoc[UITER]; +#pragma unroll + for (int t = 0; t < UITER; ++t) { + const int bi = uwarp + t * UWARP; + const int gc = bi / UCOL; + ugc[t] = gc; + uoc[t] = bi - gc * UCOL; + } + const __half *const urow = + U + (size_t)(lane >> 2) * SC4 + 8 * uq; + __syncthreads(); + PDL_ENTER(); +#define TL_STAGE_U(PUD,UB) { _Pragma("unroll") for (int t2 = 0; t2 < UITER; t2 += 2) { uint32_t u[2][4]; _Pragma("unroll") for (int d = 0; d < 2; ++d) { const int t = t2 + d; const int n = (UB) + uoc[t] * 32; const int gn = n + 8 * uq; const bool ok = gn >= 0 && gn < N4; int4 v = make_int4(0, 0, 0, 0); if (ok) v = *reinterpret_cast( urow + (size_t)ugc[t] * (8 * (size_t)SC4) + n); u[d][0] = v.x; u[d][1] = v.y; u[d][2] = v.z; u[d][3] = v.w; } _Pragma("unroll") for (int d = 0; d < 2; ++d) { const int t = t2 + d; uint32_t *q = u[d]; uint32_t x, y; x = (uq & 1) ? q[0] : q[1]; y = (uq & 1) ? q[2] : q[3]; x = __shfl_xor_sync(0xffffffffu, x, 1); y = __shfl_xor_sync(0xffffffffu, y, 1); if (uq & 1) { q[0] = x; q[2] = y; } else { q[1] = x; q[3] = y; } x = (uq & 2) ? q[0] : q[2]; y = (uq & 2) ? q[1] : q[3]; x = __shfl_xor_sync(0xffffffffu, x, 2); y = __shfl_xor_sync(0xffffffffu, y, 2); if (uq & 2) { q[0] = x; q[1] = y; } else { q[2] = x; q[3] = y; } tl_stmatrix4((PUD) + ugc[t] * TL_LBO_U + (uoc[t] * 32 + lane) * 8, q, true); } } if (uwarp < 4) { const int n = (UB) + 512; const int gn = n + 2 * (lane & 3); uint32_t q = 0; if (gn >= 0 && gn + 1 < N4) q = *reinterpret_cast( U + (size_t)(uwarp * 8 + (lane >> 2)) * SC4 + gn); tl_stmatrix1t((PUD) + uwarp * TL_LBO_U + (512 + (lane & 7)) * 8, q); } } + TL_STAGE_U(PU, (int)blockIdx.x * TL_BN - 16) + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + int ubuf = 0; + for (int tile = blockIdx.x; tile < NTILE; tile += gridDim.x, ubuf ^= 1) { + const int g0 = tile * TL_BN; + bool active_tile = false; + for (int row = max(0, g0 - 16) / L4; + row <= min(N4 - 1, g0 - 16 + TL_UW - 1) / L4; ++row) { + if (native_meta()->lengths[row] > 0) { + active_tile = true; + break; + } + } + __half *const PUc = PU + ubuf * TL_PU_E; + __half *const PUn = PU + (ubuf ^ 1) * TL_PU_E; + { + const int b0 = g0 / L4; + for (int s = 0; s < 2; ++s) { + const int sb = b0 + s; + if (sb >= B) break; + const int cv = (int)((long long)sb * L4 + native_true_tokens(sb,T4) - g0); + if (cv < 0 || cv >= TL_BN) continue; + for (int i = tid; i < 32 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + ns_u[((size_t)sb * 32 + co) * 6 + h] = + PUc[(co >> 3) * TL_LBO_U + + (cv + h + 16) * 8 + (co & 7)]; + } + } + } + if (active_tile) { + float acc[64]; +#pragma unroll + for (int i = 0; i < 64; ++i) + acc[i] = ((i & 2) == 0) ? bias_i0 : bias_i1; + const int c0 = wg * 128; + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int c = 0; c < 2; ++c) { + const uint64_t da0 = + wg_desc(AI + c * 7 * 1024, 128u, TL_SBO_A * 2u); + const uint64_t db0 = + wg_desc(PUc + (2 * c) * TL_LBO_U + (c0 + 2) * 8, + TL_LBO_U * 2u, 128u); +#pragma unroll + for (int k = 0; k < 7; ++k) { + const uint64_t da = da0 + 128 * k; + const uint64_t db = db0 + k; + WGMMA_M64N128K16_T(acc, da, db, 0, 0); + } + } + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); + if (tile + (int)gridDim.x < NTILE) + TL_STAGE_U(PUn, (tile + (int)gridDim.x) * TL_BN - 16) + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + if (er0 < 32) { + const float a0 = sh_ain[er0], r0 = sh_ari[er0]; + const float a1 = sh_ain[er0 + 8], r1 = sh_ari[er0 + 8]; +#pragma unroll + for (int i = 0; i < 64; ++i) { + const int h = (i & 2) >> 1; + const float v = acc[i]; + const float sn = __sinf((h ? a1 : a0) * v); + acc[i] = __fmaf_rn(sn * sn, h ? r1 : r0, v); + } + } else { +#pragma unroll + for (int i = 0; i < 64; ++i) { + const float v = acc[i]; + acc[i] = v > 0.f ? v : 0.01f * v; + } + } +#pragma unroll + for (int j = 0; j < 16; j += 2) { + uint32_t r[4]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + const __half2 hv = __floats2half2_rn( + acc[jj * 4 + h * 2 + 0], acc[jj * 4 + h * 2 + 1]); + r[m] = *reinterpret_cast(&hv); + } + const int m_addr = lane >> 3, row = lane & 7; + const int jj = j + (m_addr >> 1), h = m_addr & 1; + tl_stmatrix4(PM + (warp * 2 + h) * TL_LBO_M + + (c0 + jj * 8 + row) * 8, + r, true); + } + } else if (tile + (int)gridDim.x < NTILE) { + TL_STAGE_U(PUn, (tile + (int)gridDim.x) * TL_BN - 16) + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + { + const int base = g0 - 8; + const int b0 = (base >= 0) ? base / L4 : 0; + const int cm0 = (int)((long long)b0 * L4 - base); + const int cm1 = cm0 + L4; + const bool hit0 = (b0 < B) && (cm0 + 6 > 0) && (cm0 < TL_MW); + const bool hit1 = (b0 + 1 < B) && (cm1 + 6 > 0) && (cm1 < TL_MW); + if (hit0 || hit1) { + for (int s = 0; s < 2; ++s) { + if (!(s ? hit1 : hit0)) continue; + const int sb = b0 + s, cm = s ? cm1 : cm0; + for (int i = tid; i < 64 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + const int c = cm + h; + if (c >= 0 && c < TL_MW) + PM[(co >> 3) * TL_LBO_M + c * 8 + (co & 7)] = + st_m4[((size_t)sb * 64 + co) * 6 + h]; + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + } + } + { + const int b0 = g0 / L4; + for (int s = 0; s < 2; ++s) { + const int sb = b0 + s; + if (sb >= B) break; + const int cv = (int)((long long)sb * L4 + native_true_tokens(sb,T4) - g0); + if (cv < 0 || cv >= TL_BN) continue; + for (int i = tid; i < 64 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + ns_m4[((size_t)sb * 64 + co) * 6 + h] = + PM[(co >> 3) * TL_LBO_M + + (cv + h + 8) * 8 + (co & 7)]; + } + } + } + if (active_tile) { + const int c0 = wg * 128; + if (wg == 3) + tl_skip_band<120>(PUc, PM, AS, PV, c0, warp, lane, er0, sh_ask, + sh_ari_v, bias_s0, bias_s1); + else + tl_skip_band<128>(PUc, PM, AS, PV, c0, warp, lane, er0, sh_ask, + sh_ari_v, bias_s0, bias_s1); + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + { + const int b0 = g0 / L4; + const int cv0 = (int)((long long)b0 * L4 - g0), cv1 = cv0 + L4; + const bool hv0 = (b0 < B) && (cv0 + 6 > 0) && (cv0 < TL_VW); + const bool hv1 = (b0 + 1 < B) && (cv1 + 6 > 0) && (cv1 < TL_VW); + if (hv0 || hv1) { + for (int s = 0; s < 2; ++s) { + if (!(s ? hv1 : hv0)) continue; + const int sb = b0 + s, cv = s ? cv1 : cv0; + for (int i = tid; i < 32 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + const int c = cv + h; + if (c >= 0 && c < TL_VW) + PM[(co >> 3) * TL_LBO_V + c * 8 + (co & 7)] = + st_v[((size_t)sb * 32 + co) * 6 + h]; + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + } + } + { + const int b0 = g0 / L4; + for (int s = 0; s < 2; ++s) { + const int sb = b0 + s; + if (sb >= B) break; + const int cv = (int)((long long)sb * L4 + native_true_tokens(sb,T4) - g0); + if (cv < 0 || cv >= TL_BN) continue; + for (int i = tid; i < 32 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + ns_v[((size_t)sb * 32 + co) * 6 + h] = + PM[(co >> 3) * TL_LBO_V + (cv + h) * 8 + (co & 7)]; + } + } + } + __syncthreads(); + constexpr int PC_LD = 10; + if (active_tile) { + float *const PC = reinterpret_cast(PV); +#pragma unroll + for (int round = 0; round < 2; ++round) { + float acc[4]; + const int p0 = round * 256 + wg * 64; + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); + const uint64_t da0 = + wg_desc(PM + p0 * 8, TL_LBO_V * 2u, 128u); + const uint64_t db0 = wg_desc(PB, 128u, 128u); + const uint64_t da1 = + wg_desc(PM + 2 * TL_LBO_V + p0 * 8, + TL_LBO_V * 2u, 128u); + const uint64_t db1 = wg_desc(PB + 128, 128u, 128u); + WGMMA_M64N8K16_ZERO_T(acc, da0, db0, 0, 1); + WGMMA_M64N8K16_T(acc, da1, db1, 0, 1); + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + const int row0 = p0 + warp * 16 + (lane >> 2); + const int col = (lane & 3) * 2; + *reinterpret_cast(PC + row0 * PC_LD + col) = + make_float2(acc[0], acc[1]); + *reinterpret_cast(PC + (row0 + 8) * PC_LD + col) = + make_float2(acc[2], acc[3]); + } + } + __syncthreads(); + if (active_tile) { + const float *const PC = reinterpret_cast(PV); + const int p0 = tid; + const bool live = p0 < TL_BN; + const unsigned live_mask = __ballot_sync(0xffffffffu, live); + if (live) { + const float s0 = __fadd_rn(PC[(p0 + 0) * PC_LD + 0], + PC[(p0 + 1) * PC_LD + 1]); + const float s1 = __fadd_rn(PC[(p0 + 2) * PC_LD + 2], + PC[(p0 + 3) * PC_LD + 3]); + const float s2 = __fadd_rn(PC[(p0 + 4) * PC_LD + 4], + PC[(p0 + 5) * PC_LD + 5]); + const float s3 = PC[(p0 + 6) * PC_LD + 6]; + float y0 = __fadd_rn(__fadd_rn(s0, s1), + __fadd_rn(s2, s3)); + y0 = fminf(fmaxf(y0 + bpo0, -1.f), 1.f); + const float y1 = __shfl_down_sync(live_mask, y0, 1); + if ((p0 & 1) == 0) { + const int b0 = g0 / L4; + const int rr = g0 - b0 * L4; + int t0 = rr + p0, bo0 = b0; + int t1 = t0 + 1, bo1 = bo0; + if (t0 >= L4) { t0 -= L4; ++bo0; } + if (t1 >= L4) { t1 -= L4; ++bo1; } + if (bo0 == bo1 && bo0 < B && t1 < T4) { + *reinterpret_cast<__half2 *>( + wave + (size_t)bo0 * T4 + t0) = + __floats2half2_rn(y0, y1); + } else { + if (bo0 < B && t0 < T4) + wave[(size_t)bo0 * T4 + t0] = __float2half(y0); + if (bo1 < B && t1 < T4) + wave[(size_t)bo1 * T4 + t1] = __float2half(y1); + } + } + } + } + } +#undef TL_STAGE_U +} diff --git a/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wg2_kernel.inc b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wg2_kernel.inc new file mode 100644 index 0000000..4ffa7ad --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wg2_kernel.inc @@ -0,0 +1,280 @@ +__device__ __forceinline__ void cpa16(void *dsm, const void *src, bool pred) { + unsigned a = static_cast(__cvta_generic_to_shared(dsm)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(a), + "l"(src), "r"(pred ? 16 : 0) + : "memory"); +} +template +__device__ __forceinline__ void cpa16o(unsigned a, const void *src, bool pred) { + asm volatile("cp.async.cg.shared.global [%0+%3], [%1], 16, %2;\n" ::"r"(a), + "l"(src), "r"(pred ? 16 : 0), "n"(OFF) + : "memory"); +} +template +__device__ __noinline__ +void wg2_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int BK = CIB * KT; + constexpr int WG_BM = 64 * MSPLIT; + constexpr int WG_BNT = WG_BN * NSPLIT; + constexpr int NTS = 128 * MSPLIT * NSPLIT; + constexpr int WG_NT = NTS * KSPLIT; + constexpr int NCRAW = WG_BNT / 4; + constexpr int EPN = + (WG_NT % NCRAW == 0) ? WG_BNT : (((WG_BNT + 127) / 128) * 128); + constexpr int WG_CLD = EPN + 8; + constexpr bool NB8_P2 = ((WG_BN / 8) & (WG_BN / 8 - 1)) == 0; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int ST_E = SA * A_E + SB * NSPLIT * B_E; + constexpr int SBO_AE = KG * WG_SLOT; + constexpr int LBO_BE = (WG_BN / 8) * WG_SLOT; + constexpr int NB8 = WG_BN / 8; + constexpr bool ALDG = (WG_BN <= 64); + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *smem = reinterpret_cast<__half *>(arena); + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int kgrp = tid / NTS; + const int stid = tid - kgrp * NTS; + const int wgl = stid >> 7; + const int mg = wgl % MSPLIT; + const int ngid = wgl / MSPLIT; + const int warp = (stid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * WG_BM; + const int nb = blockIdx.y * WG_BNT; + const int n0 = nb + ngid * WG_BN; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + const int NCH = (Kg + BK - 1) / BK; + if (tid < WG_BM) { + const int co = m0 + tid; + if (do_snake && co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + for (int i = tid; i < EPN; i += WG_NT) { + const int n = nb + i; + colok[i] = (i < WG_BNT && n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mg * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (kgrp == 0 && eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (kgrp == 0 && eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + constexpr int KGX = ALDG ? (KG <= 8 ? 8 : (KG <= 16 ? 16 : 32)) : KG; + constexpr int ATOT = WG_BM * KGX; + constexpr int AREG = (ATOT + NTS - 1) / NTS; + constexpr int BTOT = NSPLIT * CIB * NB8; + constexpr int BREG = (BTOT + NTS - 1) / NTS; + const __half *aptr[AREG]; + int aoff[AREG]; + int akmax[AREG]; +#pragma unroll + for (int r = 0; r < AREG; ++r) { + const int idx = stid + r * NTS; + const int m = idx / KGX, v = idx - (idx / KGX) * KGX; + const int co = m0 + m; + const bool live = (ATOT % NTS == 0 || idx < ATOT) && (co < Cout) && + v < KG; + aoff[r] = (m >> 3) * SBO_AE + (v < KG ? v : 0) * WG_SLOT + (m & 7) * 8; + akmax[r] = live ? (Kg - v * 8) : 0; + aptr[r] = w + (size_t)(live ? co : 0) * Kg + v * 8; + } + const __half *bptr[BREG]; + int bcimax[BREG], bnc[BREG], bcil[BREG], bgrp[BREG]; +#pragma unroll + for (int r = 0; r < BREG; ++r) { + const int idx = stid + r * NTS; + const int gg = idx / (CIB * NB8); + const int rr = idx - gg * (CIB * NB8); + const int cil = rr / NB8, nc = rr - (rr / NB8) * NB8; + bcil[r] = cil; + bnc[r] = nc; + bgrp[r] = gg; + const bool live = (BTOT % NTS == 0 || idx < BTOT); + bcimax[r] = live ? (Cin - cil) : 0; + bptr[r] = in + (size_t)cil * SCin + nb + gg * WG_BN + nc * 8; + } + int4 br0[BREG], br1[BREG]; + int4 ar[ALDG ? AREG : 1]; + uint32_t edge[BREG]; + __half *const kaB = smem + kgrp * ST_E; + __half *const kbB = kaB + SA * A_E; + const unsigned barid = 1u + kgrp + (unsigned)(KSPLIT * native_subid()); +#define WG2_BAR() asm volatile("bar.sync %0, %1;\n" ::"r"(barid), "r"(NTS)) +#define WG2_LOADB(c) { const int _ci0 = (c) * CIB; _Pragma("unroll") for (int r = 0; r < BREG; ++r) { br0[r] = make_int4(0, 0, 0, 0); if constexpr (KT == 3) edge[r] = 0; else br1[r] = make_int4(0, 0, 0, 0); if (_ci0 < bcimax[r]) { const int4 *src = reinterpret_cast( bptr[r] + (size_t)_ci0 * SCin); br0[r] = src[0]; if constexpr (KT == 3) { if (!NB8_P2 || bnc[r] == NB8 - 1) edge[r] = src[1].x; } else { br1[r] = src[1]; } } } } +#define WG2_LOADA(c) if constexpr (ALDG) { const int _kg0 = (c) * BK; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { ar[r] = make_int4(0, 0, 0, 0); if (_kg0 < akmax[r]) ar[r] = *reinterpret_cast(aptr[r] + _kg0); } } +#define WG2_PUTA(c,AD) { _Pragma("unroll") for (int r = 0; r < AREG; ++r) { if constexpr (ALDG) { if (akmax[r] > 0) *reinterpret_cast((AD) + aoff[r]) = ar[r]; } else { const int _kg0 = (c) * BK; if (akmax[r] > 0) cpa16((AD) + aoff[r], aptr[r] + _kg0, _kg0 < akmax[r]); } } if constexpr (!ALDG) asm volatile("cp.async.commit_group;\n" ::: "memory"); } +#define WG2_STSB(BD) { _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int cil = bcil[r], nc = bnc[r]; unsigned int u[8]; u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; if constexpr (KT == 3) { if constexpr (NB8_P2) { u[4] = __shfl_down_sync(0xffffffffu, u[0], 1, NB8); if (nc == NB8 - 1) u[4] = edge[r]; } else { u[4] = edge[r]; } u[5] = u[6] = u[7] = 0; } else { u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; } if (BTOT % NTS == 0 || (stid + r * NTS) < BTOT) { _Pragma("unroll") for (int k = 0; k < KT; ++k) { const int kh = k >> 1; int4 o; if (k & 1) { o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); } else { o.x = u[kh + 0]; o.y = u[kh + 1]; o.z = u[kh + 2]; o.w = u[kh + 3]; } const int row = cil * KT + k; *reinterpret_cast((BD) + bgrp[r] * B_E + nc * WG_SLOT + (row >> 3) * LBO_BE + (row & 7) * 8) = o; } } } } +#define WG2_MMA(AD,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t da0 = wg_desc((AD) + mg * 8 * SBO_AE, WG_SLOT * 2, SBO_AE * 2); const uint64_t db0 = wg_desc((BD) + ngid * B_E, LBO_BE * 2, WG_SLOT * 2); _Pragma("unroll") for (int kk = 0; kk < BK; kk += 16) { const uint64_t da = da0 + (kk >> 3) * (WG_SLOT / 8); const uint64_t db = db0 + (kk >> 3) * (LBO_BE / 8); if constexpr (WG_BN == 64) { WGMMA_M64N64K16(acc, da, db); } else if constexpr (WG_BN == 256) { WGMMA_M64N256K16(acc, da, db); } else if constexpr (WG_BN == 104) { WGMMA_M64N104K16(acc, da, db); } else if constexpr (WG_BN == 112) { WGMMA_M64N112K16(acc, da, db); } else if constexpr (WG_BN == 120) { WGMMA_M64N120K16(acc, da, db); } else if constexpr (WG_BN == 224) { WGMMA_M64N224K16(acc, da, db); } else { static_assert(WG_BN == 128, "no wgmma macro for WG_BN"); WGMMA_M64N128K16(acc, da, db); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + constexpr int SMAX = SA > SB ? SA : SB; + WG2_LOADA(kgrp) + WG2_LOADB(kgrp) +#pragma unroll 1 + for (int s = 0; s < SMAX; ++s) { + const int c = s * KSPLIT + kgrp; + if (s < SA) { + if (c < NCH) WG2_PUTA(c, kaB + s * A_E) + if (c + KSPLIT < NCH) WG2_LOADA(c + KSPLIT) + } + if (s < SB) { + if (c < NCH) WG2_STSB(kbB + s * NSPLIT * B_E) + if (c + KSPLIT < NCH) WG2_LOADB(c + KSPLIT) + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + const int NIT = (NCH - kgrp + KSPLIT - 1) / KSPLIT; + for (int i = 0; i < NIT; ++i) { + if constexpr (!ALDG) + asm volatile("cp.async.wait_group %0;\n" ::"n"(SA - 2) : "memory"); + WG2_BAR(); + WG2_MMA(kaB + (i % SA) * A_E, kbB + (i % SB) * NSPLIT * B_E) + asm volatile("wgmma.wait_group.sync.aligned 1;\n" ::: "memory"); + if constexpr (NTS > 128) WG2_BAR(); + if (i >= 1) { + const int ca = (i + SA - 1) * KSPLIT + kgrp; + if (ca < NCH) WG2_PUTA(ca, kaB + ((i - 1) % SA) * A_E) + if (ALDG && ca + KSPLIT < NCH) WG2_LOADA(ca + KSPLIT) + const int cb = (i + SB - 1) * KSPLIT + kgrp; + if (cb < NCH) { + WG2_STSB(kbB + ((i - 1) % SB) * NSPLIT * B_E) + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if (cb + KSPLIT < NCH) WG2_LOADB(cb + KSPLIT) + } else if constexpr (ALDG) { + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + } + } + } + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); +#undef WG2_LOADB +#undef WG2_LOADA +#undef WG2_PUTA +#undef WG2_STSB +#undef WG2_MMA +#undef WG2_BAR + __syncthreads(); + float *Ck = Cs + kgrp * WG_BM * WG_CLD; + const int cb = (lane & 3) * 2; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Ck + (er0 + h * 8) * WG_CLD + + ngid * WG_BN + j * 8 + cb) = v; + } + } + __syncthreads(); + constexpr int NC = EPN / 4; + constexpr int MSTEP = WG_NT / NC; + constexpr int MITER = WG_BM / MSTEP; + constexpr int UNR = (MITER % 4 == 0) ? 4 : ((MITER % 2 == 0) ? 2 : 1); + static_assert(MITER % UNR == 0, "row batch must divide the row count"); + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + const int n = nb + lc; +#pragma unroll 1 + for (int t0 = 0; t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + v[u] = *reinterpret_cast(Cs + m * WG_CLD + lc); + if (resid != nullptr) { + const int co = m0 + m; + const int cc = co < Cout ? co : 0; + const __half2 *rp = reinterpret_cast( + resid + (size_t)cc * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + const int co = m0 + m; + if (co >= Cout) continue; + float4 vv = v[u]; +#pragma unroll + for (int g = 1; g < KSPLIT; ++g) { + float4 p = *reinterpret_cast( + Cs + g * WG_BM * WG_CLD + m * WG_CLD + lc); + vv.x += p.x; vv.y += p.y; vv.z += p.z; vv.w += p.w; + } + if (resid != nullptr) { + vv.x += __half2float(r0[u].x); vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (do_snake) { + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + } + __half *op = out + (size_t)co * SCout + n + H; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B * H; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int total = Cout * B * pad_len; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} diff --git a/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wg3_kernel.inc b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wg3_kernel.inc new file mode 100644 index 0000000..fa55b32 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wg3_kernel.inc @@ -0,0 +1,177 @@ +template +__device__ __noinline__ void wg3_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int WG_BNT = WG_BN * NSG; + constexpr int WG_CLD = WG_BNT + 8; + constexpr int WG_NT = 128 * NWG * NSG; + constexpr int WG_BM = 64 * NWG; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int SBO_AE = KG * WG_SLOT; + constexpr int LBO_BE = (WG_BN / 8) * WG_SLOT; + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + constexpr int ST_E = A_E + NSG * B_E; + __half *As0 = reinterpret_cast<__half *>(arena); + __half *Bs0 = As0 + A_E; + __half *As1 = (STAGES == 2) ? As0 + ST_E : As0; + __half *Bs1 = (STAGES == 2) ? Bs0 + ST_E : Bs0; + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int wgid = tid >> 7; + const int mgid = wgid % NWG, ngid = wgid / NWG; + const int warp = (tid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * WG_BM; + constexpr int WGROWS = 8; + const int nb = blockIdx.y * WG_BNT; + const int n0 = nb + ngid * WG_BN; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + if (tid < WG_BM) { + const int co = m0 + tid; + if (do_snake && co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + if (tid < WG_BNT) { + const int n = nb + tid; + colok[tid] = (n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mgid * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + constexpr int ATOT = WG_BM * KG; + constexpr int AREG = (ATOT + WG_NT - 1) / WG_NT; + constexpr int CIB = BK / KT; + constexpr int BTOT = NSG * CIB * (WG_BN / 8); + constexpr int BREG = (BTOT + WG_NT - 1) / WG_NT; + int4 ar[AREG]; + int4 br0[BREG], br1[BREG]; + uint32_t edge[BREG]; +#define WG_LOADG(kg0) { const int _ci0 = (kg0) / KT; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; ar[r] = make_int4(0, 0, 0, 0); if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; const int co = m0 + m, kg = (kg0) + v * 8; if (co < Cout && kg < Kg) ar[r] = *reinterpret_cast( w + (size_t)co * Kg + kg); } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; br0[r] = make_int4(0, 0, 0, 0); if constexpr (KT == 3) edge[r] = 0; else br1[r] = make_int4(0, 0, 0, 0); if (BTOT % WG_NT == 0 || idx < BTOT) { const int _g = idx / (CIB * (WG_BN / 8)); const int _r = idx - _g * (CIB * (WG_BN / 8)); const int ci = _ci0 + _r / (WG_BN / 8); const int nc = _r % (WG_BN / 8); if (ci < Cin) { const int4 *src = reinterpret_cast( in + (size_t)ci * SCin + nb + _g * WG_BN + nc * 8); br0[r] = src[0]; if constexpr (KT == 3) { if (nc == WG_BN / 8 - 1) edge[r] = src[1].x; } else { br1[r] = src[1]; } } } } } + WG_LOADG(0) +#define WG_STS(AD,BD) { _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; *reinterpret_cast( (AD) + (m >> 3) * SBO_AE + v * WG_SLOT + (m & 7) * 8) = ar[r]; } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; if (BTOT % WG_NT == 0 || idx < BTOT) { const int _g = idx / (CIB * (WG_BN / 8)); const int _r = idx - _g * (CIB * (WG_BN / 8)); const int cil = _r / (WG_BN / 8), nc = _r % (WG_BN / 8); unsigned int u[8]; u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; if constexpr (KT == 3) { u[4] = __shfl_down_sync(0xffffffffu, u[0], 1, WG_BN / 8); if (nc == WG_BN / 8 - 1) u[4] = edge[r]; u[5] = u[6] = u[7] = 0; } else { u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; } _Pragma("unroll") for (int k = 0; k < KT; ++k) { const int kh = k >> 1; int4 o; if (k & 1) { o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); } else { o.x = u[kh + 0]; o.y = u[kh + 1]; o.z = u[kh + 2]; o.w = u[kh + 3]; } const int row = cil * KT + k; *reinterpret_cast((BD) + _g * B_E + nc * WG_SLOT + (row >> 3) * LBO_BE + (row & 7) * 8) = o; } } } if constexpr (WG_BN == 64) { asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); } } +#define WG_MMA(AD,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t da0 = wg_desc((AD) + mgid * WGROWS * SBO_AE, WG_SLOT * 2, SBO_AE * 2); const uint64_t db0 = wg_desc((BD) + ngid * B_E, LBO_BE * 2, WG_SLOT * 2); _Pragma("unroll") for (int kk = 0; kk < BK; kk += 16) { const uint64_t da = da0 + (kk >> 3) * (WG_SLOT / 8); const uint64_t db = db0 + (kk >> 3) * (LBO_BE / 8); if constexpr (WG_BN == 64) { WGMMA_M64N64K16(acc, da, db); } else { WGMMA_M64N128K16(acc, da, db); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + WG_STS(As0, Bs0) + if (BK < Kg) WG_LOADG(BK) + __syncthreads(); + int cur = 0; + for (int kg0 = 0; kg0 < Kg; kg0 += BK) { + if (cur == 0) { WG_MMA(As0, Bs0) } else { WG_MMA(As1, Bs1) } + if (kg0 + BK < Kg) { + if (STAGES == 2) + asm volatile("wgmma.wait_group.sync.aligned 1;\n" ::: "memory"); + else + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + __syncthreads(); + if (cur == 0) { WG_STS(As1, Bs1) } else { WG_STS(As0, Bs0) } + if (kg0 + 2 * BK < Kg) WG_LOADG(kg0 + 2 * BK) + __syncthreads(); + } else { + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + } + if (STAGES == 2) cur ^= 1; + } +#undef WG_LOADG +#undef WG_STS +#undef WG_MMA + __syncthreads(); + const int r0 = er0; + const int cb = (lane & 3) * 2; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Cs + (r0 + h * 8) * WG_CLD + + ngid * WG_BN + j * 8 + cb) = v; + } + } + __syncthreads(); + constexpr int NC = WG_BNT / 4; + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + for (int m = lr; m < WG_BM; m += WG_NT / NC) { + const int co = m0 + m; + if (co >= Cout) continue; + const int n = nb + lc; + float4 v = *reinterpret_cast(Cs + m * WG_CLD + lc); + __half *op = out + (size_t)co * SCout + n + H; + if (resid != nullptr) { + const __half2 *rp = reinterpret_cast( + resid + (size_t)co * SCout + n + H); + __half2 r0h = rp[0], r1h = rp[1]; + v.x += __half2float(r0h.x); v.y += __half2float(r0h.y); + v.z += __half2float(r1h.x); v.w += __half2float(r1h.y); + } + if (do_snake) { + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + } + if (ok4) { + __half2 o0 = __floats2half2_rn(v.x, v.y); + __half2 o1 = __floats2half2_rn(v.z, v.w); + reinterpret_cast<__half2 *>(op)[0] = o0; + reinterpret_cast<__half2 *>(op)[1] = o1; + } else { + const float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B * H; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int total = Cout * B * pad_len; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} diff --git a/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wg4_step_widths.inc b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wg4_step_widths.inc new file mode 100644 index 0000000..2c9ec6f --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wg4_step_widths.inc @@ -0,0 +1,468 @@ +template +__device__ __noinline__ void wg4_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int CoutP, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + if constexpr (PDL && !LATE_PDL) PDL_ENTER(); + constexpr int NT = 128 * MSPLIT * NSPLIT; + constexpr int BM = 64 * MSPLIT; + constexpr int BNT = WG_BN * NSPLIT; + constexpr int NCRAW = BNT / 4; + constexpr int EPN = (NT % NCRAW == 0) ? BNT : (((BNT + 127) / 128) * 128); + constexpr int WG_CLD = EPN + 8; + constexpr int CG = CIB / 8; + constexpr int NBR = BNT + KT - 1; + constexpr int NBRP = BC8 ? ((NBR + 7) & ~7) : ((NBR + 31) & ~31); + constexpr int UCOL = NBRP / 32; + constexpr int NUNIT = CG * UCOL; + constexpr int NWARP = NT / 32; + constexpr int BITER = (NUNIT + NWARP - 1) / NWARP; + constexpr int B_E = CIB * NBRP; + constexpr int ST_E = B_E; + constexpr int LBO_B = NBRP * 8; + constexpr int ABUF = ABUF_; + static_assert(!OC8 || PREACT, + "K=7 C8 output is only used for residual-free pair heads"); + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *smem = reinterpret_cast<__half *>(arena); + __half *Hs = smem; + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int wgid = tid >> 7; + const int mg = wgid % MSPLIT, ng = wgid / MSPLIT; + const int warp = (tid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * BM; + const int nb = blockIdx.y * BNT; + const int H = KT - 1; + bool active_tile = false; + for (int row = nb / L; row <= min(N - 1, nb + BNT - 1) / L; ++row) { + if (native_meta()->lengths[row] > 0) { + active_tile = true; + break; + } + } + if (active_tile) { + const int half = Cout >> 1; + const int NCH = (Cin + CIB - 1) / CIB; + for (int i = tid; i < BM; i += NT) { + const int co = m0 + i; + if (co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[i] = a; + row_ri[i] = __frcp_rn(a + 1e-9f); + } + } + for (int i = tid; i < EPN; i += NT) { + const int n = nb + i; + colok[i] = (i < BNT && n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mg * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + const int mtile = blockIdx.x * MSPLIT + mg; + const __half *const apt = w + ((size_t)mtile * NCH * KT) * 1024 + + (size_t)(tid & 127) * 8; + uint32_t areg[ABUF][KT][4]; + const int wi = tid >> 5; + int bdst[BITER], boff[BITER]; + bool blive[BITER]; + const __half *const bbase = in + (size_t)(lane >> 2) * SCin + nb + + 8 * (lane & 3); +#pragma unroll + for (int t = 0; t < BITER; ++t) { + const int u = wi + t * NWARP; + const int gc = (u < NUNIT) ? u / UCOL : 0; + const int cu = (u < NUNIT) ? u - gc * UCOL : 0; + blive[t] = u < NUNIT; + bdst[t] = gc * LBO_B + (cu * 32 + lane) * 8; + boff[t] = gc * 8 * SCin + cu * 32; + } + const size_t SC8i = (size_t)SCin * 8; + constexpr int BQ = B_E / 8; + constexpr int BCREG = (BQ + NT - 1) / NT; + int cdst[BCREG]; + const __half *csrc[BCREG]; + bool clive[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) { + const int q = tid + r * NT; + const int gc = q / NBRP, rem = q - gc * NBRP; + clive[r] = ((BQ % NT == 0) || q < BQ) && gc < CG; + cdst[r] = clive[r] ? gc * LBO_B + rem * 8 : 0; + csrc[r] = in + (clive[r] ? (size_t)gc * SC8i + + (size_t)nb * 8 + rem * 8 : 0); + } + uint32_t br[BITER][4]; + constexpr int CPA_STAGES = 4; + constexpr int BSTAGES = BC8 ? CPA_STAGES : STAGES; + static_assert(STAGES >= CPA_STAGES, "C8 B-panel pipeline allocation"); + static_assert(ABUF == 2 && BSTAGES == 4, "wg4 unrolled pipeline shape"); + uint64_t dslot[BSTAGES]; +#pragma unroll + for (int s = 0; s < BSTAGES; ++s) + dslot[s] = wg_desc(smem + s * ST_E + ng * WG_BN * 8, LBO_B * 2u, 128u); + uint32_t csm0[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) + csm0[r] = static_cast(__cvta_generic_to_shared(smem + cdst[r])); + uint32_t bsm0[BITER]; +#pragma unroll + for (int t = 0; t < BITER; ++t) + bsm0[t] = static_cast(__cvta_generic_to_shared(smem + bdst[t])); + const size_t CSTEP = (size_t)CG * SC8i; + const __half *crol[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) crol[r] = csrc[r] + CPA_STAGES * CSTEP; + const __half *arol = apt + (size_t)ABUF * (KT * 1024); +#define WG4_CPB(c,BD) { const size_t _o = (size_t)(c) * CG * SC8i; _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16((BD) + cdst[r], csrc[r] + _o, true); if ((c) + 1 == CPA_STAGES) { _Pragma("unroll") for (int r = 0; r < BCREG; ++r) csrc[r] += (size_t)CPA_STAGES * CSTEP; } } +#define WG4_CPB_R(SLOT) { _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16o<(SLOT) * ST_E * 2>(csm0[r], csrc[r], true); _Pragma("unroll") for (int r = 0; r < BCREG; ++r) csrc[r] += CSTEP; } +#define WG4_LOADB(c) { const size_t _o = (size_t)(c) * CIB * SCin; _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { const int4 _v = *reinterpret_cast( bbase + _o + boff[t]); br[t][0] = _v.x; br[t][1] = _v.y; br[t][2] = _v.z; br[t][3] = _v.w; } } } +#define WG4_XPOSE(u) { const int _q = lane & 3; uint32_t _x, _y; _x = (_q & 1) ? u[0] : u[1]; _y = (_q & 1) ? u[2] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 1); _y = __shfl_xor_sync(0xffffffffu, _y, 1); if (_q & 1) { u[0] = _x; u[2] = _y; } else { u[1] = _x; u[3] = _y; } _x = (_q & 2) ? u[0] : u[2]; _y = (_q & 2) ? u[1] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 2); _y = __shfl_xor_sync(0xffffffffu, _y, 2); if (_q & 2) { u[0] = _x; u[1] = _y; } else { u[2] = _x; u[3] = _y; } } +#define WG4_STOREB(BD) { _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { WG4_XPOSE(br[t]) const uint32_t _a = static_cast( __cvta_generic_to_shared((BD) + bdst[t])); asm volatile( "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 " "[%0], {%1,%2,%3,%4};\n" ::"r"(_a), "r"(br[t][0]), "r"(br[t][1]), "r"(br[t][2]), "r"(br[t][3])); } } } +#define WG4_STOREB_R(SLOT) { _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { WG4_XPOSE(br[t]) asm volatile( "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 " "[%0+%5], {%1,%2,%3,%4};\n" ::"r"(bsm0[t]), "r"(br[t][0]), "r"(br[t][1]), "r"(br[t][2]), "r"(br[t][3]), "n"((SLOT) * ST_E * 2)); } } } +#define WG4_LOADA_R(BUF) { _Pragma("unroll") for (int kk = 0; kk < 4; ++kk) { const int4 _v = *reinterpret_cast(arol + kk * 1024); areg[0][kk][0] = _v.x; areg[0][kk][1] = _v.y; areg[0][kk][2] = _v.z; areg[0][kk][3] = _v.w; } } +#define WG4_MMA_R(BUF,SLOT) { const uint64_t db0 = dslot[0] + (uint64_t)(SLOT) * (ST_E / 8); asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); _Pragma("unroll") for (int kk = 0; kk < 4; ++kk) { WGMMA_RS_ANY_T(WG_BN, acc, areg[0][kk], db0 + kk, 0); } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); _Pragma("unroll") for (int kk = 0; kk < 4; ++kk) asm volatile("" : "+r"(areg[0][kk][0]), "+r"(areg[0][kk][1]), "+r"(areg[0][kk][2]), "+r"(areg[0][kk][3])); _Pragma("unroll") for (int kk = 4; kk < KT; ++kk) { const int4 _v = *reinterpret_cast(arol + kk * 1024); areg[0][kk - 4][0] = _v.x; areg[0][kk - 4][1] = _v.y; areg[0][kk - 4][2] = _v.z; areg[0][kk - 4][3] = _v.w; } asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); _Pragma("unroll") for (int kk = 4; kk < KT; ++kk) { WGMMA_RS_ANY_T(WG_BN, acc, areg[0][kk - 4], db0 + kk, 0); } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); arol += KT * 1024; } +#define WG4_LOADA(BUF,c) { if constexpr ((BUF) == 0) { const __half *_p = apt + (size_t)(c) * (KT * 1024); _Pragma("unroll") for (int kk = 0; kk < 4; ++kk) { const int4 _v = *reinterpret_cast(_p + kk * 1024); areg[0][kk][0] = _v.x; areg[0][kk][1] = _v.y; areg[0][kk][2] = _v.z; areg[0][kk][3] = _v.w; } arol = _p; } } +#define WG4_KEEPA(BUF) { _Pragma("unroll") for (int kk = 0; kk < KT - 4; ++kk) asm volatile("" : "+r"(areg[0][kk][0]), "+r"(areg[0][kk][1]), "+r"(areg[0][kk][2]), "+r"(areg[0][kk][3])); } +#define WG4_MMA(BUF,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t db0 = wg_desc((BD) + ng * WG_BN * 8, LBO_B * 2u, 128u); _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { WGMMA_RS_ANY_T(WG_BN, acc, areg[BUF][kk], db0 + kk, 0); } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + if constexpr (PDL && LATE_PDL) PDL_ENTER(); + constexpr int WGD = 1; + if constexpr (BC8) { +#pragma unroll 1 + for (int st = 0; st < CPA_STAGES; ++st) { + if (st < NCH) WG4_CPB(st, smem + st * ST_E) + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } + WG4_LOADA(0, 0) + if (ABUF > 1 && 1 < NCH) WG4_LOADA(1, 1) + if (ABUF > 2 && 2 < NCH) WG4_LOADA(2, 2) + } else { + WG4_LOADB(0) + WG4_LOADA(0, 0) + if (ABUF > 1 && 1 < NCH) WG4_LOADA(1, 1) + if (ABUF > 2 && 2 < NCH) WG4_LOADA(2, 2) + WG4_STOREB(smem) + if (1 < NCH) WG4_LOADB(1) + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + } +#define WG4_BODY(BUF,i,SLOT) { if constexpr (BC8) { asm volatile("cp.async.wait_group %0;\n" : : "n"(CPA_STAGES - 3) : "memory"); __syncthreads(); if ((i) >= 2) { if ((i) + 2 < NCH) { WG4_CPB_R(((SLOT) + 2) % CPA_STAGES) } asm volatile("cp.async.commit_group;\n" ::: "memory"); } } else { if (i) __syncthreads(); } WG4_MMA_R(BUF, SLOT) if constexpr (NSPLIT > 1 || BC8) { asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); WG4_KEEPA(0) if constexpr (NSPLIT > 1) __syncthreads(); } if constexpr (!BC8) { const int cb = (i) + 1; if (cb < NCH) { WG4_STOREB_R(((SLOT) + 1) % BSTAGES) asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); if (cb + 1 < NCH) WG4_LOADB(cb + 1) } } if ((i) + 1 < NCH) { if constexpr (NSPLIT == 1 && !BC8) { asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); WG4_KEEPA(0) } WG4_LOADA_R(0) } } + for (int i = 0; i < NCH; i += 4) { + WG4_BODY(0, i, 0) + if (i + 1 < NCH) WG4_BODY(1, i + 1, 1) + if (i + 2 < NCH) WG4_BODY(0, i + 2, 2) + if (i + 3 < NCH) WG4_BODY(1, i + 3, 3) + } + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + if constexpr (BC8) + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); +#undef WG4_CPB +#undef WG4_CPB_R +#undef WG4_STOREB_R +#undef WG4_LOADA_R +#undef WG4_MMA_R +#undef WG4_LOADB +#undef WG4_XPOSE +#undef WG4_STOREB +#undef WG4_LOADA +#undef WG4_KEEPA +#undef WG4_MMA +#undef WG4_BODY +#define TL4_SW(c,r) ((c) * BM + (((r) & ~7) ^ (((c) & 7) << 3)) + ((r) & 7)) + __syncthreads(); + const int cb2 = (lane & 3) * 2; + const bool snake = PREACT && (m0 + er0) < half; + const float sn_a0 = snake ? row_alpha[er0] : 0.f; + const float sn_r0 = snake ? row_ri[er0] : 0.f; + const float sn_a1 = snake ? row_alpha[er0 + 8] : 0.f; + const float sn_r1 = snake ? row_ri[er0 + 8] : 0.f; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + const int row = er0 + h * 8; + if constexpr (PREACT) { + if (snake) { + const float a = h ? sn_a1 : sn_a0; + const float ri = h ? sn_r1 : sn_r0; + const float s0 = __sinf(a * v.x); + const float s1 = __sinf(a * v.y); + v.x = __fmaf_rn(s0 * s0, ri, v.x); + v.y = __fmaf_rn(s1 * s1, ri, v.y); + } else { + v.x = v.x > 0.f ? v.x : 0.01f * v.x; + v.y = v.y > 0.f ? v.y : 0.01f * v.y; + } + if constexpr (OC8) { + const int c0 = ng * WG_BN + j * 8 + cb2; + Hs[TL4_SW(c0, row)] = __float2half(v.x); + Hs[TL4_SW(c0 + 1, row)] = __float2half(v.y); + } else { + *reinterpret_cast<__half2 *>( + Hs + row * WG_CLD + ng * WG_BN + j * 8 + cb2) = + __floats2half2_rn(v.x, v.y); + } + } else { + *reinterpret_cast( + Cs + row * WG_CLD + ng * WG_BN + j * 8 + cb2) = v; + } + } + } + __syncthreads(); + if constexpr (PREACT) { + if constexpr (OC8) { + constexpr int GRP = BM / 8; + constexpr int TOTG = GRP * BNT; + constexpr int GITER = (TOTG + NT - 1) / NT; + const size_t SC8o = (size_t)SCout * 8; +#pragma unroll + for (int t = 0; t < GITER; ++t) { + const int idx = tid + t * NT; + if (TOTG % NT != 0 && idx >= TOTG) continue; + const int g = idx / BNT; + const int j = idx - g * BNT; + if (!colok[j]) continue; + const int m0g = g * 8; + if (m0 + m0g >= Cout) continue; + *reinterpret_cast( + out + (size_t)((m0 + m0g) >> 3) * SC8o + + (size_t)(nb + j + H) * 8) = + *reinterpret_cast(Hs + TL4_SW(j, m0g)); + } + } else { + static_assert((BNT & 3) == 0, "four-column preactivated epilogue"); + constexpr int NC = BNT / 4; + constexpr int TOT = BM * NC; + for (int q = tid; q < TOT; q += NT) { + const int m = q / NC; + const int lc = (q - m * NC) * 4; + const int co = m0 + m; + const bool any4 = (colok[lc] | colok[lc + 1] | + colok[lc + 2] | colok[lc + 3]) != 0; + if (co >= Cout || !any4) continue; + const __half2 *sp = reinterpret_cast( + Hs + m * WG_CLD + lc); + const __half2 v0 = sp[0], v1 = sp[1]; + __half *op = out + (size_t)co * SCout + nb + lc + H; + const bool ok4 = (colok[lc] & colok[lc + 1] & + colok[lc + 2] & colok[lc + 3]) != 0; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = v0; + reinterpret_cast<__half2 *>(op)[1] = v1; + } else { + if (colok[lc + 0]) op[0] = v0.x; + if (colok[lc + 1]) op[1] = v0.y; + if (colok[lc + 2]) op[2] = v1.x; + if (colok[lc + 3]) op[3] = v1.y; + } + } + } + } else { + constexpr bool COMPACT_EPI = + (WG_BN == 136 || WG_BN == 160 || WG_BN == 208); + if constexpr (COMPACT_EPI) { + static_assert((BNT & 3) == 0, "4-column epilogue groups"); + constexpr int NC = BNT / 4; + constexpr int TOT = BM * NC; + constexpr int MITER = (TOT + NT - 1) / NT; + constexpr int UNR = (MITER >= 8) ? 8 : ((MITER >= 4) ? 4 : + ((MITER >= 2) ? 2 : 1)); +#pragma unroll 1 + for (int t0 = 0; t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int q = tid + (t0 + u) * NT; + if (q >= TOT) continue; + const int m = q / NC; + const int lc = (q - m * NC) * 4; + const int co = m0 + m; + const bool any4 = (colok[lc] | colok[lc + 1] | + colok[lc + 2] | colok[lc + 3]) != 0; + if (co >= Cout || !any4) continue; + v[u] = *reinterpret_cast(Cs + m * WG_CLD + lc); + if constexpr (!PREACT) { + const int n = nb + lc; + const __half2 *rp = reinterpret_cast( + resid + (size_t)co * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int q = tid + (t0 + u) * NT; + if (q >= TOT) continue; + const int m = q / NC; + const int lc = (q - m * NC) * 4; + const int co = m0 + m; + const bool any4 = (colok[lc] | colok[lc + 1] | + colok[lc + 2] | colok[lc + 3]) != 0; + if (co >= Cout || !any4) continue; + const int n = nb + lc; + float4 vv = v[u]; + if constexpr (!PREACT) { + vv.x += __half2float(r0[u].x); vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (co < half) { + const float a = row_alpha[m], ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + __half *op = out + (size_t)co * SCout + n + H; + const bool ok4 = (colok[lc] & colok[lc + 1] & + colok[lc + 2] & colok[lc + 3]) != 0; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + } else { + constexpr int NC = EPN / 4; + constexpr int MSTEP = NT / NC; + constexpr int MITER = BM / MSTEP; + constexpr int UNR = (MITER % 8 == 0) ? 8 + : ((MITER % 4 == 0) ? 4 + : ((MITER % 2 == 0) ? 2 : 1)); + static_assert(MITER % UNR == 0, "row batch must divide row count"); + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && + colok[lc + 2] && colok[lc + 3]; + const bool live4 = colok[lc] || colok[lc + 1] || + colok[lc + 2] || colok[lc + 3]; + const int n = nb + lc; +#pragma unroll 1 + for (int t0 = 0; live4 && t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + v[u] = *reinterpret_cast( + Cs + m * WG_CLD + lc); + if constexpr (!PREACT) { + const int co = m0 + m; + const int cc = co < Cout ? co : 0; + const __half2 *rp = reinterpret_cast( + resid + (size_t)cc * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + const int co = m0 + m; + if (co >= Cout) continue; + float4 vv = v[u]; + if constexpr (!PREACT) { + vv.x += __half2float(r0[u].x); + vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); + vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (co < half) { + const float a = row_alpha[m], ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + __half *op = out + (size_t)co * SCout + n + H; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + } + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B; + for (int idx = bxy * NT + tid; idx < total; idx += nbx * NT) { + const int b = idx / Cin, ci = idx - b * Cin; + const int x = b * L + native_true_tokens(b,T); + __half *const dst = ns + ((size_t)b * Cin + ci) * H; + if constexpr (BC8) { + const __half *const src = + in + (size_t)(ci >> 3) * ((size_t)SCin * 8) + + (size_t)x * 8 + (ci & 7); +#pragma unroll + for (int j = 0; j < H; ++j) dst[j] = src[j * 8]; + } else { + const __half *const src = in + (size_t)ci * SCin + x; +#pragma unroll + for (int j = 0; j < H; j += 2) + *reinterpret_cast(dst + j) = + *reinterpret_cast(src + j); + } + } + } + if (pad_src != nullptr) { + const int total = Cout * B; + for (int idx = bxy * NT + tid; idx < total; idx += nbx * NT) { + const int b = idx / Cout, co = idx - b * Cout; + const int x = b * L + pad_off; + const __half *const src = + pad_src + ((size_t)b * Cout + co) * pad_len; + if constexpr (OC8) { + __half *const dst = + out + (size_t)(co >> 3) * ((size_t)SCout * 8) + + (size_t)x * 8 + (co & 7); + for (int j = 0; j < pad_len; ++j) dst[j * 8] = src[j]; + } else { + __half *const dst = out + (size_t)co * SCout + x; + for (int j = 0; j < pad_len; ++j) dst[j] = src[j]; + } + } + } + } +} diff --git a/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wg5_step_widths.inc b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wg5_step_widths.inc new file mode 100644 index 0000000..94d8f0c --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wg5_step_widths.inc @@ -0,0 +1,409 @@ +template +__device__ __noinline__ +void wg5_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int KT = 3; + constexpr int BM = 64 * MSPLIT; + constexpr int BNT = WG_BN; + constexpr int NTS = 128 * MSPLIT; + constexpr int WG_NT = NTS * KSPLIT; + constexpr int NCRAW = BNT / 4; + constexpr int EPN = + (WG_NT % NCRAW == 0) ? BNT : (((BNT + 127) / 128) * 128); + constexpr int WG_CLD_R = (OC8 ? WG_BN : EPN) + 8; + constexpr int WG_CLD = (WG_CLD_R % 16 == 0) ? WG_CLD_R + 4 : WG_CLD_R; + constexpr int CG = CIB / 8; + constexpr int CSUB = CIB / 16; + constexpr int AK = KT * CSUB; + constexpr int NBRP = BC8 ? (((BNT + KT - 1) + 7) & ~7) + : (((BNT + KT - 1) + 31) & ~31); + constexpr int UCOL = NBRP / 32; + constexpr int NUNIT = CG * UCOL; + constexpr int NWARP = NTS / 32; + constexpr int BITER = (NUNIT + NWARP - 1) / NWARP; + constexpr int B_E = CIB * NBRP; + constexpr int ST_E = SB * B_E; + constexpr int LBO_B = NBRP * 8; + constexpr int ABUF = 3; + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *smem = reinterpret_cast<__half *>(arena); + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int kgrp = tid / NTS; + const int stid = tid - kgrp * NTS; + const int mg = stid >> 7; + const int warp = (stid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * BM; + const int nb = blockIdx.y * BNT; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + constexpr int NCH = (CIN + CIB - 1) / CIB; + static_assert(NCH % KSPLIT == 0, "wg5 K-split must divide the chunk count"); + bool active_tile = false; + for (int row = nb / L; row <= min(N - 1, nb + BNT - 1) / L; ++row) { + if (native_meta()->lengths[row] > 0) { + active_tile = true; + break; + } + } + if (active_tile) { + if (tid < BM) { + const int co = m0 + tid; + if (co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + for (int i = tid; i < EPN; i += WG_NT) { + const int n = nb + i; + colok[i] = (i < BNT && n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mg * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (kgrp == 0 && eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (kgrp == 0 && eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + const int mtile = blockIdx.x * MSPLIT + mg; + const __half *const apt = w + ((size_t)mtile * NCH * AK) * 1024 + + (size_t)(stid & 127) * 8; + uint32_t areg[ABUF][AK][4]; + const int wi = stid >> 5; + int bdst[BITER], boff[BITER]; + bool blive[BITER]; + const __half *const bbase = in + (size_t)(lane >> 2) * SCin + nb + + 8 * (lane & 3); +#pragma unroll + for (int t = 0; t < BITER; ++t) { + const int u = wi + t * NWARP; + const int gc = (u < NUNIT) ? u / UCOL : 0; + const int cu = (u < NUNIT) ? u - gc * UCOL : 0; + blive[t] = u < NUNIT; + bdst[t] = gc * LBO_B + (cu * 32 + lane) * 8; + boff[t] = gc * 8 * SCin + cu * 32; + } + const size_t SC8i = (size_t)SCin * 8; + constexpr int BQ = B_E / 8; + constexpr int BCREG = (BQ + NTS - 1) / NTS; + int cdst[BCREG]; + const __half *csrc[BCREG]; + bool clive[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) { + const int q = stid + r * NTS; + const int gc = q / NBRP, rem = q - gc * NBRP; + clive[r] = ((BQ % NTS == 0) || q < BQ) && gc < CG; + cdst[r] = clive[r] ? gc * LBO_B + rem * 8 : 0; + csrc[r] = in + (clive[r] ? (size_t)gc * SC8i + + (size_t)nb * 8 + rem * 8 : 0); + } + uint32_t br[BITER][4]; + __half *const kbB = smem + kgrp * ST_E; + const unsigned barid = 1u + kgrp + (unsigned)(KSPLIT * native_subid()); + uint64_t dbuf[ABUF]; +#pragma unroll + for (int s = 0; s < ABUF; ++s) + dbuf[s] = wg_desc(kbB + s * B_E, LBO_B * 2u, 128u); + uint32_t csm0[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) + csm0[r] = static_cast( + __cvta_generic_to_shared(kbB + cdst[r])); + const size_t CSTEP = (size_t)KSPLIT * CG * SC8i; + constexpr int ASTEP = KSPLIT * AK * 1024; + const __half *crol[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) + crol[r] = csrc[r] + (size_t)(ABUF * KSPLIT + kgrp) * CG * SC8i; + const __half *arol = apt + (size_t)(ABUF * KSPLIT + kgrp) * (AK * 1024); + int crolc = ABUF * KSPLIT + kgrp; +#define WG5_BAR() asm volatile("bar.sync %0, %1;\n" ::"r"(barid), "r"(NTS)) +#define WG5_CPB(c,BD) { const size_t _o = (size_t)(c) * CG * SC8i; _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16((BD) + cdst[r], csrc[r] + _o, true); if ((c) / KSPLIT == ABUF - 1) { const size_t _r = (size_t)(ABUF * KSPLIT + kgrp) * CG * SC8i; _Pragma("unroll") for (int r = 0; r < BCREG; ++r) csrc[r] += _r; } } +#define WG5_LOADA(BUF,c) { (void)(c); } +#define WG5_KEEPA(BUF) { _Pragma("unroll") for (int kk = 0; kk < AK; ++kk) asm volatile("" : "+r"(areg[BUF][kk][0]), "+r"(areg[BUF][kk][1]), "+r"(areg[BUF][kk][2]), "+r"(areg[BUF][kk][3])); } +#define WG5_LOADB(c) { const size_t _o = (size_t)(c) * CIB * SCin; _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (NUNIT == NWARP * BITER || blive[t]) { const int4 _v = *reinterpret_cast( bbase + _o + boff[t]); br[t][0] = _v.x; br[t][1] = _v.y; br[t][2] = _v.z; br[t][3] = _v.w; } } } +#define WG5_BXPOSE(u) { const int _q = lane & 3; uint32_t _x, _y; _x = (_q & 1) ? u[0] : u[1]; _y = (_q & 1) ? u[2] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 1); _y = __shfl_xor_sync(0xffffffffu, _y, 1); if (_q & 1) { u[0] = _x; u[2] = _y; } else { u[1] = _x; u[3] = _y; } _x = (_q & 2) ? u[0] : u[2]; _y = (_q & 2) ? u[1] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 2); _y = __shfl_xor_sync(0xffffffffu, _y, 2); if (_q & 2) { u[0] = _x; u[1] = _y; } else { u[2] = _x; u[3] = _y; } } +#define WG5_STOREB(BD) { _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (NUNIT == NWARP * BITER || blive[t]) { WG5_BXPOSE(br[t]) const uint32_t _a = static_cast( __cvta_generic_to_shared((BD) + bdst[t])); asm volatile( "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 " "[%0], {%1,%2,%3,%4};\n" ::"r"(_a), "r"(br[t][0]), "r"(br[t][1]), "r"(br[t][2]), "r"(br[t][3])); } } } +#define WG5_MMA(BUF) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t db0 = dbuf[0] + (uint64_t)(BUF) * (B_E / 8); _Pragma("unroll") for (int h = 0; h < CSUB; ++h) { _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { WGMMA_RS_ANY_T(WG_BN, acc, areg[BUF][h * KT + kk], db0 + kk + h * (LBO_B / 4), 0); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } +#define WG5_CPB_R(SLOT) { _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16o<(SLOT) * B_E * 2>(csm0[r], csrc[r], true); } +#define WG5_LOADA_R(BUF) { } +#define WG5_ROLL() { crolc += KSPLIT; arol += ASTEP; if constexpr (BC8) { _Pragma("unroll") for (int r = 0; r < BCREG; ++r) csrc[r] += CSTEP; } } + if constexpr (PDL) { + (void)0; + if (tid == 0) + (void)0; + } + static_assert(SB == 3 && WGD == 1 && ABUF == 3, "wg5 pipeline shape"); + if constexpr (BC8) { +#pragma unroll 1 + for (int s = 0; s < SB; ++s) { + const int c = s * KSPLIT + kgrp; + if (c < NCH) WG5_CPB(c, kbB + s * B_E) + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } + } else { + WG5_LOADB(kgrp) +#pragma unroll 1 + for (int s = 0; s < SB; ++s) { + const int c = s * KSPLIT + kgrp; + if (c < NCH) WG5_STOREB(kbB + s * B_E) + if (c + KSPLIT < NCH) WG5_LOADB(c + KSPLIT) + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + } + WG5_LOADA(0, kgrp) + if (KSPLIT + kgrp < NCH) WG5_LOADA(1, KSPLIT + kgrp) + if (2 * KSPLIT + kgrp < NCH) WG5_LOADA(2, 2 * KSPLIT + kgrp) + constexpr int NIT = NCH / KSPLIT; + static_assert(NIT % ABUF == 0, "wg5 mainloop unroll shape"); +#define WG5_BODY(BUF,i) { { const __half *_at = apt + (size_t)((i) * KSPLIT + kgrp) * (AK * 1024); _Pragma("unroll") for (int kk = 0; kk < AK; ++kk) { const int4 _v = *reinterpret_cast(_at + kk * 1024); areg[BUF][kk][0] = _v.x; areg[BUF][kk][1] = _v.y; areg[BUF][kk][2] = _v.z; areg[BUF][kk][3] = _v.w; } } if constexpr (BC8) asm volatile("cp.async.wait_group %0;\n" : : "n"(SB - 2) : "memory"); WG5_BAR(); WG5_MMA(BUF) asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); WG5_KEEPA(BUF) if constexpr (!BC8) if ((i) > 1 && crolc < NCH) WG5_LOADB(crolc) if constexpr (NTS > 128) WG5_BAR(); if ((i) >= 1) { const int nb2 = ((BUF) + 2) % ABUF; if (crolc < NCH) { if constexpr (BC8) { WG5_CPB_R(nb2) } else { WG5_STOREB(kbB + nb2 * B_E) asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); } WG5_LOADA_R(nb2) } WG5_ROLL() if constexpr (BC8) asm volatile("cp.async.commit_group;\n" ::: "memory"); } } +#pragma unroll + for (int i = 0; i < NIT; i += 3) { + WG5_BODY(0, i) + WG5_BODY(1, i + 1) + WG5_BODY(2, i + 2) + } + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + if constexpr (BC8) + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); +#undef WG5_BXPOSE +#undef WG5_CPB +#undef WG5_CPB_R +#undef WG5_LOADA_R +#undef WG5_ROLL +#undef WG5_LOADB +#undef WG5_STOREB +#undef WG5_MMA +#undef WG5_LOADA +#undef WG5_KEEPA +#undef WG5_BODY +#undef WG5_BAR + __syncthreads(); + const int cb2 = (lane & 3) * 2; + if constexpr (OC8) { + __half *Ck = smem + kgrp * BM * WG_CLD; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + *reinterpret_cast<__half2 *>(Ck + (er0 + h * 8) * WG_CLD + + j * 8 + cb2) = + __floats2half2_rn(acc[j * 4 + h * 2], + acc[j * 4 + h * 2 + 1]); + } + } + } else { + float *Ck = Cs + kgrp * BM * WG_CLD; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], + acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Ck + (er0 + h * 8) * WG_CLD + + j * 8 + cb2) = v; + } + } + } + __syncthreads(); + if constexpr (OC8) { + constexpr int GRP = BM / 8; + constexpr int TOTG = GRP * WG_BN; + constexpr int GITER = (TOTG + WG_NT - 1) / WG_NT; + const size_t SC8o = (size_t)SCout * 8; +#pragma unroll + for (int t = 0; t < GITER; ++t) { + const int idx = tid + t * WG_NT; + if (TOTG % WG_NT != 0 && idx >= TOTG) continue; + const int j = idx - (idx / WG_BN) * WG_BN, g = idx / WG_BN; + if (!colok[j]) continue; + const int m0g = g * 8; + if (m0 + m0g >= Cout) continue; + const int n = nb + j; + __half rr[8]; + if constexpr (CIN == 1536) { + if constexpr (RC8) { + *reinterpret_cast(rr) = + *reinterpret_cast( + resid + (size_t)((m0 + m0g) >> 3) * SC8o + + (size_t)(n + H) * 8); + } else { +#pragma unroll + for (int h = 0; h < 8; ++h) + rr[h] = resid[(size_t)(m0 + m0g + h) * SCout + n + H]; + } + } + uint32_t pk[4]; + const bool snake = (m0 + m0g) < half; +#pragma unroll + for (int h = 0; h < 8; h += 2) { + float f[2]; +#pragma unroll + for (int e = 0; e < 2; ++e) { + const int m = m0g + h + e; + float a = __half2float(smem[m * WG_CLD + j]); +#pragma unroll + for (int kg = 1; kg < KSPLIT; ++kg) + a += __half2float( + smem[kg * BM * WG_CLD + m * WG_CLD + j]); + if constexpr (CIN == 1536) + a += __half2float(rr[h + e]); + if (snake) { + const float sn = __sinf(row_alpha[m] * a); + a = __fmaf_rn(sn * sn, row_ri[m], a); + } else { + a = a > 0.f ? a : 0.01f * a; + } + f[e] = a; + } + const __half2 p = __floats2half2_rn(f[0], f[1]); + pk[h >> 1] = *reinterpret_cast(&p); + } + *reinterpret_cast( + out + (size_t)((m0 + m0g) >> 3) * SC8o + + (size_t)(n + H) * 8) = make_int4(pk[0], pk[1], pk[2], pk[3]); + } + } else { + constexpr int NC = EPN / 4; + constexpr int MSTEP = WG_NT / NC; + constexpr int MITER = BM / MSTEP; + constexpr int UNR = (MITER % 4 == 0) ? 4 : ((MITER % 2 == 0) ? 2 : 1); + static_assert(MITER % UNR == 0, "row batch must divide the row count"); + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + const bool live4 = colok[lc] || colok[lc + 1] || colok[lc + 2] || + colok[lc + 3]; + const int n = nb + lc; +#pragma unroll 1 + for (int t0 = 0; live4 && t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + v[u] = *reinterpret_cast(Cs + m * WG_CLD + lc); + if constexpr (CIN == 1536) { + const int co = m0 + m; + const int cc = co < Cout ? co : 0; + if constexpr (RC8) { + const __half *rp = + resid + (size_t)(cc >> 3) * ((size_t)SCout * 8) + + (size_t)(n + H) * 8 + (cc & 7); + r0[u].x = rp[0]; r0[u].y = rp[8]; + r1[u].x = rp[16]; r1[u].y = rp[24]; + } else { + const __half2 *rp = reinterpret_cast( + resid + (size_t)cc * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + const int co = m0 + m; + if (co >= Cout) continue; + float4 vv = v[u]; +#pragma unroll + for (int g = 1; g < KSPLIT; ++g) { + float4 p = *reinterpret_cast( + Cs + g * BM * WG_CLD + m * WG_CLD + lc); + vv.x += p.x; vv.y += p.y; vv.z += p.z; vv.w += p.w; + } + if constexpr (CIN == 1536) { + vv.x += __half2float(r0[u].x); vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + __half *op = out + (size_t)co * SCout + n + H; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + static_assert(CIN == 768 || CIN == 1536); + constexpr int COUT = (CIN == 768) ? 1536 : 768; + if (ns != nullptr) { + const int total = CIN * B; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int b = idx / CIN, ci = idx - b * CIN; + const int x = b * L + native_true_tokens(b,T); + __half *const dst = ns + ((size_t)b * CIN + ci) * H; + if constexpr (BC8) { + const __half *const src = + in + (size_t)(ci >> 3) * ((size_t)SCin * 8) + + (size_t)x * 8 + (ci & 7); +#pragma unroll + for (int j = 0; j < H; ++j) dst[j] = src[j * 8]; + } else { + const __half *const src = in + (size_t)ci * SCin + x; +#pragma unroll + for (int j = 0; j < H; j += 2) + *reinterpret_cast(dst + j) = + *reinterpret_cast(src + j); + } + } + } + if (pad_src != nullptr) { + const int total = COUT * B; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int b = idx / COUT, co = idx - b * COUT; + const int x = b * L + pad_off; + const __half *const src = + pad_src + ((size_t)b * COUT + co) * pad_len; + if constexpr (OC8) { + __half *const dst = + out + (size_t)(co >> 3) * ((size_t)SCout * 8) + + (size_t)x * 8 + (co & 7); + for (int j = 0; j < pad_len; ++j) dst[j * 8] = src[j]; + } else { + __half *const dst = out + (size_t)co * SCout + x; + for (int j = 0; j < pad_len; ++j) dst[j] = src[j]; + } + } + } + } +} diff --git a/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wg_kernel.inc b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wg_kernel.inc new file mode 100644 index 0000000..2d5c766 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wg_kernel.inc @@ -0,0 +1,182 @@ +__device__ __forceinline__ uint64_t wg_desc(const void *p, uint32_t lbo_bytes, + uint32_t sbo_bytes) { + uint32_t a = static_cast(__cvta_generic_to_shared(p)); + uint64_t d = static_cast((a & 0x3FFFFu) >> 4); + d |= static_cast((lbo_bytes & 0x3FFFFu) >> 4) << 16; + d |= static_cast((sbo_bytes & 0x3FFFFu) >> 4) << 32; + return d; +} +#include "wgmma_macro.inc" +constexpr int WG_SLOT = 72; +template +__device__ __noinline__ void wg_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int WG_CLD = WG_BN + 8; + constexpr int WG_NT = 128 * NWG; + constexpr int WG_BM = 64 * NWG; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int SBO_AE = KG * WG_SLOT; + constexpr int LBO_BE = (WG_BN / 8) * WG_SLOT; + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *As0 = reinterpret_cast<__half *>(arena); + __half *Bs0 = As0 + A_E; + __half *As1 = (STAGES == 2) ? Bs0 + B_E : As0; + __half *Bs1 = (STAGES == 2) ? Bs0 + B_E + A_E : Bs0; + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int wgid = tid >> 7; + const int warp = (tid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * WG_BM; + constexpr int WGROWS = 8; + const int n0 = blockIdx.y * WG_BN; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + if (tid < WG_BM) { + const int co = m0 + tid; + if (do_snake && co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + if (tid < WG_BN) { + const int n = n0 + tid; + colok[tid] = (n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = wgid * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + constexpr int ATOT = WG_BM * KG; + constexpr int AREG = (ATOT + WG_NT - 1) / WG_NT; + constexpr int CIB = BK / KT; + constexpr int BTOT = CIB * (WG_BN / 8); + constexpr int BREG = (BTOT + WG_NT - 1) / WG_NT; + int4 ar[AREG]; + int4 br0[BREG], br1[BREG]; + uint32_t edge[BREG]; +#define WG_LOADG(kg0) { const int _ci0 = (kg0) / KT; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; ar[r] = make_int4(0, 0, 0, 0); if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; const int co = m0 + m, kg = (kg0) + v * 8; if (co < Cout && kg < Kg) ar[r] = *reinterpret_cast( w + (size_t)co * Kg + kg); } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; br0[r] = make_int4(0, 0, 0, 0); if constexpr (KT == 3) edge[r] = 0; else br1[r] = make_int4(0, 0, 0, 0); if (BTOT % WG_NT == 0 || idx < BTOT) { const int ci = _ci0 + idx / (WG_BN / 8); const int nc = idx % (WG_BN / 8); if (ci < Cin) { const int4 *src = reinterpret_cast( in + (size_t)ci * SCin + n0 + nc * 8); br0[r] = src[0]; if constexpr (KT == 3) { if (nc == WG_BN / 8 - 1) edge[r] = src[1].x; } else { br1[r] = src[1]; } } } } } + WG_LOADG(0) +#define WG_STS(AD,BD) { _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; *reinterpret_cast( (AD) + (m >> 3) * SBO_AE + v * WG_SLOT + (m & 7) * 8) = ar[r]; } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; if (BTOT % WG_NT == 0 || idx < BTOT) { const int cil = idx / (WG_BN / 8), nc = idx % (WG_BN / 8); unsigned int u[8]; u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; if constexpr (KT == 3) { u[4] = __shfl_down_sync(0xffffffffu, u[0], 1, WG_BN / 8); if (nc == WG_BN / 8 - 1) u[4] = edge[r]; u[5] = u[6] = u[7] = 0; } else { u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; } _Pragma("unroll") for (int k = 0; k < KT; ++k) { const int kh = k >> 1; int4 o; if (k & 1) { o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); } else { o.x = u[kh + 0]; o.y = u[kh + 1]; o.z = u[kh + 2]; o.w = u[kh + 3]; } const int row = cil * KT + k; *reinterpret_cast((BD) + nc * WG_SLOT + (row >> 3) * LBO_BE + (row & 7) * 8) = o; } } } if constexpr (WG_BN == 64) { asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); } } +#define WG_MMA(AD,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t da0 = wg_desc((AD) + wgid * WGROWS * SBO_AE, WG_SLOT * 2, SBO_AE * 2); const uint64_t db0 = wg_desc((BD), LBO_BE * 2, WG_SLOT * 2); _Pragma("unroll") for (int kk = 0; kk < BK; kk += 16) { const uint64_t da = da0 + (kk >> 3) * (WG_SLOT / 8); const uint64_t db = db0 + (kk >> 3) * (LBO_BE / 8); if constexpr (WG_BN == 64) { WGMMA_M64N64K16(acc, da, db); } else { WGMMA_M64N128K16(acc, da, db); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + WG_STS(As0, Bs0) + if (BK < Kg) WG_LOADG(BK) + __syncthreads(); + int cur = 0; + for (int kg0 = 0; kg0 < Kg; kg0 += BK) { + if (cur == 0) { WG_MMA(As0, Bs0) } else { WG_MMA(As1, Bs1) } + if (kg0 + BK < Kg) { + if (STAGES == 2) + asm volatile("wgmma.wait_group.sync.aligned 1;\n" ::: "memory"); + else + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + __syncthreads(); + if (cur == 0) { WG_STS(As1, Bs1) } else { WG_STS(As0, Bs0) } + if (kg0 + 2 * BK < Kg) WG_LOADG(kg0 + 2 * BK) + __syncthreads(); + } else { + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + } + cur ^= 1; + } +#undef WG_LOADG +#undef WG_STS +#undef WG_MMA + __syncthreads(); + const int r0 = er0; + const int cb = (lane & 3) * 2; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Cs + (r0 + h * 8) * WG_CLD + j * 8 + cb) = v; + } + } + __syncthreads(); + constexpr int NC = WG_BN / 4; + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + for (int m = lr; m < WG_BM; m += WG_NT / NC) { + const int co = m0 + m; + if (co >= Cout) continue; + const int n = n0 + lc; + float4 v = *reinterpret_cast(Cs + m * WG_CLD + lc); + __half *op = out + (size_t)co * SCout + n + H; + if (resid != nullptr) { + const __half2 *rp = reinterpret_cast( + resid + (size_t)co * SCout + n + H); + __half2 r0h = rp[0], r1h = rp[1]; + v.x += __half2float(r0h.x); v.y += __half2float(r0h.y); + v.z += __half2float(r1h.x); v.w += __half2float(r1h.y); + } + if (do_snake) { + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + } + if (ok4) { + __half2 o0 = __floats2half2_rn(v.x, v.y); + __half2 o1 = __floats2half2_rn(v.z, v.w); + reinterpret_cast<__half2 *>(op)[0] = o0; + reinterpret_cast<__half2 *>(op)[1] = o1; + } else { + const float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B * H; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int total = Cout * B * pad_len; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} diff --git a/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wgmma_macro.inc b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wgmma_macro.inc new file mode 100644 index 0000000..db624e9 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wgmma_macro.inc @@ -0,0 +1,102 @@ +#define WGMMA_M64N8K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " "{" "%0,%1,%2,%3}" ", %4, %5, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N8K16(D,DA,DB) WGMMA_M64N8K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N8K16_ZERO_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " "{" "%0,%1,%2,%3}" ", %4, %5, p, 1, 1, " #TA ", " #TB ";\n}\n" : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) : "l"(DA), "l"(DB), "r"(0)) +#define WGMMA_M64N16K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %10, 0;\n" "wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7}" ", %8, %9, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N16K16(D,DA,DB) WGMMA_M64N16K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N24K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %14, 0;\n" "wgmma.mma_async.sync.aligned.m64n24k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11}" ", %12, %13, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N24K16(D,DA,DB) WGMMA_M64N24K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N32K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %18, 0;\n" "wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15}" ", %16, %17, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N32K16(D,DA,DB) WGMMA_M64N32K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N40K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %22, 0;\n" "wgmma.mma_async.sync.aligned.m64n40k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19}" ", %20, %21, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N40K16(D,DA,DB) WGMMA_M64N40K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N48K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %26, 0;\n" "wgmma.mma_async.sync.aligned.m64n48k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N48K16(D,DA,DB) WGMMA_M64N48K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N56K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %30, 0;\n" "wgmma.mma_async.sync.aligned.m64n56k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N56K16(D,DA,DB) WGMMA_M64N56K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N64K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %34, 0;\n" "wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31}" ", %32, %33, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N64K16(D,DA,DB) WGMMA_M64N64K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N72K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %38, 0;\n" "wgmma.mma_async.sync.aligned.m64n72k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35}" ", %36, %37, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N72K16(D,DA,DB) WGMMA_M64N72K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N80K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %42, 0;\n" "wgmma.mma_async.sync.aligned.m64n80k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39}" ", %40, %41, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N80K16(D,DA,DB) WGMMA_M64N80K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N88K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %46, 0;\n" "wgmma.mma_async.sync.aligned.m64n88k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43}" ", %44, %45, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N88K16(D,DA,DB) WGMMA_M64N88K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N96K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %50, 0;\n" "wgmma.mma_async.sync.aligned.m64n96k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47}" ", %48, %49, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N96K16(D,DA,DB) WGMMA_M64N96K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N104K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %54, 0;\n" "wgmma.mma_async.sync.aligned.m64n104k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51}" ", %52, %53, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N104K16(D,DA,DB) WGMMA_M64N104K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N112K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %58, 0;\n" "wgmma.mma_async.sync.aligned.m64n112k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55}" ", %56, %57, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N112K16(D,DA,DB) WGMMA_M64N112K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N120K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %62, 0;\n" "wgmma.mma_async.sync.aligned.m64n120k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59}" ", %60, %61, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N120K16(D,DA,DB) WGMMA_M64N120K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N128K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %66, 0;\n" "wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63}" ", %64, %65, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N128K16(D,DA,DB) WGMMA_M64N128K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N136K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %70, 0;\n" "wgmma.mma_async.sync.aligned.m64n136k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67}" ", %68, %69, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N136K16(D,DA,DB) WGMMA_M64N136K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N144K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %74, 0;\n" "wgmma.mma_async.sync.aligned.m64n144k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71}" ", %72, %73, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N144K16(D,DA,DB) WGMMA_M64N144K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N152K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %78, 0;\n" "wgmma.mma_async.sync.aligned.m64n152k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75}" ", %76, %77, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N152K16(D,DA,DB) WGMMA_M64N152K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N160K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %82, 0;\n" "wgmma.mma_async.sync.aligned.m64n160k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79}" ", %80, %81, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N160K16(D,DA,DB) WGMMA_M64N160K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N168K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %86, 0;\n" "wgmma.mma_async.sync.aligned.m64n168k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83}" ", %84, %85, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N168K16(D,DA,DB) WGMMA_M64N168K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N176K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %90, 0;\n" "wgmma.mma_async.sync.aligned.m64n176k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87}" ", %88, %89, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N176K16(D,DA,DB) WGMMA_M64N176K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N184K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %94, 0;\n" "wgmma.mma_async.sync.aligned.m64n184k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91}" ", %92, %93, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N184K16(D,DA,DB) WGMMA_M64N184K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N192K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %98, 0;\n" "wgmma.mma_async.sync.aligned.m64n192k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95}" ", %96, %97, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N192K16(D,DA,DB) WGMMA_M64N192K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N200K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %102, 0;\n" "wgmma.mma_async.sync.aligned.m64n200k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99}" ", %100, %101, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N200K16(D,DA,DB) WGMMA_M64N200K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N208K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %106, 0;\n" "wgmma.mma_async.sync.aligned.m64n208k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103}" ", %104, %105, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N208K16(D,DA,DB) WGMMA_M64N208K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N216K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %110, 0;\n" "wgmma.mma_async.sync.aligned.m64n216k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107}" ", %108, %109, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N216K16(D,DA,DB) WGMMA_M64N216K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N224K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %114, 0;\n" "wgmma.mma_async.sync.aligned.m64n224k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111}" ", %112, %113, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N224K16(D,DA,DB) WGMMA_M64N224K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N232K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %118, 0;\n" "wgmma.mma_async.sync.aligned.m64n232k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115}" ", %116, %117, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N232K16(D,DA,DB) WGMMA_M64N232K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N240K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %122, 0;\n" "wgmma.mma_async.sync.aligned.m64n240k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119}" ", %120, %121, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N240K16(D,DA,DB) WGMMA_M64N240K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N248K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %126, 0;\n" "wgmma.mma_async.sync.aligned.m64n248k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123}" ", %124, %125, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N248K16(D,DA,DB) WGMMA_M64N248K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N256K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %130, 0;\n" "wgmma.mma_async.sync.aligned.m64n256k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123,%124,%125,%126,%127}" ", %128, %129, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]), "+f"(D[124]), "+f"(D[125]), "+f"(D[126]), "+f"(D[127]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N256K16(D,DA,DB) WGMMA_M64N256K16_T(D, DA, DB, 0, 1) +#define WGMMA_ANY_T(NW,D,DA,DB,TA,TB) do { if constexpr (NW == 8) { WGMMA_M64N8K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 16) { WGMMA_M64N16K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 24) { WGMMA_M64N24K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 32) { WGMMA_M64N32K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 40) { WGMMA_M64N40K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 48) { WGMMA_M64N48K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 56) { WGMMA_M64N56K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 64) { WGMMA_M64N64K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 72) { WGMMA_M64N72K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 80) { WGMMA_M64N80K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 88) { WGMMA_M64N88K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 96) { WGMMA_M64N96K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 104) { WGMMA_M64N104K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 112) { WGMMA_M64N112K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 120) { WGMMA_M64N120K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 128) { WGMMA_M64N128K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 136) { WGMMA_M64N136K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 144) { WGMMA_M64N144K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 152) { WGMMA_M64N152K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 160) { WGMMA_M64N160K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 168) { WGMMA_M64N168K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 176) { WGMMA_M64N176K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 184) { WGMMA_M64N184K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 192) { WGMMA_M64N192K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 200) { WGMMA_M64N200K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 208) { WGMMA_M64N208K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 216) { WGMMA_M64N216K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 224) { WGMMA_M64N224K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 232) { WGMMA_M64N232K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 240) { WGMMA_M64N240K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 248) { WGMMA_M64N248K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 256) { WGMMA_M64N256K16_T(D, DA, DB, TA, TB); } } while (0) +#define WGMMA_ANY(NW,D,DA,DB) WGMMA_ANY_T(NW, D, DA, DB, 0, 1) +#define WGMMA_RS_M64N96K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %53, 0;\n" "wgmma.mma_async.sync.aligned.m64n96k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47}" ", {%48,%49,%50,%51}, %52, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N104K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %57, 0;\n" "wgmma.mma_async.sync.aligned.m64n104k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51}" ", {%52,%53,%54,%55}, %56, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N112K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %61, 0;\n" "wgmma.mma_async.sync.aligned.m64n112k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55}" ", {%56,%57,%58,%59}, %60, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N120K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %65, 0;\n" "wgmma.mma_async.sync.aligned.m64n120k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59}" ", {%60,%61,%62,%63}, %64, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N128K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %69, 0;\n" "wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63}" ", {%64,%65,%66,%67}, %68, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_ANY_T(NW,D,A,DB,TB) do { if constexpr (NW == 96) { WGMMA_RS_M64N96K16_T(D, A, DB, TB); } else if constexpr (NW == 104) { WGMMA_RS_M64N104K16_T(D, A, DB, TB); } else if constexpr (NW == 112) { WGMMA_RS_M64N112K16_T(D, A, DB, TB); } else if constexpr (NW == 120) { WGMMA_RS_M64N120K16_T(D, A, DB, TB); } else if constexpr (NW == 128) { WGMMA_RS_M64N128K16_T(D, A, DB, TB); } } while (0) +#define WGMMA_RS_M64N8K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %9, 0;\n" "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " "{" "%0,%1,%2,%3}" ", {%4,%5,%6,%7}, %8, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N16K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %13, 0;\n" "wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7}" ", {%8,%9,%10,%11}, %12, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N24K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %17, 0;\n" "wgmma.mma_async.sync.aligned.m64n24k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11}" ", {%12,%13,%14,%15}, %16, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N32K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %21, 0;\n" "wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15}" ", {%16,%17,%18,%19}, %20, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N40K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %25, 0;\n" "wgmma.mma_async.sync.aligned.m64n40k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N48K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %29, 0;\n" "wgmma.mma_async.sync.aligned.m64n48k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N56K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %33, 0;\n" "wgmma.mma_async.sync.aligned.m64n56k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N64K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %37, 0;\n" "wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31}" ", {%32,%33,%34,%35}, %36, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N72K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %41, 0;\n" "wgmma.mma_async.sync.aligned.m64n72k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35}" ", {%36,%37,%38,%39}, %40, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N80K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %45, 0;\n" "wgmma.mma_async.sync.aligned.m64n80k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39}" ", {%40,%41,%42,%43}, %44, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N88K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %49, 0;\n" "wgmma.mma_async.sync.aligned.m64n88k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43}" ", {%44,%45,%46,%47}, %48, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N136K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %73, 0;\n" "wgmma.mma_async.sync.aligned.m64n136k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67}" ", {%68,%69,%70,%71}, %72, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N144K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %77, 0;\n" "wgmma.mma_async.sync.aligned.m64n144k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71}" ", {%72,%73,%74,%75}, %76, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N152K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %81, 0;\n" "wgmma.mma_async.sync.aligned.m64n152k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75}" ", {%76,%77,%78,%79}, %80, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N160K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %85, 0;\n" "wgmma.mma_async.sync.aligned.m64n160k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79}" ", {%80,%81,%82,%83}, %84, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N168K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %89, 0;\n" "wgmma.mma_async.sync.aligned.m64n168k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83}" ", {%84,%85,%86,%87}, %88, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N176K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %93, 0;\n" "wgmma.mma_async.sync.aligned.m64n176k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87}" ", {%88,%89,%90,%91}, %92, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N184K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %97, 0;\n" "wgmma.mma_async.sync.aligned.m64n184k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91}" ", {%92,%93,%94,%95}, %96, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N192K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %101, 0;\n" "wgmma.mma_async.sync.aligned.m64n192k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95}" ", {%96,%97,%98,%99}, %100, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N200K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %105, 0;\n" "wgmma.mma_async.sync.aligned.m64n200k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99}" ", {%100,%101,%102,%103}, %104, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N208K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %109, 0;\n" "wgmma.mma_async.sync.aligned.m64n208k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103}" ", {%104,%105,%106,%107}, %108, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N216K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %113, 0;\n" "wgmma.mma_async.sync.aligned.m64n216k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107}" ", {%108,%109,%110,%111}, %112, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N224K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %117, 0;\n" "wgmma.mma_async.sync.aligned.m64n224k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111}" ", {%112,%113,%114,%115}, %116, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N232K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %121, 0;\n" "wgmma.mma_async.sync.aligned.m64n232k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115}" ", {%116,%117,%118,%119}, %120, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N240K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %125, 0;\n" "wgmma.mma_async.sync.aligned.m64n240k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119}" ", {%120,%121,%122,%123}, %124, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N248K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %129, 0;\n" "wgmma.mma_async.sync.aligned.m64n248k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123}" ", {%124,%125,%126,%127}, %128, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N256K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %133, 0;\n" "wgmma.mma_async.sync.aligned.m64n256k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123,%124,%125,%126,%127}" ", {%128,%129,%130,%131}, %132, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]), "+f"(D[124]), "+f"(D[125]), "+f"(D[126]), "+f"(D[127]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#undef WGMMA_RS_ANY_T +#define WGMMA_RS_ANY_T(NW,D,A,DB,TB) do { if constexpr (NW == 8) { WGMMA_RS_M64N8K16_T(D, A, DB, TB); } else if constexpr (NW == 16) { WGMMA_RS_M64N16K16_T(D, A, DB, TB); } else if constexpr (NW == 24) { WGMMA_RS_M64N24K16_T(D, A, DB, TB); } else if constexpr (NW == 32) { WGMMA_RS_M64N32K16_T(D, A, DB, TB); } else if constexpr (NW == 40) { WGMMA_RS_M64N40K16_T(D, A, DB, TB); } else if constexpr (NW == 48) { WGMMA_RS_M64N48K16_T(D, A, DB, TB); } else if constexpr (NW == 56) { WGMMA_RS_M64N56K16_T(D, A, DB, TB); } else if constexpr (NW == 64) { WGMMA_RS_M64N64K16_T(D, A, DB, TB); } else if constexpr (NW == 72) { WGMMA_RS_M64N72K16_T(D, A, DB, TB); } else if constexpr (NW == 80) { WGMMA_RS_M64N80K16_T(D, A, DB, TB); } else if constexpr (NW == 88) { WGMMA_RS_M64N88K16_T(D, A, DB, TB); } else if constexpr (NW == 96) { WGMMA_RS_M64N96K16_T(D, A, DB, TB); } else if constexpr (NW == 104) { WGMMA_RS_M64N104K16_T(D, A, DB, TB); } else if constexpr (NW == 112) { WGMMA_RS_M64N112K16_T(D, A, DB, TB); } else if constexpr (NW == 120) { WGMMA_RS_M64N120K16_T(D, A, DB, TB); } else if constexpr (NW == 128) { WGMMA_RS_M64N128K16_T(D, A, DB, TB); } else if constexpr (NW == 136) { WGMMA_RS_M64N136K16_T(D, A, DB, TB); } else if constexpr (NW == 144) { WGMMA_RS_M64N144K16_T(D, A, DB, TB); } else if constexpr (NW == 152) { WGMMA_RS_M64N152K16_T(D, A, DB, TB); } else if constexpr (NW == 160) { WGMMA_RS_M64N160K16_T(D, A, DB, TB); } else if constexpr (NW == 168) { WGMMA_RS_M64N168K16_T(D, A, DB, TB); } else if constexpr (NW == 176) { WGMMA_RS_M64N176K16_T(D, A, DB, TB); } else if constexpr (NW == 184) { WGMMA_RS_M64N184K16_T(D, A, DB, TB); } else if constexpr (NW == 192) { WGMMA_RS_M64N192K16_T(D, A, DB, TB); } else if constexpr (NW == 200) { WGMMA_RS_M64N200K16_T(D, A, DB, TB); } else if constexpr (NW == 208) { WGMMA_RS_M64N208K16_T(D, A, DB, TB); } else if constexpr (NW == 216) { WGMMA_RS_M64N216K16_T(D, A, DB, TB); } else if constexpr (NW == 224) { WGMMA_RS_M64N224K16_T(D, A, DB, TB); } else if constexpr (NW == 232) { WGMMA_RS_M64N232K16_T(D, A, DB, TB); } else if constexpr (NW == 240) { WGMMA_RS_M64N240K16_T(D, A, DB, TB); } else if constexpr (NW == 248) { WGMMA_RS_M64N248K16_T(D, A, DB, TB); } else if constexpr (NW == 256) { WGMMA_RS_M64N256K16_T(D, A, DB, TB); } } while (0) diff --git a/deploy/nemotron_native/kernels/codec/small/kernel_kf_codec.h b/deploy/nemotron_native/kernels/codec/small/kernel_kf_codec.h new file mode 100644 index 0000000..9a8b4d5 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/small/kernel_kf_codec.h @@ -0,0 +1,13 @@ +#pragma once +#include +struct CodecArgs { + const void *inp[41]; + const void *st[26]; + void *ns[26]; + void *waveform; + void *scratch; + int B; + int F; +}; +void codec_launch(const CodecArgs &a, void *stream); +size_t codec_scratch_bytes(int B, int F); diff --git a/deploy/nemotron_native/kernels/codec/small/legacy_codec.cuh b/deploy/nemotron_native/kernels/codec/small/legacy_codec.cuh new file mode 100644 index 0000000..cb128c4 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/small/legacy_codec.cuh @@ -0,0 +1,1388 @@ +#include "kernel_kf_codec.h" +#include +#include +#include +#include +using namespace nvcuda; +template __device__ void prepare_attribute(F,int) {} +#ifndef TC2_QPT +#define TC2_QPT 4 +#endif +#define DIVUP(a,b) (((a) + (b) - 1) / (b)) +__device__ __forceinline__ float half_snake(float v, int c, int half, + const float *__restrict__ alpha) { + if (c < half) { + const float a = alpha[c]; + const float s = __sinf(a * v); + return __fmaf_rn(s, s / (a + 1e-9f), v); + } + return v > 0.f ? v : 0.01f * v; +} +#ifndef WG_ST3 +#define WG_ST3 1 +#define WG_MB3 3 +#define WG_ST7 1 +#define WG_MB7 3 +#define WG_BN3 128 +#define WG_BN7 128 +#endif +#define PDL_ENTER() do { (void)0; if (threadIdx.x == 0) (void)0; } while (0) +#define PDL_TRIGGER() do { if (threadIdx.x == 0) (void)0; } while (0) +template +__device__ __forceinline__ void pdl_launch(KFn kern, dim3 grid, dim3 blk, size_t shmem, + cudaStream_t s, KArgs... args) { + native_invoke(kern, grid, blk, shmem, args...); +} +template +__device__ __forceinline__ void pdl_launch_pf(KFn kern, dim3 grid, dim3 blk, size_t shmem, + cudaStream_t s, const void *pf, size_t pfb, KArgs... args) { + native_invoke_pf(kern, grid, blk, shmem, pf, pfb, args...); +} +#include "wg_kernel.inc" +#include "wg2_kernel.inc" +#include "wg3_kernel.inc" +#include "wg4_step_widths.inc" +#include "wg5_step_widths.inc" +#include "repack3.inc" +#include "tail_kernel.inc" +template +__device__ __noinline__ void dense_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + PDL_TRIGGER(); + constexpr int BK = 16 * KT; + constexpr int NT = NWARPS * 32; + constexpr int WTM = BM / WM; + constexpr int WTN = BN / WN; + constexpr int FM = WTM / 16, FN = WTN / 16; + constexpr int LDA = BK + 8, LDB = BN + 8; + constexpr int GEMM_SM = BM * LDA + BK * LDB; + constexpr int EPI_SM = NWARPS * WTM * WTN * 2; + constexpr int SM_H = GEMM_SM > EPI_SM ? GEMM_SM : EPI_SM; + __half *smem=reinterpret_cast<__half*>(native_arena()); + char *colok=reinterpret_cast(native_aux()); + __half *As = smem; + __half *Bs = smem + BM * LDA; + float *Cw = reinterpret_cast(smem); + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int wm = warp / WN, wn = warp % WN; + const int m0 = blockIdx.x * BM; + const int n0 = blockIdx.y * BN; + const int H = KT - 1; + const int Kg = Cin * KT; + if (tid < BN) { + const int n = n0 + tid; + colok[tid] = (n < N && (n % L) < T) ? 1 : 0; + } + wmma::fragment acc[FM][FN]; +#pragma unroll + for (int i = 0; i < FM; ++i) +#pragma unroll + for (int j = 0; j < FN; ++j) wmma::fill_fragment(acc[i][j], 0.f); + constexpr int ATOT = BM * (BK / 8); + constexpr int AREG = (ATOT + NT - 1) / NT; + constexpr int BTOT = 16 * (BN / 8); + constexpr int BREG = (BTOT + NT - 1) / NT; + int4 ar[AREG]; + int4 br0[BREG], br1[BREG]; +#define CODEC_LOADG(kg0) { const int _ci0 = (kg0) / KT; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * NT; ar[r] = make_int4(0, 0, 0, 0); if (ATOT % NT == 0 || idx < ATOT) { const int m = idx / (BK / 8), v = idx % (BK / 8); const int co = m0 + m, kg = (kg0) + v * 8; if (co < Cout && kg < Kg) ar[r] = *reinterpret_cast(w + (size_t)co * Kg + kg); } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * NT; br0[r] = make_int4(0, 0, 0, 0); br1[r] = make_int4(0, 0, 0, 0); if (BTOT % NT == 0 || idx < BTOT) { const int cil = idx / (BN / 8), nc = idx % (BN / 8); const int ci = _ci0 + cil; if (ci < Cin) { const int4 *src = reinterpret_cast( in + (size_t)ci * SCin + n0 + nc * 8); br0[r] = src[0]; br1[r] = src[1]; } } } } + CODEC_LOADG(0) + for (int kg0 = 0; kg0 < Kg; kg0 += BK) { + __syncthreads(); +#pragma unroll + for (int r = 0; r < AREG; ++r) { + const int idx = tid + r * NT; + if (ATOT % NT == 0 || idx < ATOT) + *reinterpret_cast(As + (idx / (BK / 8)) * LDA + + (idx % (BK / 8)) * 8) = ar[r]; + } +#pragma unroll + for (int r = 0; r < BREG; ++r) { + const int idx = tid + r * NT; + if (BTOT % NT != 0 && idx >= BTOT) continue; + const int cil = idx / (BN / 8), nc = idx % (BN / 8); + unsigned int u[8]; + u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; + u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; +#pragma unroll + for (int k = 0; k < KT; ++k) { + const int kh = k >> 1; + int4 o; + if (k & 1) { + o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); + o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); + o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); + o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); + } else { + o.x = u[kh + 0]; o.y = u[kh + 1]; + o.z = u[kh + 2]; o.w = u[kh + 3]; + } + *reinterpret_cast(Bs + (cil * KT + k) * LDB + nc * 8) = o; + } + } + __syncthreads(); + if (kg0 + BK < Kg) CODEC_LOADG(kg0 + BK) +#pragma unroll + for (int kf = 0; kf < BK; kf += 16) { + wmma::fragment af[FM]; + wmma::fragment bf[FN]; +#pragma unroll + for (int i = 0; i < FM; ++i) + wmma::load_matrix_sync(af[i], As + (wm * WTM + i * 16) * LDA + kf, LDA); +#pragma unroll + for (int j = 0; j < FN; ++j) + wmma::load_matrix_sync(bf[j], Bs + kf * LDB + wn * WTN + j * 16, LDB); +#pragma unroll + for (int i = 0; i < FM; ++i) +#pragma unroll + for (int j = 0; j < FN; ++j) + wmma::mma_sync(acc[i][j], af[i], bf[j], acc[i][j]); + } + } +#undef CODEC_LOADG + __syncthreads(); + float *Cme = Cw + warp * WTM * WTN; +#pragma unroll + for (int i = 0; i < FM; ++i) +#pragma unroll + for (int j = 0; j < FN; ++j) + wmma::store_matrix_sync(Cme + (i * 16) * WTN + j * 16, acc[i][j], WTN, + wmma::mem_row_major); + __syncwarp(); + const int half = Cout >> 1; + const int lane = tid & 31; + for (int e = lane; e < WTM * WTN; e += 32) { + const int mm = e / WTN, nn = e % WTN; + const int co = m0 + wm * WTM + mm; + const int nl = wn * WTN + nn; + if (co >= Cout || !colok[nl]) continue; + const int n = n0 + nl; + float v = Cme[e] + bias[co]; + if (resid != nullptr) + v += __half2float(resid[(size_t)co * SCout + n + H]); + if (do_snake) v = half_snake(v, co, half, alpha); + out[(size_t)co * SCout + n + H] = __float2half(v); + } + if (blockIdx.y == 0) { + const int nbx = gridDim.x; + if (ns != nullptr) { + const int B = N / L; + const int total = Cin * B * H; + for (int idx = blockIdx.x * NT + tid; idx < total; idx += nbx * NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int B = N / L; + const int total = Cout * B * pad_len; + for (int idx = blockIdx.x * NT + tid; idx < total; idx += nbx * NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} +template +__device__ __noinline__ void transpose_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + __half *__restrict__ out, const __half *__restrict__ pad_src, + __half *__restrict__ ns, int Cin, int Cout, int Tin, + int Lin, int SCin, int jbase, int Lout, int SCout, int Hout, int pad_off, + int pad_len, int B) { + if constexpr (PDL) PDL_ENTER(); + constexpr int Ktaps = 2 * STRIDE; + constexpr int QPT = QPT_; + constexpr int VPT = QPT * STRIDE; + static_assert(QPT % 2 == 0, "half2 input loads need an even frame count"); + static_assert(VPT % 2 == 0, "half2 staging needs an even slot"); + float *ws=reinterpret_cast(native_aux()); + __half *sm=reinterpret_cast<__half*>(native_arena()); + const int co = blockIdx.y; + const int b = blockIdx.z; + const int tid = threadIdx.x, BLK = blockDim.x; + for (int i = tid; i < G * Ktaps; i += BLK) + ws[i] = __half2float(w[(size_t)(co * G + i / Ktaps) * Ktaps + (i % Ktaps)]); + __syncthreads(); + const float bs = bias[0 + co]; + const int half = Cout >> 1; + const bool snake = co < half; + const float av = snake ? alpha[co] : 0.f; + const float ri = snake ? __frcp_rn(av + 1e-9f) : 0.f; + const __half *ic = in + (size_t)(co * G) * SCin + b * Lin + jbase; + __half *oc = out + (size_t)co * SCout + b * Lout + Hout; + constexpr bool VEC16 = (VPT % 8) == 0; + const int ophase = + VEC16 ? 0 : (int)((((uintptr_t)oc) >> 1) & 7u); + __half *const sbase = sm + ophase; + const int hd = (8 - ophase) & 7; + const int lane = tid & 31; + for (int qb = 1 + blockIdx.x * BLK * QPT; qb <= Tin; + qb += gridDim.x * BLK * QPT) { + const int q = qb + QPT * tid; + const bool a0 = (q <= Tin); + float jq[QPT][G], jp0[G]; +#pragma unroll + for (int g = 0; g < G; ++g) { +#pragma unroll + for (int e2 = 0; e2 < QPT / 2; ++e2) { + const __half2 v = a0 ? *reinterpret_cast( + ic + (size_t)g * SCin + q + 2 * e2) + : __floats2half2_rn(0.f, 0.f); + jq[2 * e2][g] = __half2float(v.x); + jq[2 * e2 + 1][g] = __half2float(v.y); + } + } +#pragma unroll + for (int g = 0; g < G; ++g) { + const float up = __shfl_up_sync(0xffffffffu, jq[QPT - 1][g], 1); + jp0[g] = (lane == 0) + ? (a0 ? __half2float(ic[(size_t)g * SCin + q - 1]) : 0.f) + : up; + } + __half *const sp = sbase + tid * VPT; + __align__(16) __half vreg[VPT]; +#pragma unroll + for (int r = 0; r < STRIDE; ++r) { + float wc[G], wp[G]; +#pragma unroll + for (int g = 0; g < G; ++g) { + wc[g] = ws[g * Ktaps + r]; + wp[g] = ws[g * Ktaps + r + STRIDE]; + } +#pragma unroll + for (int e = 0; e < QPT; ++e) { + float v = bs; +#pragma unroll + for (int g = 0; g < G; ++g) { + const float cur = jq[e][g]; + const float prv = (e == 0) ? jp0[g] : jq[e - 1][g]; + v = __fmaf_rn(wc[g], cur, __fmaf_rn(wp[g], prv, v)); + } + if (snake) { + const float sn = __sinf(av * v); + v = __fmaf_rn(sn * sn, ri, v); + } else { + v = v > 0.f ? v : 0.01f * v; + } + vreg[e * STRIDE + r] = __float2half(v); + } + } + if constexpr (VEC16) { +#pragma unroll + for (int m = 0; m < VPT / 8; ++m) + *reinterpret_cast(sp + m * 8) = + *reinterpret_cast(vreg + m * 8); + } else { +#pragma unroll + for (int i = 0; i < VPT; ++i) sp[i] = vreg[i]; + } + __syncthreads(); + const int nval = min(BLK * QPT, Tin - qb + 1); + const int tot = nval * STRIDE; + __half *dst = oc + (size_t)(qb - 1) * STRIDE; + if constexpr (VEC16) { + const int tot2 = tot >> 1; + for (int i = tid; i < tot2; i += BLK) + reinterpret_cast<__half2 *>(dst)[i] = + reinterpret_cast(sbase)[i]; + if ((tot & 1) && tid == 0) dst[tot - 1] = sbase[tot - 1]; + } else { + const int head = hd < tot ? hd : tot; + for (int i = tid; i < head; i += BLK) dst[i] = sbase[i]; + const int nb4 = (tot - head) >> 3; + const int4 *s4 = reinterpret_cast(sbase + head); + int4 *d4 = reinterpret_cast(dst + head); + for (int i = tid; i < nb4; i += BLK) d4[i] = s4[i]; + for (int i = head + (nb4 << 3) + tid; i < tot; i += BLK) + dst[i] = sbase[i]; + } + __syncthreads(); + } + if (blockIdx.x == 0) { + if (ns != nullptr) { + for (int ci = co + tid * Cout; ci < Cin; ci += BLK * Cout) + ns[(size_t)b * Cin + ci] = + in[(size_t)ci * SCin + b * Lin + jbase + native_true_tokens(b,Tin)]; + } + if (pad_src != nullptr) { + for (int j = tid; j < pad_len; j += BLK) + out[(size_t)co * SCout + b * Lout + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } +} +template +__device__ __noinline__ void transpose_conv_early_bpack_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + __half *__restrict__ out, const __half *__restrict__ pad_src, + __half *__restrict__ ns, int Cin, int Cout, int Tin, + int Lin, int SCin, int jbase, int Lout, int SCout, int Hout, int pad_off, + int pad_len, int B) { + if constexpr (PDL) PDL_ENTER(); + static_assert(BPACK <= 32); + static_assert(LPB == 8 || LPB == 16 || LPB == 32); + constexpr int Ktaps = 2 * STRIDE; + constexpr int WARP = 32; + constexpr int NT = WARP * BPACK; + constexpr int BPB = WARP / LPB; + float *ws=reinterpret_cast(native_aux()); + __half *sm=reinterpret_cast<__half*>(native_arena()); + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int lane = tid & 31; + const int sub = lane % LPB; + const int bsub = lane / LPB; + const int co = blockIdx.y; + const int b = blockIdx.z * (BPACK * BPB) + warp * BPB + bsub; + const bool active = b < B; + for (int i = tid; i < G * Ktaps; i += NT) + ws[i] = __half2float(w[(size_t)(co * G + i / Ktaps) * Ktaps + + (i % Ktaps)]); + __syncthreads(); + const int qb = 1 + blockIdx.x * LPB; + const int q = qb + sub; + __half *sp = sm + (warp * WARP + lane) * STRIDE; + const bool live = active && q <= Tin; + { + const float bs = bias[co]; + const int half = Cout >> 1; + const bool snake = co < half; + const float av = snake ? alpha[co] : 0.f; + const float ri = snake ? __frcp_rn(av + 1e-9f) : 0.f; + const __half *ic = in + (size_t)(co * G) * SCin + b * Lin + jbase; + float jq[G], jp[G]; +#pragma unroll + for (int g = 0; g < G; ++g) + jq[g] = live ? __half2float(ic[(size_t)g * SCin + q]) : 0.f; +#pragma unroll + for (int g = 0; g < G; ++g) { + const float up = __shfl_up_sync(0xffffffffu, jq[g], 1, LPB); + jp[g] = (sub == 0) + ? (live ? __half2float(ic[(size_t)g * SCin + q - 1]) : 0.f) + : up; + } + if (live) +#pragma unroll + for (int r = 0; r < STRIDE; ++r) { + float v = bs; +#pragma unroll + for (int g = 0; g < G; ++g) + v = __fmaf_rn(ws[g * Ktaps + r], jq[g], + __fmaf_rn(ws[g * Ktaps + r + STRIDE], jp[g], v)); + if (snake) { + const float sn = __sinf(av * v); + v = __fmaf_rn(sn * sn, ri, v); + } else { + v = v > 0.f ? v : 0.01f * v; + } + sp[r] = __float2half(v); + } + } + __syncwarp(); + if (active) { + const int nval = min(LPB, Tin - (qb - 1)); + const int tot2 = (nval * STRIDE) >> 1; + const __half *src = sm + (warp * WARP + bsub * LPB) * STRIDE; + __half *dst = out + (size_t)co * SCout + b * Lout + Hout + + (size_t)(qb - 1) * STRIDE; + for (int i = sub; i < tot2; i += LPB) + reinterpret_cast<__half2 *>(dst)[i] = + reinterpret_cast(src)[i]; + if ((nval * STRIDE) & 1) + if (sub == 0) dst[nval * STRIDE - 1] = src[nval * STRIDE - 1]; + } + if (blockIdx.x == 0 && active) { + if (ns != nullptr) { + for (int ci = co + sub * Cout; ci < Cin; ci += LPB * Cout) + ns[(size_t)b * Cin + ci] = + in[(size_t)ci * SCin + b * Lin + jbase + native_true_tokens(b,Tin)]; + } + if (pad_src != nullptr) { + for (int j = sub; j < pad_len; j += LPB) + out[(size_t)co * SCout + b * Lout + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } +} +template +__device__ __noinline__ void transpose_conv_bpack2_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + __half *__restrict__ out, const __half *__restrict__ pad_src, + __half *__restrict__ ns, int Cin, int Cout, int Tin, + int Lin, int SCin, int jbase, int Lout, int SCout, int Hout, int pad_off, + int pad_len, int B) { + if constexpr (PDL) PDL_ENTER(); + static_assert(BPACK <= 32); + static_assert(LPB == 8 || LPB == 16 || LPB == 32); + constexpr int Ktaps = 2 * STRIDE; + constexpr int WARP = 32; + constexpr int NT = WARP * BPACK; + constexpr int BPB = WARP / LPB; + constexpr int QPT = 2; + constexpr int VPT = QPT * STRIDE; + float *ws=reinterpret_cast(native_aux()); + __half *sm=reinterpret_cast<__half*>(native_arena()); + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int lane = tid & 31; + const int sub = lane % LPB; + const int bsub = lane / LPB; + const int co = blockIdx.y; + const int b = blockIdx.z * (BPACK * BPB) + warp * BPB + bsub; + const bool active = b < B; + for (int i = tid; i < G * Ktaps; i += NT) + ws[i] = __half2float(w[(size_t)(co * G + i / Ktaps) * Ktaps + + (i % Ktaps)]); + __syncthreads(); + const int qb = 1 + blockIdx.x * (LPB * QPT); + const int q = qb + sub * QPT; + __half *sp = sm + (warp * WARP + lane) * VPT; + const bool a0 = active && (q <= Tin); + { + const float bs = bias[co]; + const int half = Cout >> 1; + const bool snake = co < half; + const float av = snake ? alpha[co] : 0.f; + const float ri = snake ? __frcp_rn(av + 1e-9f) : 0.f; + const __half *ic = in + (size_t)(co * G) * SCin + b * Lin + jbase; + float jq0[G], jq1[G], jp0[G]; +#pragma unroll + for (int g = 0; g < G; ++g) { + const __half2 v = a0 ? *reinterpret_cast( + ic + (size_t)g * SCin + q) + : __floats2half2_rn(0.f, 0.f); + jq0[g] = __half2float(v.x); + jq1[g] = __half2float(v.y); + } +#pragma unroll + for (int g = 0; g < G; ++g) { + const float up = __shfl_up_sync(0xffffffffu, jq1[g], 1, LPB); + jp0[g] = (sub == 0) + ? (a0 ? __half2float(ic[(size_t)g * SCin + q - 1]) + : 0.f) + : up; + } + if (a0) { + float vv[VPT]; +#pragma unroll + for (int r = 0; r < STRIDE; ++r) { + float wc[G], wp[G]; +#pragma unroll + for (int g = 0; g < G; ++g) { + wc[g] = ws[g * Ktaps + r]; + wp[g] = ws[g * Ktaps + r + STRIDE]; + } +#pragma unroll + for (int e = 0; e < QPT; ++e) { + float v = bs; +#pragma unroll + for (int g = 0; g < G; ++g) { + const float cur = (e == 0) ? jq0[g] : jq1[g]; + const float prv = (e == 0) ? jp0[g] : jq0[g]; + v = __fmaf_rn(wc[g], cur, __fmaf_rn(wp[g], prv, v)); + } + if (snake) { + const float sn = __sinf(av * v); + v = __fmaf_rn(sn * sn, ri, v); + } else { + v = v > 0.f ? v : 0.01f * v; + } + vv[e * STRIDE + r] = v; + } + } + __half2 *sp2 = reinterpret_cast<__half2 *>(sp); +#pragma unroll + for (int i = 0; i < VPT / 2; ++i) + sp2[i] = __floats2half2_rn(vv[2 * i], vv[2 * i + 1]); + } + } + __syncwarp(); + if (active) { + const int nval = min(LPB * QPT, Tin - (qb - 1)); + const int tot = nval * STRIDE; + const __half *src = sm + (warp * WARP + bsub * LPB) * VPT; + __half *dst = out + (size_t)co * SCout + b * Lout + Hout + + (size_t)(qb - 1) * STRIDE; + for (int i = sub; i < (tot >> 1); i += LPB) + reinterpret_cast<__half2 *>(dst)[i] = + reinterpret_cast(src)[i]; + if (tot & 1) + if (sub == 0) dst[tot - 1] = src[tot - 1]; + } + if (blockIdx.x == 0 && active) { + if (ns != nullptr) { + for (int ci = co + sub * Cout; ci < Cin; ci += LPB * Cout) + ns[(size_t)b * Cin + ci] = + in[(size_t)ci * SCin + b * Lin + jbase + native_true_tokens(b,Tin)]; + } + if (pad_src != nullptr) { + for (int j = sub; j < pad_len; j += LPB) + out[(size_t)co * SCout + b * Lout + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } +} +#define TCONV_BPACK(G,STRIDE,BPACK,LPB,TIN,COUT,BB,...) pdl_launch(transpose_conv_early_bpack_kernel, dim3(DIVUP(TIN, LPB), COUT, DIVUP(BB, BPACK * (32 / LPB))), dim3(32 * BPACK), (size_t)(32 * (BPACK)) * (STRIDE) * 2 + 64, s, __VA_ARGS__) +#define TCONV_BP2(G,STRIDE,LPB,TIN,COUT,BB,...) pdl_launch(transpose_conv_bpack2_kernel, dim3(DIVUP(TIN, (LPB) * 2), COUT, DIVUP(BB, 32)), dim3(32 * (LPB)), (size_t)(32 * (LPB)) * (4 * (STRIDE)) + 64, s, __VA_ARGS__) +#define TCONV_BPACK_DISPATCH1(G,STRIDE,BPACK,TIN,COUT,BB,...) do { if ((TIN) >= 32) TCONV_BP2(G, STRIDE, 16, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) >= 16) TCONV_BP2(G, STRIDE, 8, TIN, COUT, BB, __VA_ARGS__); else TCONV_BPACK(G, STRIDE, BPACK, 16, TIN, COUT, BB, __VA_ARGS__); } while (0) +#define TCONV_BPACK_DISPATCH(G,STRIDE,BPACK,TIN,COUT,BB,...) do { if ((TIN) >= 32) TCONV_BP2(G, STRIDE, 16, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) >= 16) TCONV_BP2(G, STRIDE, 8, TIN, COUT, BB, __VA_ARGS__); else TCONV_BPACK(G, STRIDE, BPACK, 16, TIN, COUT, BB, __VA_ARGS__); } while (0) +#define TCONV_BPACK_DISPATCH8(G,STRIDE,BPACK,TIN,COUT,BB,...) do { if ((TIN) >= 32) TCONV_BP2(G, STRIDE, 16, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) >= 16) TCONV_BP2(G, STRIDE, 8, TIN, COUT, BB, __VA_ARGS__); else if ((TIN) <= 8) TCONV_BPACK(G, STRIDE, BPACK, 8, TIN, COUT, BB, __VA_ARGS__); else TCONV_BPACK(G, STRIDE, BPACK, 16, TIN, COUT, BB, __VA_ARGS__); } while (0) +__device__ __noinline__ void pack_latent_kernel(const __half *__restrict__ latent, + const __half *__restrict__ state, + __half *__restrict__ out, int Cin, int T, + int L, int SC, int H, int B) { + const int total = Cin * B * L; + for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; + idx += gridDim.x * blockDim.x) { + const int j = idx % L, r = idx / L; + const int b = r % B, ci = r / B; + out[(size_t)ci * SC + b * L + j] = + (j < H) ? state[((size_t)b * Cin + ci) * H + j] + : latent[((size_t)b * Cin + ci) * T + (j - H)]; + } +} +#define PC_VEC 8 +__device__ __noinline__ void post_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, __half *__restrict__ y, + __half *__restrict__ ns, int Cin, int T, int L, int SC, int KT, int B, + int N) { + float *pw=reinterpret_cast(native_arena()); + for (int i = threadIdx.x; i < Cin * KT; i += 256) pw[i] = __half2float(w[i]); + __syncthreads(); + const int H = KT - 1; + const float bs = bias[0]; + for (int n = (blockIdx.x * 256 + threadIdx.x) * PC_VEC; n < N; + n += gridDim.x * 256 * PC_VEC) { + float v[PC_VEC]; +#pragma unroll + for (int i = 0; i < PC_VEC; ++i) v[i] = bs; + for (int ci = 0; ci < Cin; ++ci) { + const __half *xr = in + (size_t)ci * SC + n; + const float *wr = pw + ci * KT; + const int4 lo = *reinterpret_cast(xr); + const int4 hi = *reinterpret_cast(xr + 8); + unsigned u[8]; + u[0] = lo.x; u[1] = lo.y; u[2] = lo.z; u[3] = lo.w; + u[4] = hi.x; u[5] = hi.y; u[6] = hi.z; u[7] = hi.w; + float x[PC_VEC + 6]; +#pragma unroll + for (int j = 0; j < (PC_VEC + 6) / 2; ++j) { + const float2 f = __half22float2( + *reinterpret_cast(&u[j])); + x[2 * j] = f.x; + x[2 * j + 1] = f.y; + } +#pragma unroll + for (int i = 0; i < PC_VEC; ++i) + for (int k = 0; k < 7; ++k) + if (k < KT) v[i] = __fmaf_rn(wr[k], x[i + k], v[i]); + } + const int b0 = n / L, t0 = n - (n / L) * L; +#pragma unroll + for (int i = 0; i < PC_VEC; ++i) { + int t = t0 + i, b = b0; + if (t >= L) { t -= L; b = b0 + 1; } + if (t < T && b < B) + y[(size_t)b * T + t] = + __float2half(fminf(fmaxf(v[i], -1.f), 1.f)); + } + } + if (ns != nullptr) { + const int total = B * Cin * H; + for (int idx = blockIdx.x * 256 + threadIdx.x; idx < total; + idx += gridDim.x * 256) { + const int j = idx % H, r = idx / H; + const int ci = r % Cin, b = r / Cin; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SC + b * L + native_true_tokens(b,T) + j]; + } + } +} +#define RP_JOBS 6 +struct RepackSet { + const __half *src[RP_JOBS]; + __half *dst[RP_JOBS]; + int Cout[RP_JOBS], CoutP[RP_JOBS], Cin[RP_JOBS]; + int base[RP_JOBS], nch[RP_JOBS], rs[RP_JOBS]; + int k3base; + const __half *k3src[2]; + __half *k3dst[2]; +}; +__host__ __device__ inline int native_RP_COUT(int j) {const int v[6]={768, 384, 256, 128, 64, 32};return v[j];} +__host__ __device__ inline int native_RP_CIN(int j) {const int v[6]={384, 768, 128, 256, 32, 64};return v[j];} +__host__ __device__ inline int native_RP_COUTP(int j) {const int v[6]={768, 384, 256, 128, 128, 128};return v[j];} +template +__device__ __noinline__ void repack_kernel(RepackSet js) { + const int bx = blockIdx.x; + int job, job_bx; + if (bx < 1152) { + job = 0; job_bx = 0; + } else if (bx < 2304) { + job = 1; job_bx = 1152; + } else if (bx < 2432) { + job = 2; job_bx = 2304; + } else if (bx < 2560) { + job = 3; job_bx = 2432; + } else if (bx < 2576) { + job = 4; job_bx = 2560; + } else { + job = 5; job_bx = 2576; + } + const int Cout = js.Cout[job], CoutP = js.CoutP[job], Cin = js.Cin[job]; + const int idx = (bx - job_bx) * blockDim.x + threadIdx.x; + const int co = idx / Cin, ci = idx - co * Cin; + const int c = ci / CIB, cl = ci - c * CIB; + __half *d = js.dst[job] + (size_t)c * (KT * CoutP * CIB) + + (size_t)(co >> 3) * (CIB * 8) + (cl >> 3) * 64 + (co & 7) * 8 + + (cl & 7); + const __half *sp = js.src[job] + (size_t)(co * Cin + ci) * KT; + const bool live = co < Cout; +#pragma unroll + for (int k = 0; k < KT; ++k) + d[(size_t)k * CoutP * CIB] = live ? sp[k] : __float2half(0.f); +} +template +__device__ __noinline__ void repack_rs_kernel( + const RepackSet &js, const __half *__restrict__ latent, const __half *__restrict__ latent_state, + __half *__restrict__ latent_out, int latent_C, int latent_T, int latent_L, + int latent_SC, int latent_H, int latent_B) { + PDL_TRIGGER(); + constexpr int CO_TILE = 16; + constexpr int CI_TILE = 64; + constexpr int ROW_H = CI_TILE * KT + 8; + constexpr int ROW4 = CI_TILE * KT / 8; + __half *tile=reinterpret_cast<__half*>(native_arena()); + const int bx = blockIdx.x; + { + const int idx = bx + threadIdx.x * gridDim.x; + const int total = latent_C * latent_B * latent_L; + if (idx < total) { + const int j = idx % latent_L, r = idx / latent_L; + const int b = r % latent_B, ci = r / latent_B; + latent_out[(size_t)ci * latent_SC + b * latent_L + j] = + (j < latent_H) + ? latent_state[((size_t)b * latent_C + ci) * latent_H + j] + : latent[((size_t)b * latent_C + ci) * latent_T + + (j - latent_H)]; + } + } + if (bx >= js.k3base) { + constexpr int K3 = 3; + constexpr int R3_H = CI_TILE * K3 + 8; + constexpr int R3_4 = CI_TILE * K3 / 8; + const int r = bx - js.k3base; + const int conv = r / 1152; + const int loc = r - conv * 1152; + const int Cout3 = conv ? 768 : 1536; + const int Cin3 = conv ? 1536 : 768; + const int nci3 = Cin3 / CI_TILE; + const int cot3 = loc / nci3; + if (cot3 * CO_TILE >= Cout3) return; + const int co03 = cot3 * CO_TILE; + const int ci03 = (loc - cot3 * nci3) * CI_TILE; + const int NCH3 = Cin3 >> 4; + const __half *const s3 = js.k3src[conv]; + __half *const d3 = js.k3dst[conv]; + const int t3 = threadIdx.x; +#pragma unroll + for (int n = 0; n < CO_TILE * R3_4 / 128; ++n) { + const int i = t3 + n * 128; + const int co = i / R3_4, q = i - co * R3_4; + reinterpret_cast(tile + co * R3_H)[q] = + *reinterpret_cast( + s3 + ((size_t)(co03 + co) * Cin3 + ci03) * K3 + q * 8); + } + __syncthreads(); + const int w3 = t3 >> 5, l3 = t3 & 31; + const int rA3 = l3 >> 2, q43 = l3 & 3; + const int cb3 = co03 >> 4; + __half *const db3 = d3 + ((size_t)(cb3 >> 2) * NCH3 * K3) * 1024 + + ((cb3 & 3) * 32 + l3) * 8; +#pragma unroll + for (int u = w3; u < (CI_TILE / 16) * K3; u += 4) { + const int cl = u / K3, tap = u - cl * K3; + const int cgi = (ci03 >> 4) + cl; + const __half *sa = tile + rA3 * R3_H + (cl * 16 + 2 * q43) * K3 + tap; + const __half *sb = sa + 8 * R3_H; + uint32_t o[4]; + o[0] = __half_as_ushort(sa[0]) | + ((uint32_t)__half_as_ushort(sa[K3]) << 16); + o[1] = __half_as_ushort(sb[0]) | + ((uint32_t)__half_as_ushort(sb[K3]) << 16); + o[2] = __half_as_ushort(sa[8 * K3]) | + ((uint32_t)__half_as_ushort(sa[9 * K3]) << 16); + o[3] = __half_as_ushort(sb[8 * K3]) | + ((uint32_t)__half_as_ushort(sb[9 * K3]) << 16); + *reinterpret_cast(db3 + (size_t)(cgi * K3 + tap) * 1024) = + make_int4(o[0], o[1], o[2], o[3]); + } + return; + } + int job = 0; +#pragma unroll + for (int q = 1; q < RP_JOBS; ++q) if (bx >= js.base[q]) job = q; + const int local = bx - js.base[job]; + const int Cout = js.Cout[job], Cin = js.Cin[job]; + const int nci = (Cin + CI_TILE - 1) / CI_TILE; + const int cot = local / nci; + const int co0 = cot * CO_TILE; + const int ci0 = (local - cot * nci) * CI_TILE; + const int NCH = js.nch[job]; + const int t = threadIdx.x; + const int valid4 = (min(Cin - ci0, CI_TILE) * KT) >> 3; + for (int i = t; i < CO_TILE * ROW4; i += 128) { + const int co = i / ROW4, q = i - co * ROW4; + int4 v = make_int4(0, 0, 0, 0); + if (co0 + co < Cout && q < valid4) + v = *reinterpret_cast( + js.src[job] + ((size_t)(co0 + co) * Cin + ci0) * KT + q * 8); + reinterpret_cast(tile + co * ROW_H)[q] = v; + } + __syncthreads(); + const int warp = t >> 5, lane = t & 31; + const int rA = lane >> 2, q4 = lane & 3; + const int cb = co0 >> 4; + const int mtile = cb >> 2, wsub = cb & 3; + __half *const dbase = js.dst[job] + + ((size_t)mtile * NCH * KT) * 1024 + (wsub * 32 + lane) * 8; + if (!js.rs[job]) { + const int CoutP = js.CoutP[job]; + constexpr int UN2 = 16 * (CI_TILE / 8) * KT; +#pragma unroll + for (int n = 0; n < UN2 / 128; ++n) { + const int u = t + n * 128; + const int col = u & 15, cig = (u >> 4) & (CI_TILE / 8 - 1); + const int tap = u / (16 * (CI_TILE / 8)); + const int gci = ci0 + cig * 8; + if (gci >= Cin) continue; + const int gco = co0 + col; + uint32_t o[4]; + const __half *sp = tile + col * ROW_H; +#pragma unroll + for (int e = 0; e < 4; ++e) + o[e] = __half_as_ushort(sp[(cig * 8 + 2 * e) * KT + tap]) | + ((uint32_t)__half_as_ushort( + sp[(cig * 8 + 2 * e + 1) * KT + tap]) << 16); + __half *d = js.dst[job] + + (size_t)(gci >> 4) * ((size_t)KT * CoutP * 16) + + (size_t)tap * (CoutP * 16) + (size_t)(gco >> 3) * 128 + + ((gci >> 3) & 1) * 64 + (gco & 7) * 8; + *reinterpret_cast(d) = make_int4(o[0], o[1], o[2], o[3]); + } + return; + } + constexpr int UNITS = (CI_TILE / 16) * KT; +#pragma unroll + for (int u = warp; u < UNITS; u += 4) { + const int cl = u / KT, tap = u - cl * KT; + const int cgi = (ci0 >> 4) + cl; + if (cgi >= NCH) continue; + const __half *sa = tile + rA * ROW_H + (cl * 16 + 2 * q4) * KT + tap; + const __half *sb = sa + 8 * ROW_H; + uint32_t o[4]; + o[0] = __half_as_ushort(sa[0]) | + ((uint32_t)__half_as_ushort(sa[KT]) << 16); + o[1] = __half_as_ushort(sb[0]) | + ((uint32_t)__half_as_ushort(sb[KT]) << 16); + o[2] = __half_as_ushort(sa[8 * KT]) | + ((uint32_t)__half_as_ushort(sa[9 * KT]) << 16); + o[3] = __half_as_ushort(sb[8 * KT]) | + ((uint32_t)__half_as_ushort(sb[9 * KT]) << 16); + *reinterpret_cast(dbase + (size_t)(cgi * KT + tap) * 1024) = + make_int4(o[0], o[1], o[2], o[3]); + } +} +__device__ __noinline__ void repack3_rs_kernel( + const __half *__restrict__ win, const __half *__restrict__ wsk, + __half *__restrict__ din, __half *__restrict__ dsk, + const __half *__restrict__ asrc, __half *__restrict__ adst, int SCA, + int ncol, int ntile) { + PDL_TRIGGER(); + if (blockIdx.y == 12) { + if (adst == nullptr) return; + for (int u = blockIdx.x; u < 96 * ntile; u += gridDim.x) { + const int g = u / ntile; + const int col = (u - g * ntile) * 128 + (int)threadIdx.x; + if (col >= ncol) continue; + __half v[8]; +#pragma unroll + for (int h = 0; h < 8; ++h) + v[h] = asrc[(size_t)(g * 8 + h) * SCA + col]; + *reinterpret_cast(adst + (size_t)g * ((size_t)SCA * 8) + + (size_t)col * 8) = + *reinterpret_cast(v); + } + return; + } + constexpr int KT = 3; + constexpr int CO_TILE = 16; + constexpr int CI_TILE = 64; + constexpr int ROW_H = CI_TILE * KT + 8; + constexpr int ROW4 = CI_TILE * KT / 8; + __half *tile=reinterpret_cast<__half*>(native_arena()); + const int jy = blockIdx.y; + const int conv = jy & 1, b = 5 - (jy >> 1); + const int Cout = conv ? 768 : 1536; + const int Cin = conv ? 1536 : 768; + const __half *src = (conv ? wsk : win) + (size_t)(b + 1) * Cout * Cin * KT; + __half *dst = (conv ? dsk : din) + (size_t)b * ((size_t)Cout * Cin * KT); + const int nci = Cin / CI_TILE; + const int cot = blockIdx.x / nci; + const int co0 = cot * CO_TILE; + const int ci0 = (blockIdx.x - cot * nci) * CI_TILE; + const int NCH = Cin >> 4; + const int t = threadIdx.x; +#pragma unroll + for (int n = 0; n < CO_TILE * ROW4 / 128; ++n) { + const int i = t + n * 128; + const int co = i / ROW4, q = i - co * ROW4; + reinterpret_cast(tile + co * ROW_H)[q] = + *reinterpret_cast( + src + ((size_t)(co0 + co) * Cin + ci0) * KT + q * 8); + } + __syncthreads(); + const int warp = t >> 5, lane = t & 31; + const int rA = lane >> 2, q4 = lane & 3; + const int cb = co0 >> 4; + __half *const dbase = dst + ((size_t)(cb >> 2) * NCH * KT) * 1024 + + ((cb & 3) * 32 + lane) * 8; +#pragma unroll + for (int u = warp; u < (CI_TILE / 16) * KT; u += 4) { + const int cl = u / KT, tap = u - cl * KT; + const int cgi = (ci0 >> 4) + cl; + const __half *sa = tile + rA * ROW_H + (cl * 16 + 2 * q4) * KT + tap; + const __half *sb = sa + 8 * ROW_H; + uint32_t o[4]; + o[0] = __half_as_ushort(sa[0]) | + ((uint32_t)__half_as_ushort(sa[KT]) << 16); + o[1] = __half_as_ushort(sb[0]) | + ((uint32_t)__half_as_ushort(sb[KT]) << 16); + o[2] = __half_as_ushort(sa[8 * KT]) | + ((uint32_t)__half_as_ushort(sa[9 * KT]) << 16); + o[3] = __half_as_ushort(sb[8 * KT]) | + ((uint32_t)__half_as_ushort(sb[9 * KT]) << 16); + *reinterpret_cast(dbase + (size_t)(cgi * KT + tap) * 1024) = + make_int4(o[0], o[1], o[2], o[3]); + } +} +__device__ __noinline__ void repack_activation_only(const __half *asrc,__half *adst,int SCA,int ncol,int ntile){ + if (adst == nullptr) return; + for (int u = blockIdx.x; u < 96 * ntile; u += gridDim.x) { + const int g = u / ntile; + const int col = (u - g * ntile) * 128 + (int)threadIdx.x; + if (col >= ncol) continue; + __half v[8]; +#pragma unroll + for (int h = 0; h < 8; ++h) + v[h] = asrc[(size_t)(g * 8 + h) * SCA + col]; + *reinterpret_cast(adst + (size_t)g * ((size_t)SCA * 8) + + (size_t)col * 8) = + *reinterpret_cast(v); + } + return; +} +namespace { +struct Stack { + int T, L, N, SC, H; +}; +inline __host__ __device__ int tc_blk(int Tin, int qpt = 2) { + int best = 256, bw = 1 << 30; + for (int b = 256; b >= 32; b >>= 1) { + const int per = qpt * b; + const int waste = (Tin + per - 1) / per * per - Tin; + if (waste < bw) { bw = waste; best = b; } + } + return best; +} +inline __host__ __device__ int roundup(int a, int b) { return (a + b - 1) / b * b; } +__host__ __device__ Stack mk(int B, int T, int K) { + Stack s; + s.T = T; + s.H = K - 1; + s.L = s.H + T; + s.N = B * s.L; + s.SC = roundup(s.N, 128) + 128; + return s; +} +struct Plan { + Stack s0, s1, s2, s3, s4; + __half *LB, *X0, *M0, *Y0; + __half *A, *Ac, *Bb, *M1; + __half *P, *Q, *M2; + __half *R, *S, *M3; + __half *U; + __half *W7[RP_JOBS]; + __half *W3i, *W3s; + size_t total; +}; +static constexpr int kTuneB = 32; +inline __host__ __device__ int nref(const Stack &st) { return kTuneB * st.L; } +inline __host__ __device__ bool use_rp3i(int, int F) { return (size_t)kTuneB * (4 * F + 2) > 400; } +inline __host__ __device__ bool use_rp3s(int, int F) { return (size_t)kTuneB * (4 * F + 2) > 400; } +__host__ __device__ Plan build_plan(int B, int F, __half *base) { + Plan p; + p.s0 = mk(B, 2 * F, 3); + p.s1 = mk(B, 4 * F, 3); + p.s2 = mk(B, 36 * F, 7); + p.s3 = mk(B, 252 * F, 7); + p.s4 = mk(B, 1764 * F, 7); + size_t off = 0; + auto alloc = [&](int ch, const Stack &s) -> __half * { + __half *r = base ? base + off : nullptr; + off += (size_t)ch * s.SC; + return r; + }; + p.LB = alloc(40, p.s0); + p.X0 = alloc(768, p.s0); + p.M0 = alloc(1536, p.s0); + p.Y0 = alloc(768, p.s0); + p.A = alloc(768, p.s1); + p.Ac = use_rp3i(B, F) ? alloc(768, p.s1) : nullptr; + p.Bb = alloc(768, p.s1); + p.M1 = alloc(1536, p.s1); + p.P = alloc(384, p.s2); + p.Q = alloc(384, p.s2); + p.M2 = alloc(768, p.s2); + p.R = alloc(128, p.s3); + p.S = alloc(128, p.s3); + p.M3 = alloc(256, p.s3); + p.U = alloc(32, p.s4); + for (int j = 0; j < RP_JOBS; ++j) { + p.W7[j] = base ? base + off : nullptr; + off += (size_t)native_RP_COUTP(j) * native_RP_CIN(j) * 7; + } + p.W3i = nullptr; + if (use_rp3i(B, F)) { + p.W3i = base ? base + off + (size_t)3 * 1536 * 768 : nullptr; + off += (size_t)7 * 3 * 1536 * 768; + } + p.W3s = nullptr; + if (use_rp3s(B, F)) { + p.W3s = base ? base + off + (size_t)3 * 768 * 1536 : nullptr; + off += (size_t)7 * 3 * 768 * 1536; + } + p.total = off; + return p; +} +__host__ __device__ const __half *H16(const void *p) { return static_cast(p); } +__host__ __device__ const float *F32(const void *p) { return static_cast(p); } +} +size_t codec_scratch_bytes(int B, int F) { + return build_plan(B, F, nullptr).total * sizeof(__half); +} +template +__device__ __forceinline__ void wg_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, bool prep = false) { + constexpr int WG_BM = 64 * NWG, WG_NT = 128 * NWG, WG_CLD = WG_BN + 8; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int GEMM_B = STAGES * (A_E + B_E) * 2; + constexpr int EPI_B = WG_BM * WG_CLD * 4; + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute(wg_conv_kernel, SH); + return; + } + native_invoke(wg_conv_kernel,dim3(dim3(DIVUP(Cout, WG_BM), DIVUP(st.N, WG_BN))),dim3(WG_NT),SH, + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg2_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, + bool prep = false) { + constexpr int BK = CIB * KT; + constexpr int WG_BM = 64 * MSPLIT; + constexpr int WG_NT = 128 * MSPLIT * KSPLIT * NSPLIT; + constexpr int WG_BNT = WG_BN * NSPLIT; + constexpr int NCRAW = WG_BNT / 4; + constexpr int EPN = + (WG_NT % NCRAW == 0) ? WG_BNT : (((WG_BNT + 127) / 128) * 128); + constexpr int WG_CLD = EPN + 8; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int GEMM_B = KSPLIT * (SA * A_E + SB * NSPLIT * B_E) * 2; + constexpr int EPI_B = KSPLIT * WG_BM * WG_CLD * 4; + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute( + wg2_conv_kernel, + SH); + return; + } + native_invoke_pf(wg2_conv_kernel,dim3(dim3(DIVUP(Cout, WG_BM), DIVUP(st.N, WG_BNT))),dim3(WG_NT),SH, + w, (size_t)KT * Cin * Cout * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg3_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, + bool prep = false) { + constexpr int WG_BM = 64 * NWG, WG_NT = 128 * NWG * NSG; + constexpr int WG_BNT = WG_BN * NSG, WG_CLD = WG_BNT + 8; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int GEMM_B = STAGES * (A_E + NSG * B_E) * 2; + constexpr int EPI_B = WG_BM * WG_CLD * 4; + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute(wg3_conv_kernel, SH); + return; + } + native_invoke(wg3_conv_kernel,dim3(dim3(DIVUP(Cout, WG_BM), DIVUP(st.N, WG_BNT))),dim3(WG_NT),SH, + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg4_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, int CoutP, + const Stack &st, int SCin, int do_snake, int pad_off, + int pad_len, bool prep = false) { + constexpr int NT = 128 * MSPLIT * NSPLIT; + constexpr int BM = 64 * MSPLIT, BNT = WG_BN * NSPLIT; + constexpr int NBR = BNT + KT - 1; + constexpr int NBRP = BC8 ? ((NBR + 7) & ~7) : ((NBR + 31) & ~31); + constexpr int NCRAW = BNT / 4; + constexpr int EPN = + (NT % NCRAW == 0) ? BNT : (((BNT + 127) / 128) * 128); + constexpr int B_E = CIB * NBRP; + constexpr int GEMM_B = STAGES * B_E * 2; + constexpr int EPI_B = BM * (EPN + 8) * (PREACT ? 2 : 4); + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + if (prep) { + prepare_attribute( + wg4_conv_kernel, + SH); + return; + } + pdl_launch_pf(wg4_conv_kernel, + dim3(DIVUP(Cout, BM), DIVUP(st.N, BNT)), dim3(NT), SH, s, + w, (size_t)KT * CoutP * Cin * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, + st.T, st.L, SCin, st.SC, st.N, do_snake, pad_off, pad_len); +} +template +__device__ __forceinline__ void wg5_launch(cudaStream_t s, const __half *in, const __half *w, + const float *bias, const float *alpha, + const __half *resid, __half *out, const __half *pad_src, + __half *ns, int Cin, int Cout, const Stack &st, int SCin, + int do_snake, int pad_off, int pad_len, + bool prep = false) { + constexpr int BM = 64 * MSPLIT; + constexpr int NT = 128 * MSPLIT * KSPLIT; + constexpr int NCRAW = WG_BN / 4; + constexpr int EPN = (NT % NCRAW == 0) ? WG_BN : (((WG_BN + 127) / 128) * 128); + constexpr int NBRP = BC8 ? (((WG_BN + 2) + 7) & ~7) + : (((WG_BN + 2) + 31) & ~31); + constexpr int B_E = CIB * NBRP; + constexpr int GEMM_B = KSPLIT * SB * B_E * 2; + constexpr int EPI_COLS = OC8 ? WG_BN : EPN; + constexpr int EPI_CLD_R = EPI_COLS + 8; + constexpr int EPI_CLD = (EPI_CLD_R % 16 == 0) ? EPI_CLD_R + 4 : EPI_CLD_R; + constexpr int EPI_B = KSPLIT * BM * EPI_CLD * (OC8 ? 2 : 4); + constexpr int SH = GEMM_B > EPI_B ? GEMM_B : EPI_B; + static_assert(SH <= 227 * 1024, "wg5 shared over budget"); + if (prep) { + prepare_attribute( + wg5_conv_kernel, + SH); + return; + } + if constexpr (PDL) { + cudaLaunchConfig_t cfg = {}; + cudaLaunchAttribute attr[1] = {}; + native_config_grid(cfg) = dim3(DIVUP(Cout, BM), DIVUP(st.N, WG_BN)); + native_config_threads(cfg) = dim3(NT); + cfg.dynamicSmemBytes = SH; + cfg.stream = s; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = 1; + cfg.attrs = attr; + cfg.numAttrs = 1; + native_cudaLaunchKernelEx_pf( + &cfg, + wg5_conv_kernel, + w, (size_t)3 * Cin * Cout * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, + SCin, st.SC, st.N, do_snake, pad_off, pad_len); + } else { + native_invoke_pf(wg5_conv_kernel,dim3(dim3(DIVUP(Cout, BM), DIVUP(st.N, WG_BN))),dim3(NT),SH, + w, (size_t)3 * Cin * Cout * sizeof(__half), + in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, + st.L, SCin, st.SC, st.N, do_snake, pad_off, pad_len); + } +} +#ifndef D7W +#define D7W 96 +#endif +#define D7_M2_C8I(BN,LATE) wg4_launch<7, 16, 2, 1, D7W, 4, 1, false, 2, true, false, LATE>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_M2P_C8O(BN,LATE) wg4_launch<7, 16, 2, 1, D7W, 4, 1, true, 2, false, true, LATE>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_M2N2(BN) wg4_launch<7, 16, 2, 2, BN, 4, 1, false, 2>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_M1(BN) wg4_launch<7, 16, 1, 1, BN, 4, 2, false, 2>( s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, SCin, 1, pad_off, pad_len, prep) +#define D7_C8I(W) \ + wg4_launch<7, 16, 2, 1, W, 4, 1, false, 2, true, false, false>( \ + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, \ + SCin, 1, pad_off, pad_len, prep) +#define D7_C8O(W) \ + wg4_launch<7, 16, 2, 1, W, 4, 1, true, 2, false, true, false>( \ + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, CoutP, st, \ + SCin, 1, pad_off, pad_len, prep) +template +__device__ __forceinline__ int d7_band_cost(int N, int mtiles, int resident) { + const int v = mtiles * DIVUP(N, BN); + const int sub = (DIVUP(v, resident) >= 2) ? 2 : 1; + return DIVUP(DIVUP(v, sub), resident) * sub * (128 + BN + 6); +} +__device__ __forceinline__ int d7_band(int N, int mtiles, int resident) { + int best = 64, bc = d7_band_cost<64>(N, mtiles, resident); + const int c96 = d7_band_cost<96>(N, mtiles, resident); + if (c96 < bc) { bc = c96; best = 96; } + if (d7_band_cost<112>(N, mtiles, resident) < bc) best = 112; + return best; +} +__device__ __forceinline__ void dense7(cudaStream_t s, int job, const __half *in, const __half *w, + const float *bias, const float *alpha, const __half *resid, + __half *out, const __half *pad_src, __half *ns, + const Stack &st, int SCin, int pad_off, int pad_len, + bool prep = false) { + const int Cin = native_RP_CIN(job), Cout = native_RP_COUT(job), CoutP = native_RP_COUTP(job); + const int b_ = d7_band(st.N, DIVUP(Cout, 128), native_resident()); + if ((job & 1) == 0) { + if (b_ == 64) D7_C8O(64); + else if (b_ == 96) D7_C8O(96); + else D7_C8O(112); + } else { + if (b_ == 64) D7_C8I(64); + else if (b_ == 96) D7_C8I(96); + else D7_C8I(112); + } +} +__device__ __forceinline__ int wg5_band_skip(int N, int mtiles, int resident) { + const int wmin = DIVUP(mtiles * DIVUP(N, 104), resident); + if (DIVUP(mtiles * DIVUP(N, 24), resident) == wmin) return 24; + if (DIVUP(mtiles * DIVUP(N, 40), resident) == wmin) return 40; + if (DIVUP(mtiles * DIVUP(N, 48), resident) == wmin) return 48; + if (DIVUP(mtiles * DIVUP(N, 64), resident) == wmin) return 64; + return 104; +} +__device__ __forceinline__ int wg5_band_main_long(int N, int mtiles, + int resident) { + const int wmin = DIVUP(mtiles * DIVUP(N, 104), resident); + if (DIVUP(mtiles * DIVUP(N, 64), resident) == wmin) return 64; + return 104; +} +__device__ __forceinline__ int wg5_band_main_short(int N, int mtiles, + int resident) { + const int wmin = DIVUP(mtiles * DIVUP(N, 96), resident); + if (DIVUP(mtiles * DIVUP(N, 56), resident) == wmin) return 56; + if (DIVUP(mtiles * DIVUP(N, 80), resident) == wmin) return 80; + return 96; +} +#define DENSE_LAUNCH(KT,BM,BN,NW,WM,WN) native_invoke(dense_conv_kernel,dim3(dim3(DIVUP(Cout, BM), DIVUP(st.N, BN))),dim3(NW * 32),0, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st.T, st.L, SCin, st.SC, st.N, do_snake, pad_off, pad_len) +__device__ __forceinline__ void dense(cudaStream_t s, int KT, const __half *in, const __half *w, + const float *bias, const float *alpha, const __half *resid, + __half *out, const __half *pad_src, __half *ns, int Cin, + int Cout, const Stack &st, int SCin, int do_snake, + int pad_off, int pad_len, const __half *w3 = nullptr, + int layout_mode = 0, bool prep = false) { + if (KT == 3 && w3 != nullptr) { + if (Cin == 1536) { +#define WG5S(BN,O,R,PDL) wg5_launch<16, 1, 4, BN, 4, 3, 1, 1, true, O, R, 1536, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5SW(BN,O,R,PDL) wg5_launch<32, 1, 2, BN, 4, 3, 1, 1, true, O, R, 1536, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5S3(O,R) do { const int n_ = nref(st); if (n_ > 500 && n_ <= 640) { WG5SW(56, O, R, false); } else { const int b_ = wg5_band_skip(st.N, DIVUP(Cout, 64), native_resident()); if (b_ == 24) WG5S(24, O, R, true); else if (b_ == 40) WG5S(40, O, R, true); else if (b_ == 48) WG5S(48, O, R, true); else if (b_ == 64) WG5S(64, O, R, true); else WG5S(104, O, R, true); } } while (0) + if (layout_mode == 1) WG5S3(true, false); + else if (layout_mode == 2) WG5S3(true, true ); + else if (layout_mode == 3) WG5S3(false, true ); + else WG5S3(false, false); +#undef WG5S3 +#undef WG5SW +#undef WG5S + } else { +#define WG5I(BN,BC,MS,PDL) wg5_launch<16, MS, 2, BN, 4, 3, 1, 1, BC, true, false, 768, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5I4(BN,BC,PDL) wg5_launch<16, 1, 4, BN, 4, 3, 1, 1, BC, true, false, 768, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5I2W(BN,BC,PDL) wg5_launch<32, 1, 2, BN, 4, 3, 1, 1, BC, true, false, 768, PDL>( s, in, w3, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, SCin, do_snake, pad_off, pad_len, prep) +#define WG5I3(BC) do { const int n_ = nref(st); if (n_ > 640) { const int b_ = wg5_band_main_long(st.N, DIVUP(Cout, 128), native_resident()); if (b_ == 64) WG5I(64, BC, 2, true); else WG5I(104, BC, 2, true); } else if (n_ > 500) { WG5I2W(64, BC, false); } else { const int b_ = wg5_band_main_short(st.N, DIVUP(Cout, 64), native_resident()); if (b_ == 56) WG5I4(56, BC, true); else if (b_ == 80) WG5I4(80, BC, true); else WG5I4(96, BC, true); } } while (0) + if (layout_mode == 1) WG5I3(true); + else WG5I3(false); +#undef WG5I3 +#undef WG5I2W +#undef WG5I4 +#undef WG5I + } + } else if (KT == 3) { + const int n_ = nref(st); + if (Cin == 1536 && Cout == 768) { + if (n_ <= 704) + wg2_launch<3, 16, 1, 4, 1, 64, 2, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else + wg2_launch<3, 16, 1, 4, 1, 128, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + } else if (Cin == 768 && Cout == 1536) { + if (n_ > 1000) + wg2_launch<3, 16, 1, 2, 1, 128, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ > 640 && n_ <= 768) + wg2_launch<3, 16, 1, 2, 1, 64, 2, 2, 2>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ <= 320) + wg2_launch<3, 16, 1, 4, 1, 64, 2, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ > 384 && n_ <= 512) + wg2_launch<3, 16, 1, 4, 1, 112, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else if (n_ > 512 && n_ <= 640) + wg2_launch<3, 16, 1, 4, 1, 120, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + else + wg2_launch<3, 16, 1, 4, 1, 128, 4, 2, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, + st, SCin, do_snake, pad_off, pad_len, prep); + } else { + wg2_launch<3, 16, 1, 1, 1, 64, 3, 2, 2>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, + SCin, do_snake, pad_off, pad_len, prep); + } + } else if (Cout >= 128 && + ((Cin >= 384 && st.T == 288) || + (Cin < 384 && st.T == 2016))) { + wg3_launch<7, 112, 2, 2, 1, 128, 2>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, + SCin, do_snake, pad_off, pad_len); + } else { + wg3_launch<7, 112, 1, 4, 1, 128, 1>( + s, in, w, bias, alpha, resid, out, pad_src, ns, Cin, Cout, st, + SCin, do_snake, pad_off, pad_len); + } +} +__device__ void codec_program(const CodecArgs &a, const Plan &p, const RepackSet &js, int repack_grid) { + cudaStream_t s = nullptr; + const int B = a.B; + const void *const *W = a.inp; + const __half *const *ST = reinterpret_cast(a.st); + __half *const *NS = reinterpret_cast<__half *const *>(a.ns); + { + native_invoke(pack_latent_kernel,dim3(DIVUP(40*B*p.s0.L,256)),dim3(256),0, + H16(W[0]), ST[0], p.LB, 40, p.s0.T, p.s0.L, p.s0.SC, 2, B); + } + dense(s, 3, p.LB, H16(W[1]), F32(W[2]), nullptr, nullptr, p.X0, ST[1], NS[0], + 40, 768, p.s0, p.s0.SC, 0, 0, 2); + const __half *k3iw = H16(W[3]); + const float *k3ib = F32(W[4]); + const float *k3ia = F32(W[5]); + const __half *k3sw = H16(W[6]); + const float *k3sb = F32(W[7]); + const float *k3sa = F32(W[8]); + const size_t IW = (size_t)1536 * 768 * 3, SW = (size_t)768 * 1536 * 3; + dense(s, 3, p.X0, k3iw, k3ib, k3ia, nullptr, p.M0, ST[2], NS[1], 768, 1536, + p.s0, p.s0.SC, 1, 0, 2, p.W3i ? p.W3i - IW : nullptr, 0); + dense(s, 3, p.M0, k3sw, k3sb, k3sa, p.X0, p.Y0, ST[3], NS[2], 1536, 768, + p.s0, p.s0.SC, 1, 1, 1, p.W3s ? p.W3s - SW : nullptr, 0); + TCONV_BPACK_DISPATCH8(1, 2, 8, p.s0.T, 768, B, + p.Y0, H16(W[9]), F32(W[10]), F32(W[11]), p.A, ST[4], NS[3], 768, 768, + p.s0.T, p.s0.L, p.s0.SC, 1, p.s1.L, p.s1.SC, 2, 0, 2, B); + if (p.W3i != nullptr) + native_invoke(repack_activation_only,dim3(min(1152,96*DIVUP(p.s1.SC,128))),dim3(128),0, + p.A, p.Ac, p.s1.SC, p.s1.SC, DIVUP(p.s1.SC,128)); + __half *cur = p.A, *nxt = p.Bb; + for (int i = 0; i < 6; ++i) { + const int bi = i + 1; + const __half *ain = (i == 0 && p.Ac) ? p.Ac : cur; + dense(s, 3, ain, k3iw + bi * IW, k3ib + bi * 1536, k3ia + bi * 768, + nullptr, p.M1, ST[5 + 2 * i], NS[4 + 2 * i], 768, 1536, p.s1, + p.s1.SC, 1, 0, 2, p.W3i ? p.W3i + (bi - 1) * IW : nullptr, + (i > 0 || p.Ac) ? 1 : 0); + const __half *ps = (i < 5) ? ST[6 + 2 * i] : ST[16]; + const int po = (i < 5) ? 0 : 1, pl = 1 + (i < 5); + dense(s, 3, p.M1, k3sw + bi * SW, k3sb + bi * 768, k3sa + bi * 384, cur, + nxt, ps, NS[5 + 2 * i], 1536, 768, p.s1, p.s1.SC, 1, po, pl, + p.W3s ? p.W3s + (bi - 1) * SW : nullptr, + i == 0 ? 1 : (i < 5 ? 2 : 3)); + __half *t = cur; cur = nxt; nxt = t; + } + TCONV_BPACK_DISPATCH1(2, 9, 16, p.s1.T, 384, B, + cur, H16(W[12]), F32(W[13]), F32(W[14]), p.P, ST[17], NS[16], 768, 384, + p.s1.T, p.s1.L, p.s1.SC, 1, p.s2.L, p.s2.SC, 6, 0, 6, B); + dense7(s, 0, p.P, p.W7[0], F32(W[22]), F32(W[23]), nullptr, p.M2, ST[18], + NS[17], p.s2, p.s2.SC, 0, 6); + dense7(s, 1, p.M2, p.W7[1], F32(W[25]), F32(W[26]), p.P, p.Q, ST[19], + NS[18], p.s2, p.s2.SC, 5, 1); + TCONV_BPACK_DISPATCH(3, 7, 16, p.s2.T, 128, B, + p.Q, H16(W[15]), F32(W[16]), F32(W[17]), p.R, ST[20], NS[19], 384, 128, + p.s2.T, p.s2.L, p.s2.SC, 5, p.s3.L, p.s3.SC, 6, 0, 6, B); + dense7(s, 2, p.R, p.W7[2], F32(W[28]), F32(W[29]), nullptr, p.M3, ST[21], + NS[20], p.s3, p.s3.SC, 0, 6); + dense7(s, 3, p.M3, p.W7[3], F32(W[31]), F32(W[32]), p.R, p.S, ST[22], + NS[21], p.s3, p.s3.SC, 5, 1); + if (p.s3.T >= 1500) + pdl_launch(transpose_conv_kernel<4, 7, 8, true>, dim3(1, 32, B), + dim3(tc_blk(p.s3.T, 8)), + (size_t)tc_blk(p.s3.T, 8) * (8 * 7) * 2 + 64, s, + p.S, H16(W[18]), F32(W[19]), F32(W[20]), p.U, ST[23], + NS[22], 128, 32, p.s3.T, p.s3.L, p.s3.SC, 5, p.s4.L, + p.s4.SC, 6, 0, 6, B); + else + pdl_launch(transpose_conv_kernel<4, 7, 4, true>, dim3(1, 32, B), + dim3(tc_blk(p.s3.T, 4)), + (size_t)tc_blk(p.s3.T, 4) * (4 * 7) * 2 + 64, s, + p.S, H16(W[18]), F32(W[19]), F32(W[20]), p.U, ST[23], + NS[22], 128, 32, p.s3.T, p.s3.L, p.s3.SC, 5, p.s4.L, + p.s4.SC, 6, 0, 6, B); + { + const int ntile = DIVUP(p.s4.N, TL_BN); + const int grid = ntile < 132 ? ntile : 132; + pdl_launch_pf(tail_fused_kernel, dim3(grid), dim3(TL_NT), TL_SHARED, s, + p.U, (size_t)32 * p.s4.SC * sizeof(__half), + p.U, p.W7[4], p.W7[5], H16(W[39]), F32(W[34]), F32(W[35]), + F32(W[37]), F32(W[38]), F32(W[40]), ST[24], ST[25], + static_cast<__half *>(a.waveform), NS[23], NS[24], NS[25], B, + p.s4.T, p.s4.L, p.s4.SC, p.s4.N, ntile); + } +} diff --git a/deploy/nemotron_native/kernels/codec/small/native_stage1.cu b/deploy/nemotron_native/kernels/codec/small/native_stage1.cu new file mode 100644 index 0000000..7118404 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/small/native_stage1.cu @@ -0,0 +1,310 @@ +#include "native_stage1.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#ifndef NATIVE_CODEC_PREFETCH +#define NATIVE_CODEC_PREFETCH 1 +#endif +#ifndef NATIVE_CODEC_PACK_SUBCTAS +#define NATIVE_CODEC_PACK_SUBCTAS 1 +#endif +namespace cg=cooperative_groups; +static_assert(SMALL_NATIVE_BATCH == 2 || SMALL_NATIVE_BATCH == 4 || SMALL_NATIVE_BATCH == 8); +constexpr int PHYSICAL_THREADS=512,MAIN_ARENA=227840,AUX_BYTES=4096,META_OFFSET=231936,TOTAL_SHARED=232192; +constexpr int SLICE_BYTES=MAIN_ARENA+AUX_BYTES; +constexpr int MAX_SUBCTA=8; +struct NativeMeta { + dim3 grid,threads; + uint3 block[MAX_SUBCTA]; + NativeStage1Status *status;const int32_t *lengths;int padded_frames; + int submask,tshift,tmask,slice; + int resident; +}; +static_assert(sizeof(NativeMeta)<=TOTAL_SHARED-META_OFFSET); +__device__ __forceinline__ char *native_base(){extern __shared__ __align__(16) char arena[];return arena;} +__device__ __forceinline__ NativeMeta *native_meta(){return reinterpret_cast(native_base()+META_OFFSET);} +__device__ __forceinline__ int native_subid(){const NativeMeta *m=native_meta();return (int)((threadIdx.x>>m->tshift)&(unsigned)m->submask);} +__device__ __forceinline__ char *native_arena(){const NativeMeta *m=native_meta();return native_base()+native_subid()*m->slice;} +__device__ __forceinline__ char *native_aux(){const NativeMeta *m=native_meta();return native_base()+native_subid()*m->slice+(m->slice-AUX_BYTES);} +__device__ __forceinline__ uint3 native_tid(){uint3 r;r.x=threadIdx.x&(unsigned)native_meta()->tmask;r.y=0u;r.z=0u;return r;} +__device__ __forceinline__ dim3 native_block(){const uint3 b=native_meta()->block[native_subid()];return dim3(b.x,b.y,b.z);} +__device__ __forceinline__ dim3 native_grid(){return native_meta()->grid;} +__device__ __forceinline__ dim3 native_threads(){return native_meta()->threads;} +__device__ __forceinline__ int native_resident(){return native_meta()->resident;} +__device__ __forceinline__ void native_local_sync(){ + const NativeMeta *m=native_meta();unsigned n=m->threads.x,id=15u-(unsigned)native_subid(); + asm volatile("bar.sync %0, %1;"::"r"(id),"r"(n):"memory"); +} +__device__ __forceinline__ dim3 &native_config_grid(cudaLaunchConfig_t &c){return c.gridDim;} +__device__ __forceinline__ dim3 &native_config_threads(cudaLaunchConfig_t &c){return c.blockDim;} +__device__ __forceinline__ int native_true_tokens(int row,int padded_tokens){ + const NativeMeta *m=native_meta(); + return (padded_tokens/m->padded_frames)*m->lengths[row]; +} +__device__ __forceinline__ void native_warm(const char *p,size_t n,int k,int nk,int me){ + if(p==nullptr||n<16||nk<=0)return; + const size_t a=(size_t)p&15u;p-=a;n=(n+a)&~(size_t)15; + const uint4 *q=reinterpret_cast(p); + const size_t cnt=n>>4,step=(size_t)nk*PHYSICAL_THREADS; + uint32_t acc=0u; + for(size_t i=(size_t)k*PHYSICAL_THREADS+(size_t)me;i __device__ __forceinline__ void native_invoke_pf(F fn,dim3 g,dim3 t,size_t shared,const void *pf,size_t pfb,A&&... args){ + const int me=threadIdx.x,owner=blockIdx.x,owners=gridDim.x; + const cg::grid_group all=cg::this_grid();NativeMeta *meta=native_meta(); + if(t.x>512||t.y!=1||t.z!=1||t.x%32||shared>MAIN_ARENA){ + if(me==0&&owner==0)meta->status->code=10; + all.sync();return; + } + int sub=1,tshift=9; + const unsigned tx=t.x; + const int vtotal=(int)(g.x*g.y*g.z); + if(NATIVE_CODEC_PACK_SUBCTAS && (tx&(tx-1))==0u){ + int s=0;while((1u<MAX_SUBCTA)cand=MAX_SUBCTA; + const int want=(vtotal+owners-1)/owners; + while(cand>1&&cand>want)cand>>=1; + while(cand>1&&shared>(size_t)(SLICE_BYTES/cand-AUX_BYTES))cand>>=1; + sub=cand; + } + const int slice=(SLICE_BYTES/sub)&~127; + if(me==0){meta->grid=g;meta->threads=t;meta->submask=sub-1;meta->tshift=tshift; + meta->tmask=(int)((1u<slice=slice;} + const int total=vtotal,grp=me>>tshift; + const bool holder=me=busy&&pf!=nullptr)native_warm((const char*)pf,pfb,owner-busy,owners-busy,me); + for(int base=owner*sub;baseblock[grp]=make_uint3(i%g.x,(i/g.x)%g.y,i/(g.x*g.y)); + __syncthreads(); + if(holder&&i(args)...); + asm volatile("cp.async.wait_all;":::"memory"); + if(t.x>=128)asm volatile("wgmma.wait_group.sync.aligned 0;":::"memory"); + } + __syncthreads(); + } + all.sync(); +} +template __device__ __forceinline__ void native_invoke(F fn,dim3 g,dim3 t,size_t shared,A&&... args){ + native_invoke_pf(fn,g,t,shared,nullptr,(size_t)0,static_cast(args)...); +} +template __device__ __forceinline__ void native_cudaLaunchKernelEx(cudaLaunchConfig_t *c,F fn,A... args){ + native_invoke(fn,c->gridDim,c->blockDim,c->dynamicSmemBytes,args...); +} +template __device__ __forceinline__ void native_cudaLaunchKernelEx_pf(cudaLaunchConfig_t *c,F fn,const void *pf,size_t pfb,A... args){ + native_invoke_pf(fn,c->gridDim,c->blockDim,c->dynamicSmemBytes,pf,pfb,args...); +} +#define blockIdx native_block() +#define gridDim native_grid() +#define blockDim native_threads() +#define threadIdx native_tid() +#define __syncthreads() native_local_sync() +#include "legacy_codec.cuh" +#undef __syncthreads +#undef threadIdx +#undef blockDim +#undef gridDim +#undef blockIdx +namespace { +struct FastArgs { + CodecArgs codec;Plan plan;RepackSet repack;int repack_grid;int input_frames;int blocks; + const int64_t *codes;const int32_t *lengths,*slots; + const __half *state[26];__half *next[26];float *pcm;int32_t *samples; + NativeStage1Status *status;int nslots; + int channels[26],histories[26]; +}; +static_assert(sizeof(FastArgs)<=4096,"precomputed kernel argument must fit the conservative parameter budget"); +struct Handle {FastArgs args{};Plan plans[17]{};RepackSet repacks[17]{};int repack_grids[17]{};void *scratch=nullptr;size_t bytes=0;void *immutable=nullptr;size_t immutable_bytes=0;int setup_launches=0;std::vector packed;int max_frames=0,blocks=0,device=-1;}; +thread_local std::string error; +bool check(cudaError_t e,const char *where){if(e==cudaSuccess)return true;error=std::string(where)+": "+cudaGetErrorString(e);return false;} +__device__ void reject(NativeStage1Status *s,int code,int row,int frame,int detail){ + if(atomicCAS(&s->code,0,code)==0){s->row=row;s->frame=frame;s->detail=detail;} +} +__global__ __launch_bounds__(512,1) void stage1_fast_persistent(const __grid_constant__ FastArgs a){ + const auto all=cg::this_grid();size_t me=(size_t)blockIdx.x*512+threadIdx.x,step=(size_t)gridDim.x*512; + if(threadIdx.x==0){native_meta()->status=a.status;native_meta()->lengths=a.lengths;native_meta()->padded_frames=a.codec.F;native_meta()->resident=a.blocks;} + if(me==0)*a.status=NativeStage1Status{0,-1,-1,0}; + all.sync();const int B=a.codec.B,F=a.codec.F,inputF=a.input_frames; + for(int row=me;rowinputF)reject(a.status,8,row,-1,n); + if(slot<0||slot>=a.nslots)reject(a.status,2,row,-1,slot); + for(int j=0;j=0&&n<=inputF)?n*1764:0; + } + for(size_t i=me;i<(size_t)B*inputF;i+=step){ + int row=i/inputF,t=i%inputF;bool any=false; + for(int c=0;c<16;++c){int64_t v=a.codes[i*16+c];if(v<0||v>=1024)reject(a.status,4,row,t,c);any|=v!=0;} + if(t=a.lengths[row]&&any)reject(a.status,6,row,t,0); + } + for(int p=0;p<26;++p)for(size_t i=me;i<(size_t)a.nslots*a.channels[p]*a.histories[p];i+=step)a.next[p][i]=a.state[p][i]; + for(size_t i=me;i<(size_t)B*inputF*1764;i+=step)a.pcm[i]=0.f; + all.sync();if(a.status->code)return; + for(int p=0;p<26;++p){ + const size_t stride=(size_t)a.channels[p]*a.histories[p]; + __half *dest=(__half*)a.codec.st[p]; + for(size_t i=me;i<(size_t)B*stride;i+=step)dest[i]=a.state[p][(size_t)a.slots[i/stride]*stride+i%stride]; + } + __half *latent=(__half*)a.codec.inp[0];int T=2*F; + for(size_t i=me;i<(size_t)B*40*T;i+=step){ + int t=i%T,c=(i/T)%40,row=i/((size_t)40*T);float v=0.f; + if(t/2>(2*(c%5)))&3)-2)*.5f;} + latent[i]=__float2half_rn(v); + } + all.sync(); + codec_program(a.codec,a.plan,a.repack,a.repack_grid); + if(a.status->code)return; + for(int p=0;p<26;++p){ + const size_t stride=(size_t)a.channels[p]*a.histories[p];const __half *src=(const __half*)a.codec.ns[p]; + for(size_t i=me;i<(size_t)B*stride;i+=step){int row=i/stride;if(a.lengths[row])a.next[p][(size_t)a.slots[row]*stride+i%stride]=src[i];} + } + const __half *wave=(const __half*)a.codec.waveform; + for(size_t i=me;i<(size_t)B*inputF*1764;i+=step){ + const int row=i/(inputF*1764),t=i%(inputF*1764); + if(tstatus=a.status;native_meta()->resident=a.blocks;} + if(blockIdx.x==0&&threadIdx.x==0)*a.status=NativeStage1Status{0,-1,-1,0}; + all.sync(); + native_invoke(repack_rs_kernel<7>,dim3(a.repack_grid),dim3(128),MAIN_ARENA, + a.repack,H16(a.codec.inp[0]),H16(a.codec.st[0]),a.plan.LB, + 40,a.plan.s0.T,a.plan.s0.L,a.plan.s0.SC,2,a.codec.B); + native_invoke(repack3_rs_kernel,dim3(1152,12),dim3(128),MAIN_ARENA, + H16(a.codec.inp[3]),H16(a.codec.inp[6]),a.plan.W3i,a.plan.W3s, + a.plan.A,a.plan.Ac,a.plan.s1.SC,a.plan.s1.SC,DIVUP(a.plan.s1.SC,128)); +} +bool stack(Handle *h,int slot,const NativeStage1Weights *w,int first,int field,size_t elements){ + const int ids[7]={first,first+3,first+5,first+7,first+9,first+11,first+13}; + const size_t stride=elements*(field==0?sizeof(__half):sizeof(float));void *ptr=nullptr; + if(!check(cudaMalloc(&ptr,7*stride),"packed-weight allocation"))return false; + h->packed.push_back(ptr);h->args.codec.inp[slot]=ptr; + for(int i=0;i<7;++i){const void *src=field==0?w[ids[i]].weight:field==1?(const void*)w[ids[i]].bias:(const void*)w[ids[i]].alpha; + if(!check(cudaMemcpy((char*)ptr+i*stride,src,stride,cudaMemcpyDeviceToDevice),"immutable weight stack"))return false;} + return true; +} +void release(Handle *h){if(h->immutable)cudaFree(h->immutable);if(h->scratch)cudaFree(h->scratch);for(void *p:h->packed)cudaFree(p);delete h;} +} +extern "C" void *native_stage1_create(int b,int maxf,int slots,int blocks,const NativeStage1Weights *w){ + error.clear();if((b!=SMALL_NATIVE_BATCH)||maxf<1||maxf>16||slots4096||blocks<0||!w){error="invalid create arguments";return nullptr;} + for(int p=0;p<26;++p)if(!w[p].weight||!w[p].bias||(p!=0&&p!=25&&!w[p].alpha)){error="missing weight pointer";return nullptr;} + Handle *h=new(std::nothrow)Handle;if(!h){error="host allocation";return nullptr;} + h->max_frames=maxf;h->args.nslots=slots;h->args.codec.B=b; + const int cs[26]={40,768,1536,768,768,1536,768,1536,768,1536,768,1536,768,1536,768,1536,768,384,768,384,128,256,128,32,64,32}; + const int hs[26]={2,2,2,1,2,2,2,2,2,2,2,2,2,2,2,2,1,6,6,1,6,6,1,6,6,6}; + cudaDeviceProp prop{};int occupancy=0,coop=0; + if(!check(cudaGetDevice(&h->device),"device")||!check(cudaGetDeviceProperties(&prop,h->device),"properties")|| + !check(cudaDeviceGetAttribute(&coop,cudaDevAttrCooperativeLaunch,h->device),"cooperative support")|| + !check(cudaFuncSetAttribute(stage1_fast_persistent,cudaFuncAttributeMaxDynamicSharedMemorySize,TOTAL_SHARED),"shared-memory optin")|| + !check(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&occupancy,stage1_fast_persistent,512,TOTAL_SHARED),"occupancy")) {release(h);return nullptr;} + if(prop.major!=9||prop.minor!=0||!coop||occupancy<1){error="requires resident SM90 cooperative kernel";release(h);return nullptr;} + h->blocks=blocks?blocks:occupancy*prop.multiProcessorCount; + if(h->blocks>occupancy*prop.multiProcessorCount){error="grid exceeds cooperative occupancy";release(h);return nullptr;} + h->args.blocks=h->blocks; + const int capacityF=maxf<=3?3:(maxf<=5?5:std::max(8,maxf)); + size_t decoder=codec_scratch_bytes(b,capacityF),halfs=(size_t)b*40*2*capacityF+(size_t)b*1764*capacityF; + for(int p=0;p<26;++p){h->args.channels[p]=cs[p];h->args.histories[p]=hs[p];halfs+=2*(size_t)b*cs[p]*hs[p];} + h->bytes=decoder+halfs*sizeof(__half); + if(!check(cudaMalloc(&h->scratch,h->bytes),"scratch allocation")){release(h);return nullptr;} + if(!check(cudaMemset(h->scratch,0,h->bytes),"initial scratch padding")){release(h);return nullptr;} + h->args.codec.scratch=h->scratch;__half *q=(__half*)((char*)h->scratch+decoder); + h->args.codec.inp[0]=q;q+=(size_t)b*40*2*capacityF;h->args.codec.waveform=q;q+=(size_t)b*1764*capacityF; + for(int p=0;p<26;++p){h->args.codec.st[p]=q;q+=(size_t)b*cs[p]*hs[p];h->args.codec.ns[p]=q;q+=(size_t)b*cs[p]*hs[p];} + const void **in=h->args.codec.inp; + in[1]=w[0].weight;in[2]=w[0].bias; + if(!stack(h,3,w,1,0,(size_t)1536*768*3)||!stack(h,4,w,1,1,1536)||!stack(h,5,w,1,2,768)|| + !stack(h,6,w,2,0,(size_t)768*1536*3)||!stack(h,7,w,2,1,768)||!stack(h,8,w,2,2,384)){release(h);return nullptr;} + const int upops[4]={3,16,19,22};for(int i=0;i<4;++i){in[9+3*i]=w[upops[i]].weight;in[10+3*i]=w[upops[i]].bias;in[11+3*i]=w[upops[i]].alpha;} + for(int i=0;i<3;++i){int p=17+3*i;in[21+6*i]=w[p].weight;in[22+6*i]=w[p].bias;in[23+6*i]=w[p].alpha; + in[24+6*i]=w[p+1].weight;in[25+6*i]=w[p+1].bias;in[26+6*i]=w[p+1].alpha;} + in[39]=w[25].weight;in[40]=w[25].bias; + __half *immutable_w7[6]{};__half *immutable_w3i=nullptr,*immutable_w3s=nullptr; + if(capacityF>=3){ + size_t wh=0;for(int j=0;j<6;++j)wh+=(size_t)native_RP_COUTP(j)*native_RP_CIN(j)*7; + const size_t IW=(size_t)3*1536*768;wh+=14*IW; + h->immutable_bytes=wh*sizeof(__half); + if(!check(cudaMalloc(&h->immutable,h->immutable_bytes+sizeof(NativeStage1Status)),"immutable weight allocation")){release(h);return nullptr;} + if(!check(cudaMemset(h->immutable,0,h->immutable_bytes+sizeof(NativeStage1Status)),"immutable padding initialization")){release(h);return nullptr;} + __half *q=(__half*)h->immutable; + for(int j=0;j<6;++j){immutable_w7[j]=q;q+=(size_t)native_RP_COUTP(j)*native_RP_CIN(j)*7;} + immutable_w3i=q+IW;q+=7*IW;immutable_w3s=q+IW;q+=7*IW; + if((char*)q!=(char*)h->immutable+h->immutable_bytes){error="immutable layout size mismatch";release(h);return nullptr;} + } + for(int f:{3,5,8}) if(f<=capacityF){ + Plan p=build_plan(b,f,static_cast<__half*>(h->scratch)); + for(int j=0;j<6;++j)p.W7[j]=immutable_w7[j];p.W3i=immutable_w3i;p.W3s=immutable_w3s; + const void *const *W=h->args.codec.inp; + RepackSet js{}; + const void *SRC[RP_JOBS] = {W[21], W[24], W[27], W[30], W[33], W[36]}; + for (int j = 0; j < RP_JOBS; ++j) { + js.src[j] = H16(SRC[j]); + js.dst[j] = p.W7[j]; + js.Cout[j] = native_RP_COUT(j); + js.CoutP[j] = native_RP_COUTP(j); + js.Cin[j] = native_RP_CIN(j); + } + int acc = 0; + for (int j = 0; j < RP_JOBS; ++j) { + js.base[j] = acc; + js.nch[j] = native_RP_CIN(j) / 16; + js.rs[j] = (j < 4) ? 1 : 0; + acc += (native_RP_COUTP(j) / 16) * DIVUP(native_RP_CIN(j), 64); + } + js.k3base = acc; + js.k3src[0] = H16(W[3]); + js.k3src[1] = H16(W[6]); + js.k3dst[0] = p.W3i ? p.W3i - (size_t)3 * 1536 * 768 : nullptr; + js.k3dst[1] = p.W3s ? p.W3s - (size_t)3 * 768 * 1536 : nullptr; + const int k3n = p.W3i ? 2304 : 0; + h->plans[f]=p;h->repacks[f]=js;h->repack_grids[f]=acc+k3n; + } + if(capacityF>=3){ + FastArgs a=h->args;a.codec.F=3;a.plan=h->plans[3];a.repack=h->repacks[3];a.repack_grid=h->repack_grids[3]; + a.status=(NativeStage1Status*)((char*)h->immutable+h->immutable_bytes); + int setup_occupancy=0; + if(!check(cudaFuncSetAttribute(stage1_prepare_immutable,cudaFuncAttributeMaxDynamicSharedMemorySize,TOTAL_SHARED),"setup shared-memory optin")|| + !check(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&setup_occupancy,stage1_prepare_immutable,512,TOTAL_SHARED),"setup occupancy")|| + setup_occupancy<1||h->blocks>setup_occupancy*prop.multiProcessorCount){ + if(error.empty())error="setup grid exceeds cooperative occupancy";release(h);return nullptr;} + void *params[]={&a}; + if(!check(cudaLaunchCooperativeKernel((void*)stage1_prepare_immutable,dim3(h->blocks),dim3(512),params,TOTAL_SHARED,nullptr),"immutable setup launch")|| + !check(cudaStreamSynchronize(nullptr),"immutable setup completion")){release(h);return nullptr;} + NativeStage1Status status{}; + if(!check(cudaMemcpy(&status,a.status,sizeof(status),cudaMemcpyDeviceToHost),"immutable setup status")){release(h);return nullptr;} + if(status.code){error="immutable setup device status "+std::to_string(status.code);release(h);return nullptr;} + h->setup_launches=1; + } + return h; +} +extern "C" int native_stage1_run(void *handle,const int64_t *codes,const int32_t *lengths,const int32_t *slots, + const void *const *states,void *const *next,float *pcm,int32_t *samples,NativeStage1Status *status,int f,void *stream){ + error.clear();Handle *h=(Handle*)handle; + if(!h||!codes||!lengths||!slots||!states||!next||!pcm||!samples||!status){error="null run argument";return 1;} + if(f<1||f>8||f>h->max_frames){error="true-state variant supports public F1..8";return 1;} + const int paddedF=f<=3?3:(f<=5?5:8); + int device=-1;if(!check(cudaGetDevice(&device),"device")||device!=h->device){error="device mismatch";return 1;} + FastArgs a=h->args;a.input_frames=f;a.plan=h->plans[paddedF];a.repack=h->repacks[paddedF];a.repack_grid=h->repack_grids[paddedF];a.codec.F=paddedF;a.codes=codes;a.lengths=lengths;a.slots=slots;a.pcm=pcm;a.samples=samples;a.status=status; + for(int p=0;p<26;++p){if(!states[p]||!next[p]||states[p]==next[p]){error="invalid/aliased state";return 1;}a.state[p]=(const __half*)states[p];a.next[p]=(__half*)next[p];} + void *params[]={&a};if(!check(cudaLaunchCooperativeKernel((void*)stage1_fast_persistent,dim3(h->blocks),dim3(512),params,TOTAL_SHARED,(cudaStream_t)stream),"fast Stage1 launch"))return 1; + return check(cudaPeekAtLastError(),"immediate launch status")?0:1; +} +extern "C" int native_stage1_destroy(void *p){error.clear();if(!p)return 0;release((Handle*)p);return check(cudaPeekAtLastError(),"destroy status")?0:1;} +extern "C" const char *native_stage1_last_error(){return error.c_str();} +extern "C" size_t native_stage1_scratch_bytes(void *p){return p?((Handle*)p)->bytes:0;} +extern "C" int native_stage1_resident_blocks(void *p){return p?((Handle*)p)->blocks:0;} +extern "C" int native_stage1_setup_kernel_launches(void *p){return p?((Handle*)p)->setup_launches:0;} +extern "C" size_t native_stage1_immutable_weight_bytes(void *p){return p?((Handle*)p)->immutable_bytes:0;} diff --git a/deploy/nemotron_native/kernels/codec/small/native_stage1.h b/deploy/nemotron_native/kernels/codec/small/native_stage1.h new file mode 100644 index 0000000..a95b5f0 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/small/native_stage1.h @@ -0,0 +1,26 @@ +#pragma once +#include +#include +struct NativeStage1Weights { const void *weight; const float *bias; const float *alpha; }; +struct NativeStage1Status { int32_t code, row, frame, detail; }; +#ifdef __cplusplus +extern "C" { +#endif +void *native_stage1_create(int batch, int max_frames, int state_slots, + int resident_blocks, const NativeStage1Weights *weights26); +int native_stage1_run(void *handle, const int64_t *codes, const int32_t *lengths, + const int32_t *slots, const void *const *states26, + void *const *next_states26, float *pcm, int32_t *samples, + NativeStage1Status *status, int frames, void *stream); +int native_stage1_run_profile(void *handle, const int64_t *codes, const int32_t *lengths, + const int32_t *slots, const void *const *states26, + void *const *next_states26, float *pcm, int32_t *samples, + NativeStage1Status *status, int frames, void *stream, + uint64_t *phase_timestamps30); +int native_stage1_destroy(void *handle); +const char *native_stage1_last_error(); +size_t native_stage1_scratch_bytes(void *handle); +int native_stage1_resident_blocks(void *handle); +#ifdef __cplusplus +} +#endif diff --git a/deploy/nemotron_native/kernels/codec/small/repack3.inc b/deploy/nemotron_native/kernels/codec/small/repack3.inc new file mode 100644 index 0000000..2f61c93 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/small/repack3.inc @@ -0,0 +1,66 @@ +__device__ __noinline__ void repack3_kernel( + const __half *__restrict__ win, const __half *__restrict__ wsk, + __half *__restrict__ din, __half *__restrict__ dsk) { + int4 (*sm)[49]=reinterpret_cast(native_arena()); + const int job = blockIdx.y; + const int Cout = job ? 768 : 1536; + const int Cin = job ? 1536 : 768; + const __half *src = job ? wsk : win; + __half *dst = job ? dsk : din; + const int GPB = Cin >> 3; + const int Kg = Cin * 3; + const int GT = GPB >> 4; + const int RT = Cout >> 5; + const int t = threadIdx.x; + for (int tile = blockIdx.x; tile < 6 * RT * GT; tile += gridDim.x) { + const int gblk = tile % GT; + const int rest = tile / GT; + const int corow = rest % RT; + const int b = rest / RT; + const int co0 = corow * 32; + const int g0 = gblk * 16; + const __half *sbase = src + (size_t)(b + 1) * Cout * Kg + + (size_t)co0 * Kg + g0 * 24; +#pragma unroll + for (int j = 0; j < 6; ++j) { + const int idx = t + j * 256; + const int r = idx / 48, c = idx - (idx / 48) * 48; + sm[r][c] = *reinterpret_cast(sbase + (size_t)r * Kg + + c * 8); + } + __syncthreads(); + __half *dbase = dst + (size_t)b * (3 * (size_t)Cout * Cin); +#pragma unroll + for (int p = 0; p < 2; ++p) { + const int u = t + p * 256; + const int col = (u & 7) + 8 * (u >> 7); + const int g = g0 + ((u >> 3) & 15); + const int co = co0 + col; + const int c = g >> 1, v = g & 1; + const int4 a = sm[col][((u >> 3) & 15) * 3]; + const int4 e = sm[col][((u >> 3) & 15) * 3 + 1]; + const int4 f = sm[col][((u >> 3) & 15) * 3 + 2]; + uint32_t h[12]; + h[0] = a.x; h[1] = a.y; h[2] = a.z; h[3] = a.w; + h[4] = e.x; h[5] = e.y; h[6] = e.z; h[7] = e.w; + h[8] = f.x; h[9] = f.y; h[10] = f.z; h[11] = f.w; + __half *d = dbase + (size_t)c * (48 * Cout) + + (size_t)(co >> 3) * 128 + v * 64 + (co & 7) * 8; +#pragma unroll + for (int tp = 0; tp < 3; ++tp) { + uint32_t o[4]; +#pragma unroll + for (int jj = 0; jj < 4; ++jj) { + const uint32_t g0w = h[3 * jj], g1w = h[3 * jj + 1], + g2w = h[3 * jj + 2]; + o[jj] = (tp == 0) ? __byte_perm(g0w, g1w, 0x7610) + : (tp == 1) ? __byte_perm(g0w, g2w, 0x5432) + : __byte_perm(g1w, g2w, 0x7610); + } + *reinterpret_cast(d + (size_t)tp * (16 * Cout)) = + make_int4(o[0], o[1], o[2], o[3]); + } + } + __syncthreads(); + } +} diff --git a/deploy/nemotron_native/kernels/codec/small/tail_kernel.inc b/deploy/nemotron_native/kernels/codec/small/tail_kernel.inc new file mode 100644 index 0000000..8794f42 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/small/tail_kernel.inc @@ -0,0 +1,492 @@ +#define TL_NT 512 +#define TL_MW 512 +#define TL_MLD TL_MW +#define TL_VW 504 +#define TL_UW 520 +#define TL_BN 496 +#define TL_VLD 520 +#define TL_LBO_U (TL_UW * 8) +#define TL_LBO_M (TL_MLD * 8) +#define TL_LBO_V (TL_VLD * 8) +#define TL_SBO_A 128 +#define TL_PU_E (4 * TL_UW * 8) +#define TL_PM_E (8 * TL_MLD * 8) +#define TL_PV_E (32 * TL_VLD) +#define TL_AI_E (14 * 1024) +#define TL_AS_E (32 * 512 + 512) +#define TL_SHARED ((2 * TL_PU_E + TL_PM_E + TL_PV_E + TL_AI_E + TL_AS_E) * 2) +__device__ __forceinline__ void tl_stmatrix4(void *dst, const uint32_t r[4], + bool trans) { + const uint32_t a = static_cast(__cvta_generic_to_shared(dst)); + if (trans) + asm volatile( + "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1,%2,%3,%4};\n" ::"r"( + a), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3])); + else + asm volatile( + "stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1,%2,%3,%4};\n" ::"r"( + a), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3])); +} +__device__ __forceinline__ void tl_ldmatrix4t(uint32_t r[4], const void *src) { + const uint32_t a = static_cast(__cvta_generic_to_shared(src)); + asm volatile( + "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(a)); +} +__device__ __forceinline__ void tl_stmatrix1t(void *dst, uint32_t r) { + const uint32_t a = static_cast(__cvta_generic_to_shared(dst)); + asm volatile( + "stmatrix.sync.aligned.m8n8.x1.trans.shared.b16 [%0], {%1};\n" :: + "r"(a), "r"(r)); +} +template +__device__ __forceinline__ void tl_skip_band( + const __half *PU, __half *PM, const __half *AS, + __half *PV, int c0, int warp, int lane, int er0, const float *sh_ask, + const float *sh_ari_v, float b0v, float b1v) { + constexpr int NB8 = BW / 8; + float acc[BW / 2]; +#pragma unroll + for (int i = 0; i < BW / 2; ++i) acc[i] = ((i & 2) == 0) ? b0v : b1v; + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int c = 0; c < 4; ++c) { + const uint64_t da0 = + wg_desc(AS + c * 8 * 512, 128u, TL_SBO_A * 2u); + const uint64_t db0 = + wg_desc(PM + (2 * c) * TL_LBO_M + (c0 + 2) * 8, + TL_LBO_M * 2u, 128u); +#pragma unroll + for (int p = 0; p < 4; ++p) { + const uint64_t da = da0 + 64 * (2 * p); + const uint64_t db = db0 + 2 * p; + if constexpr (BW == 120) { + WGMMA_M64N120K16_T(acc, da, db, 0, 0); + } else { + WGMMA_M64N128K16_T(acc, da, db, 0, 0); + } + } + } + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + const int m_addr = lane >> 3, row = lane & 7; + const int cg = (warp & 1) * 2; + if (er0 >= 32) { +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = __floats2half2_rn( + acc[jj * 4 + h * 2 + 0], acc[jj * 4 + h * 2 + 1]); + r[m] = *reinterpret_cast(&hv); + } else { + r[m] = 0; + } + } + const int jj = j + (m_addr >> 1), h = m_addr & 1; + tl_stmatrix4(PV + (cg + h) * TL_LBO_V + (c0 + jj * 8 + row) * 8, + r, true); + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + if (er0 < 32) { +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; + const int jja = j + (m_addr >> 1), ha = m_addr & 1; + tl_ldmatrix4t(r, PV + (cg + ha) * TL_LBO_V + + (c0 + jja * 8 + row + 1) * 8); +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = *reinterpret_cast<__half2 *>(&r[m]); + acc[jj * 4 + h * 2 + 0] += __half2float(hv.x); + acc[jj * 4 + h * 2 + 1] += __half2float(hv.y); + } + } + } +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; + const int jja = j + (m_addr >> 1), ha = m_addr & 1; + tl_ldmatrix4t(r, PU + (cg + ha) * TL_LBO_U + + (c0 + jja * 8 + row + 16) * 8); +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = *reinterpret_cast<__half2 *>(&r[m]); + acc[jj * 4 + h * 2 + 0] += __half2float(hv.x); + acc[jj * 4 + h * 2 + 1] += __half2float(hv.y); + } + } + } + } + if (er0 >= 32) return; + if (er0 < 16) { + const float a0 = sh_ask[er0], r0 = sh_ari_v[er0]; + const float a1 = sh_ask[er0 + 8], r1 = sh_ari_v[er0 + 8]; +#pragma unroll + for (int i = 0; i < BW / 2; ++i) { + const int h = (i & 2) >> 1; + const float v = acc[i]; + const float sn = __sinf((h ? a1 : a0) * v); + acc[i] = __fmaf_rn(sn * sn, h ? r1 : r0, v); + } + } else { +#pragma unroll + for (int i = 0; i < BW / 2; ++i) { + const float v = acc[i]; + acc[i] = v > 0.f ? v : 0.01f * v; + } + } +#pragma unroll + for (int j = 0; j < NB8; j += 2) { + uint32_t r[4]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + if (jj < NB8) { + const __half2 hv = __floats2half2_rn( + acc[jj * 4 + h * 2 + 0], acc[jj * 4 + h * 2 + 1]); + r[m] = *reinterpret_cast(&hv); + } else { + r[m] = 0; + } + } + const int jj = j + (m_addr >> 1), h = m_addr & 1; + tl_stmatrix4(PM + (cg + h) * TL_LBO_V + + (c0 + jj * 8 + row) * 8, + r, true); + } +} +__device__ __noinline__ void tail_fused_kernel( + const __half *__restrict__ U, const __half *__restrict__ win, + const __half *__restrict__ wsk, const __half *__restrict__ wpo, + const float *__restrict__ bin, const float *__restrict__ ain, + const float *__restrict__ bsk, const float *__restrict__ ask, + const float *__restrict__ bpo, const __half *__restrict__ st_m4, + const __half *__restrict__ st_v, __half *__restrict__ wave, + __half *__restrict__ ns_u, __half *__restrict__ ns_m4, + __half *__restrict__ ns_v, int B, int T4, int L4, int SC4, int N4, + int NTILE) { + char *tl_arena=reinterpret_cast(native_arena()); + __half *const PU = reinterpret_cast<__half *>(tl_arena); + __half *const PM = PU + 2 * TL_PU_E; + __half *const PV = PM + TL_PM_E; + __half *const AI = PV + TL_PV_E; + __half *const AS = AI + TL_AI_E; + __half *const PB = AS + 32 * 512; + float *sh_ain=(float*)native_aux(), *sh_ari=sh_ain+32, *sh_ask=sh_ari+32, *sh_ari_v=sh_ask+16; + const int tid = threadIdx.x; + const int wg = tid >> 7; + const int warp = (tid & 127) >> 5; + const int lane = tid & 31; + const int er0 = warp * 16 + (lane >> 2); + const int cb2 = (lane & 3) * 2; + for (int i = tid; i < 32; i += TL_NT) { + const float a = ain[i]; + sh_ain[i] = a; + sh_ari[i] = __frcp_rn(a + 1e-9f); + } + for (int i = tid; i < 16; i += TL_NT) { + const float a = ask[i]; + sh_ask[i] = a; + sh_ari_v[i] = __frcp_rn(a + 1e-9f); + } + for (int i = tid; i < TL_AI_E / 8; i += TL_NT) { + const int blk = i / 128, off = (i - blk * 128) * 8; + const int c = blk / 7, k = blk - c * 7; + *reinterpret_cast(AI + blk * 1024 + off) = + *reinterpret_cast(win + (size_t)c * 14336 + + (size_t)k * 2048 + off); + } + for (int i = tid; i < 32 * 64; i += TL_NT) { + const int blk = i / 64, off = (i - blk * 64) * 8; + const int c = blk >> 3, k = blk & 7; + *reinterpret_cast(AS + blk * 512 + off) = + (k < 7) ? *reinterpret_cast( + wsk + (size_t)c * 14336 + (size_t)k * 2048 + off) + : make_int4(0, 0, 0, 0); + } + for (int i = tid; i < 4 * (TL_VLD - TL_VW) * 8; i += TL_NT) { + const int g = i / ((TL_VLD - TL_VW) * 8); + PV[g * TL_LBO_V + TL_VW * 8 + (i - g * ((TL_VLD - TL_VW) * 8))] = + __float2half(0.f); + } + for (int i = tid; i < 512; i += TL_NT) PB[i] = __float2half(0.f); + for (int i = tid; i < 2 * 16 * 8; i += TL_NT) { + const int c = i / (16 * 8); + const int r = i - c * 16 * 8; + const int ki = r / 8, k = r & 7; + PB[c * 128 + (ki >> 3) * 64 + (ki & 7) * 8 + k] = + (k < 7) ? wpo[(c * 16 + ki) * 7 + k] : __float2half(0.f); + } + const float bias_i0 = bin[er0], bias_i1 = bin[er0 + 8]; + const float bias_s0 = (er0 < 32) ? bsk[er0] : 0.f; + const float bias_s1 = (er0 + 8 < 32) ? bsk[er0 + 8] : 0.f; + const float bpo0 = bpo[0]; + constexpr int UCOL = TL_UW / 32; + constexpr int UUNIT = 4 * UCOL; + constexpr int UWARP = TL_NT / 32; + constexpr int UITER = (UUNIT + UWARP - 1) / UWARP; + const int uwarp = tid >> 5; + const int uq = lane & 3; + int ugc[UITER], uoc[UITER]; +#pragma unroll + for (int t = 0; t < UITER; ++t) { + const int bi = uwarp + t * UWARP; + const int gc = bi / UCOL; + ugc[t] = gc; + uoc[t] = bi - gc * UCOL; + } + const __half *const urow = + U + (size_t)(lane >> 2) * SC4 + 8 * uq; + __syncthreads(); + PDL_ENTER(); +#define TL_STAGE_U(PUD,UB) { _Pragma("unroll") for (int t2 = 0; t2 < UITER; t2 += 2) { uint32_t u[2][4]; _Pragma("unroll") for (int d = 0; d < 2; ++d) { const int t = t2 + d; const int n = (UB) + uoc[t] * 32; const int gn = n + 8 * uq; const bool ok = gn >= 0 && gn < N4; int4 v = make_int4(0, 0, 0, 0); if (ok) v = *reinterpret_cast( urow + (size_t)ugc[t] * (8 * (size_t)SC4) + n); u[d][0] = v.x; u[d][1] = v.y; u[d][2] = v.z; u[d][3] = v.w; } _Pragma("unroll") for (int d = 0; d < 2; ++d) { const int t = t2 + d; uint32_t *q = u[d]; uint32_t x, y; x = (uq & 1) ? q[0] : q[1]; y = (uq & 1) ? q[2] : q[3]; x = __shfl_xor_sync(0xffffffffu, x, 1); y = __shfl_xor_sync(0xffffffffu, y, 1); if (uq & 1) { q[0] = x; q[2] = y; } else { q[1] = x; q[3] = y; } x = (uq & 2) ? q[0] : q[2]; y = (uq & 2) ? q[1] : q[3]; x = __shfl_xor_sync(0xffffffffu, x, 2); y = __shfl_xor_sync(0xffffffffu, y, 2); if (uq & 2) { q[0] = x; q[1] = y; } else { q[2] = x; q[3] = y; } tl_stmatrix4((PUD) + ugc[t] * TL_LBO_U + (uoc[t] * 32 + lane) * 8, q, true); } } if (uwarp < 4) { const int n = (UB) + 512; const int gn = n + 2 * (lane & 3); uint32_t q = 0; if (gn >= 0 && gn + 1 < N4) q = *reinterpret_cast( U + (size_t)(uwarp * 8 + (lane >> 2)) * SC4 + gn); tl_stmatrix1t((PUD) + uwarp * TL_LBO_U + (512 + (lane & 7)) * 8, q); } } + TL_STAGE_U(PU, (int)blockIdx.x * TL_BN - 16) + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + int ubuf = 0; + for (int tile = blockIdx.x; tile < NTILE; tile += gridDim.x, ubuf ^= 1) { + const int g0 = tile * TL_BN; + __half *const PUc = PU + ubuf * TL_PU_E; + __half *const PUn = PU + (ubuf ^ 1) * TL_PU_E; + { + const int b0 = g0 / L4; + for (int s = 0; s < 2; ++s) { + const int sb = b0 + s; + if (sb >= B) break; + const int cv = (int)((long long)sb * L4 + native_true_tokens(sb,T4) - g0); + if (cv < 0 || cv >= TL_BN) continue; + for (int i = tid; i < 32 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + ns_u[((size_t)sb * 32 + co) * 6 + h] = + PUc[(co >> 3) * TL_LBO_U + + (cv + h + 16) * 8 + (co & 7)]; + } + } + } + { + float acc[64]; +#pragma unroll + for (int i = 0; i < 64; ++i) + acc[i] = ((i & 2) == 0) ? bias_i0 : bias_i1; + const int c0 = wg * 128; + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int c = 0; c < 2; ++c) { + const uint64_t da0 = + wg_desc(AI + c * 7 * 1024, 128u, TL_SBO_A * 2u); + const uint64_t db0 = + wg_desc(PUc + (2 * c) * TL_LBO_U + (c0 + 2) * 8, + TL_LBO_U * 2u, 128u); +#pragma unroll + for (int k = 0; k < 7; ++k) { + const uint64_t da = da0 + 128 * k; + const uint64_t db = db0 + k; + WGMMA_M64N128K16_T(acc, da, db, 0, 0); + } + } + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); + if (tile + (int)gridDim.x < NTILE) + TL_STAGE_U(PUn, (tile + (int)gridDim.x) * TL_BN - 16) + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + if (er0 < 32) { + const float a0 = sh_ain[er0], r0 = sh_ari[er0]; + const float a1 = sh_ain[er0 + 8], r1 = sh_ari[er0 + 8]; +#pragma unroll + for (int i = 0; i < 64; ++i) { + const int h = (i & 2) >> 1; + const float v = acc[i]; + const float sn = __sinf((h ? a1 : a0) * v); + acc[i] = __fmaf_rn(sn * sn, h ? r1 : r0, v); + } + } else { +#pragma unroll + for (int i = 0; i < 64; ++i) { + const float v = acc[i]; + acc[i] = v > 0.f ? v : 0.01f * v; + } + } +#pragma unroll + for (int j = 0; j < 16; j += 2) { + uint32_t r[4]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int jj = j + (m >> 1), h = m & 1; + const __half2 hv = __floats2half2_rn( + acc[jj * 4 + h * 2 + 0], acc[jj * 4 + h * 2 + 1]); + r[m] = *reinterpret_cast(&hv); + } + const int m_addr = lane >> 3, row = lane & 7; + const int jj = j + (m_addr >> 1), h = m_addr & 1; + tl_stmatrix4(PM + (warp * 2 + h) * TL_LBO_M + + (c0 + jj * 8 + row) * 8, + r, true); + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + { + const int base = g0 - 8; + const int b0 = (base >= 0) ? base / L4 : 0; + const int cm0 = (int)((long long)b0 * L4 - base); + const int cm1 = cm0 + L4; + const bool hit0 = (b0 < B) && (cm0 + 6 > 0) && (cm0 < TL_MW); + const bool hit1 = (b0 + 1 < B) && (cm1 + 6 > 0) && (cm1 < TL_MW); + if (hit0 || hit1) { + for (int s = 0; s < 2; ++s) { + if (!(s ? hit1 : hit0)) continue; + const int sb = b0 + s, cm = s ? cm1 : cm0; + for (int i = tid; i < 64 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + const int c = cm + h; + if (c >= 0 && c < TL_MW) + PM[(co >> 3) * TL_LBO_M + c * 8 + (co & 7)] = + st_m4[((size_t)sb * 64 + co) * 6 + h]; + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + } + } + { + const int b0 = g0 / L4; + for (int s = 0; s < 2; ++s) { + const int sb = b0 + s; + if (sb >= B) break; + const int cv = (int)((long long)sb * L4 + native_true_tokens(sb,T4) - g0); + if (cv < 0 || cv >= TL_BN) continue; + for (int i = tid; i < 64 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + ns_m4[((size_t)sb * 64 + co) * 6 + h] = + PM[(co >> 3) * TL_LBO_M + + (cv + h + 8) * 8 + (co & 7)]; + } + } + } + { + const int c0 = wg * 128; + if (wg == 3) + tl_skip_band<120>(PUc, PM, AS, PV, c0, warp, lane, er0, sh_ask, + sh_ari_v, bias_s0, bias_s1); + else + tl_skip_band<128>(PUc, PM, AS, PV, c0, warp, lane, er0, sh_ask, + sh_ari_v, bias_s0, bias_s1); + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + { + const int b0 = g0 / L4; + const int cv0 = (int)((long long)b0 * L4 - g0), cv1 = cv0 + L4; + const bool hv0 = (b0 < B) && (cv0 + 6 > 0) && (cv0 < TL_VW); + const bool hv1 = (b0 + 1 < B) && (cv1 + 6 > 0) && (cv1 < TL_VW); + if (hv0 || hv1) { + for (int s = 0; s < 2; ++s) { + if (!(s ? hv1 : hv0)) continue; + const int sb = b0 + s, cv = s ? cv1 : cv0; + for (int i = tid; i < 32 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + const int c = cv + h; + if (c >= 0 && c < TL_VW) + PM[(co >> 3) * TL_LBO_V + c * 8 + (co & 7)] = + st_v[((size_t)sb * 32 + co) * 6 + h]; + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + } + } + { + const int b0 = g0 / L4; + for (int s = 0; s < 2; ++s) { + const int sb = b0 + s; + if (sb >= B) break; + const int cv = (int)((long long)sb * L4 + native_true_tokens(sb,T4) - g0); + if (cv < 0 || cv >= TL_BN) continue; + for (int i = tid; i < 32 * 6; i += TL_NT) { + const int co = i / 6, h = i - co * 6; + ns_v[((size_t)sb * 32 + co) * 6 + h] = + PM[(co >> 3) * TL_LBO_V + (cv + h) * 8 + (co & 7)]; + } + } + } + __syncthreads(); + constexpr int PC_LD = 10; + { + float *const PC = reinterpret_cast(PV); +#pragma unroll + for (int round = 0; round < 2; ++round) { + float acc[4]; + const int p0 = round * 256 + wg * 64; + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); + const uint64_t da0 = + wg_desc(PM + p0 * 8, TL_LBO_V * 2u, 128u); + const uint64_t db0 = wg_desc(PB, 128u, 128u); + const uint64_t da1 = + wg_desc(PM + 2 * TL_LBO_V + p0 * 8, + TL_LBO_V * 2u, 128u); + const uint64_t db1 = wg_desc(PB + 128, 128u, 128u); + WGMMA_M64N8K16_ZERO_T(acc, da0, db0, 0, 1); + WGMMA_M64N8K16_T(acc, da1, db1, 0, 1); + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + const int row0 = p0 + warp * 16 + (lane >> 2); + const int col = (lane & 3) * 2; + *reinterpret_cast(PC + row0 * PC_LD + col) = + make_float2(acc[0], acc[1]); + *reinterpret_cast(PC + (row0 + 8) * PC_LD + col) = + make_float2(acc[2], acc[3]); + } + } + __syncthreads(); + { + const float *const PC = reinterpret_cast(PV); + const int p0 = tid; + const bool live = p0 < TL_BN; + const unsigned live_mask = __ballot_sync(0xffffffffu, live); + if (live) { + const float s0 = __fadd_rn(PC[(p0 + 0) * PC_LD + 0], + PC[(p0 + 1) * PC_LD + 1]); + const float s1 = __fadd_rn(PC[(p0 + 2) * PC_LD + 2], + PC[(p0 + 3) * PC_LD + 3]); + const float s2 = __fadd_rn(PC[(p0 + 4) * PC_LD + 4], + PC[(p0 + 5) * PC_LD + 5]); + const float s3 = PC[(p0 + 6) * PC_LD + 6]; + float y0 = __fadd_rn(__fadd_rn(s0, s1), + __fadd_rn(s2, s3)); + y0 = fminf(fmaxf(y0 + bpo0, -1.f), 1.f); + const float y1 = __shfl_down_sync(live_mask, y0, 1); + if ((p0 & 1) == 0) { + const int b0 = g0 / L4; + const int rr = g0 - b0 * L4; + int t0 = rr + p0, bo0 = b0; + int t1 = t0 + 1, bo1 = bo0; + if (t0 >= L4) { t0 -= L4; ++bo0; } + if (t1 >= L4) { t1 -= L4; ++bo1; } + if (bo0 == bo1 && bo0 < B && t1 < T4) { + *reinterpret_cast<__half2 *>( + wave + (size_t)bo0 * T4 + t0) = + __floats2half2_rn(y0, y1); + } else { + if (bo0 < B && t0 < T4) + wave[(size_t)bo0 * T4 + t0] = __float2half(y0); + if (bo1 < B && t1 < T4) + wave[(size_t)bo1 * T4 + t1] = __float2half(y1); + } + } + } + } + } +#undef TL_STAGE_U +} diff --git a/deploy/nemotron_native/kernels/codec/small/wg2_kernel.inc b/deploy/nemotron_native/kernels/codec/small/wg2_kernel.inc new file mode 100644 index 0000000..4ffa7ad --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/small/wg2_kernel.inc @@ -0,0 +1,280 @@ +__device__ __forceinline__ void cpa16(void *dsm, const void *src, bool pred) { + unsigned a = static_cast(__cvta_generic_to_shared(dsm)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(a), + "l"(src), "r"(pred ? 16 : 0) + : "memory"); +} +template +__device__ __forceinline__ void cpa16o(unsigned a, const void *src, bool pred) { + asm volatile("cp.async.cg.shared.global [%0+%3], [%1], 16, %2;\n" ::"r"(a), + "l"(src), "r"(pred ? 16 : 0), "n"(OFF) + : "memory"); +} +template +__device__ __noinline__ +void wg2_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int BK = CIB * KT; + constexpr int WG_BM = 64 * MSPLIT; + constexpr int WG_BNT = WG_BN * NSPLIT; + constexpr int NTS = 128 * MSPLIT * NSPLIT; + constexpr int WG_NT = NTS * KSPLIT; + constexpr int NCRAW = WG_BNT / 4; + constexpr int EPN = + (WG_NT % NCRAW == 0) ? WG_BNT : (((WG_BNT + 127) / 128) * 128); + constexpr int WG_CLD = EPN + 8; + constexpr bool NB8_P2 = ((WG_BN / 8) & (WG_BN / 8 - 1)) == 0; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int ST_E = SA * A_E + SB * NSPLIT * B_E; + constexpr int SBO_AE = KG * WG_SLOT; + constexpr int LBO_BE = (WG_BN / 8) * WG_SLOT; + constexpr int NB8 = WG_BN / 8; + constexpr bool ALDG = (WG_BN <= 64); + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *smem = reinterpret_cast<__half *>(arena); + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int kgrp = tid / NTS; + const int stid = tid - kgrp * NTS; + const int wgl = stid >> 7; + const int mg = wgl % MSPLIT; + const int ngid = wgl / MSPLIT; + const int warp = (stid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * WG_BM; + const int nb = blockIdx.y * WG_BNT; + const int n0 = nb + ngid * WG_BN; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + const int NCH = (Kg + BK - 1) / BK; + if (tid < WG_BM) { + const int co = m0 + tid; + if (do_snake && co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + for (int i = tid; i < EPN; i += WG_NT) { + const int n = nb + i; + colok[i] = (i < WG_BNT && n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mg * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (kgrp == 0 && eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (kgrp == 0 && eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + constexpr int KGX = ALDG ? (KG <= 8 ? 8 : (KG <= 16 ? 16 : 32)) : KG; + constexpr int ATOT = WG_BM * KGX; + constexpr int AREG = (ATOT + NTS - 1) / NTS; + constexpr int BTOT = NSPLIT * CIB * NB8; + constexpr int BREG = (BTOT + NTS - 1) / NTS; + const __half *aptr[AREG]; + int aoff[AREG]; + int akmax[AREG]; +#pragma unroll + for (int r = 0; r < AREG; ++r) { + const int idx = stid + r * NTS; + const int m = idx / KGX, v = idx - (idx / KGX) * KGX; + const int co = m0 + m; + const bool live = (ATOT % NTS == 0 || idx < ATOT) && (co < Cout) && + v < KG; + aoff[r] = (m >> 3) * SBO_AE + (v < KG ? v : 0) * WG_SLOT + (m & 7) * 8; + akmax[r] = live ? (Kg - v * 8) : 0; + aptr[r] = w + (size_t)(live ? co : 0) * Kg + v * 8; + } + const __half *bptr[BREG]; + int bcimax[BREG], bnc[BREG], bcil[BREG], bgrp[BREG]; +#pragma unroll + for (int r = 0; r < BREG; ++r) { + const int idx = stid + r * NTS; + const int gg = idx / (CIB * NB8); + const int rr = idx - gg * (CIB * NB8); + const int cil = rr / NB8, nc = rr - (rr / NB8) * NB8; + bcil[r] = cil; + bnc[r] = nc; + bgrp[r] = gg; + const bool live = (BTOT % NTS == 0 || idx < BTOT); + bcimax[r] = live ? (Cin - cil) : 0; + bptr[r] = in + (size_t)cil * SCin + nb + gg * WG_BN + nc * 8; + } + int4 br0[BREG], br1[BREG]; + int4 ar[ALDG ? AREG : 1]; + uint32_t edge[BREG]; + __half *const kaB = smem + kgrp * ST_E; + __half *const kbB = kaB + SA * A_E; + const unsigned barid = 1u + kgrp + (unsigned)(KSPLIT * native_subid()); +#define WG2_BAR() asm volatile("bar.sync %0, %1;\n" ::"r"(barid), "r"(NTS)) +#define WG2_LOADB(c) { const int _ci0 = (c) * CIB; _Pragma("unroll") for (int r = 0; r < BREG; ++r) { br0[r] = make_int4(0, 0, 0, 0); if constexpr (KT == 3) edge[r] = 0; else br1[r] = make_int4(0, 0, 0, 0); if (_ci0 < bcimax[r]) { const int4 *src = reinterpret_cast( bptr[r] + (size_t)_ci0 * SCin); br0[r] = src[0]; if constexpr (KT == 3) { if (!NB8_P2 || bnc[r] == NB8 - 1) edge[r] = src[1].x; } else { br1[r] = src[1]; } } } } +#define WG2_LOADA(c) if constexpr (ALDG) { const int _kg0 = (c) * BK; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { ar[r] = make_int4(0, 0, 0, 0); if (_kg0 < akmax[r]) ar[r] = *reinterpret_cast(aptr[r] + _kg0); } } +#define WG2_PUTA(c,AD) { _Pragma("unroll") for (int r = 0; r < AREG; ++r) { if constexpr (ALDG) { if (akmax[r] > 0) *reinterpret_cast((AD) + aoff[r]) = ar[r]; } else { const int _kg0 = (c) * BK; if (akmax[r] > 0) cpa16((AD) + aoff[r], aptr[r] + _kg0, _kg0 < akmax[r]); } } if constexpr (!ALDG) asm volatile("cp.async.commit_group;\n" ::: "memory"); } +#define WG2_STSB(BD) { _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int cil = bcil[r], nc = bnc[r]; unsigned int u[8]; u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; if constexpr (KT == 3) { if constexpr (NB8_P2) { u[4] = __shfl_down_sync(0xffffffffu, u[0], 1, NB8); if (nc == NB8 - 1) u[4] = edge[r]; } else { u[4] = edge[r]; } u[5] = u[6] = u[7] = 0; } else { u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; } if (BTOT % NTS == 0 || (stid + r * NTS) < BTOT) { _Pragma("unroll") for (int k = 0; k < KT; ++k) { const int kh = k >> 1; int4 o; if (k & 1) { o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); } else { o.x = u[kh + 0]; o.y = u[kh + 1]; o.z = u[kh + 2]; o.w = u[kh + 3]; } const int row = cil * KT + k; *reinterpret_cast((BD) + bgrp[r] * B_E + nc * WG_SLOT + (row >> 3) * LBO_BE + (row & 7) * 8) = o; } } } } +#define WG2_MMA(AD,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t da0 = wg_desc((AD) + mg * 8 * SBO_AE, WG_SLOT * 2, SBO_AE * 2); const uint64_t db0 = wg_desc((BD) + ngid * B_E, LBO_BE * 2, WG_SLOT * 2); _Pragma("unroll") for (int kk = 0; kk < BK; kk += 16) { const uint64_t da = da0 + (kk >> 3) * (WG_SLOT / 8); const uint64_t db = db0 + (kk >> 3) * (LBO_BE / 8); if constexpr (WG_BN == 64) { WGMMA_M64N64K16(acc, da, db); } else if constexpr (WG_BN == 256) { WGMMA_M64N256K16(acc, da, db); } else if constexpr (WG_BN == 104) { WGMMA_M64N104K16(acc, da, db); } else if constexpr (WG_BN == 112) { WGMMA_M64N112K16(acc, da, db); } else if constexpr (WG_BN == 120) { WGMMA_M64N120K16(acc, da, db); } else if constexpr (WG_BN == 224) { WGMMA_M64N224K16(acc, da, db); } else { static_assert(WG_BN == 128, "no wgmma macro for WG_BN"); WGMMA_M64N128K16(acc, da, db); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + constexpr int SMAX = SA > SB ? SA : SB; + WG2_LOADA(kgrp) + WG2_LOADB(kgrp) +#pragma unroll 1 + for (int s = 0; s < SMAX; ++s) { + const int c = s * KSPLIT + kgrp; + if (s < SA) { + if (c < NCH) WG2_PUTA(c, kaB + s * A_E) + if (c + KSPLIT < NCH) WG2_LOADA(c + KSPLIT) + } + if (s < SB) { + if (c < NCH) WG2_STSB(kbB + s * NSPLIT * B_E) + if (c + KSPLIT < NCH) WG2_LOADB(c + KSPLIT) + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + const int NIT = (NCH - kgrp + KSPLIT - 1) / KSPLIT; + for (int i = 0; i < NIT; ++i) { + if constexpr (!ALDG) + asm volatile("cp.async.wait_group %0;\n" ::"n"(SA - 2) : "memory"); + WG2_BAR(); + WG2_MMA(kaB + (i % SA) * A_E, kbB + (i % SB) * NSPLIT * B_E) + asm volatile("wgmma.wait_group.sync.aligned 1;\n" ::: "memory"); + if constexpr (NTS > 128) WG2_BAR(); + if (i >= 1) { + const int ca = (i + SA - 1) * KSPLIT + kgrp; + if (ca < NCH) WG2_PUTA(ca, kaB + ((i - 1) % SA) * A_E) + if (ALDG && ca + KSPLIT < NCH) WG2_LOADA(ca + KSPLIT) + const int cb = (i + SB - 1) * KSPLIT + kgrp; + if (cb < NCH) { + WG2_STSB(kbB + ((i - 1) % SB) * NSPLIT * B_E) + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if (cb + KSPLIT < NCH) WG2_LOADB(cb + KSPLIT) + } else if constexpr (ALDG) { + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + } + } + } + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); +#undef WG2_LOADB +#undef WG2_LOADA +#undef WG2_PUTA +#undef WG2_STSB +#undef WG2_MMA +#undef WG2_BAR + __syncthreads(); + float *Ck = Cs + kgrp * WG_BM * WG_CLD; + const int cb = (lane & 3) * 2; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Ck + (er0 + h * 8) * WG_CLD + + ngid * WG_BN + j * 8 + cb) = v; + } + } + __syncthreads(); + constexpr int NC = EPN / 4; + constexpr int MSTEP = WG_NT / NC; + constexpr int MITER = WG_BM / MSTEP; + constexpr int UNR = (MITER % 4 == 0) ? 4 : ((MITER % 2 == 0) ? 2 : 1); + static_assert(MITER % UNR == 0, "row batch must divide the row count"); + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + const int n = nb + lc; +#pragma unroll 1 + for (int t0 = 0; t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + v[u] = *reinterpret_cast(Cs + m * WG_CLD + lc); + if (resid != nullptr) { + const int co = m0 + m; + const int cc = co < Cout ? co : 0; + const __half2 *rp = reinterpret_cast( + resid + (size_t)cc * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + const int co = m0 + m; + if (co >= Cout) continue; + float4 vv = v[u]; +#pragma unroll + for (int g = 1; g < KSPLIT; ++g) { + float4 p = *reinterpret_cast( + Cs + g * WG_BM * WG_CLD + m * WG_CLD + lc); + vv.x += p.x; vv.y += p.y; vv.z += p.z; vv.w += p.w; + } + if (resid != nullptr) { + vv.x += __half2float(r0[u].x); vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (do_snake) { + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + } + __half *op = out + (size_t)co * SCout + n + H; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B * H; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int total = Cout * B * pad_len; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} diff --git a/deploy/nemotron_native/kernels/codec/small/wg3_kernel.inc b/deploy/nemotron_native/kernels/codec/small/wg3_kernel.inc new file mode 100644 index 0000000..fa55b32 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/small/wg3_kernel.inc @@ -0,0 +1,177 @@ +template +__device__ __noinline__ void wg3_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int WG_BNT = WG_BN * NSG; + constexpr int WG_CLD = WG_BNT + 8; + constexpr int WG_NT = 128 * NWG * NSG; + constexpr int WG_BM = 64 * NWG; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int SBO_AE = KG * WG_SLOT; + constexpr int LBO_BE = (WG_BN / 8) * WG_SLOT; + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + constexpr int ST_E = A_E + NSG * B_E; + __half *As0 = reinterpret_cast<__half *>(arena); + __half *Bs0 = As0 + A_E; + __half *As1 = (STAGES == 2) ? As0 + ST_E : As0; + __half *Bs1 = (STAGES == 2) ? Bs0 + ST_E : Bs0; + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int wgid = tid >> 7; + const int mgid = wgid % NWG, ngid = wgid / NWG; + const int warp = (tid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * WG_BM; + constexpr int WGROWS = 8; + const int nb = blockIdx.y * WG_BNT; + const int n0 = nb + ngid * WG_BN; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + if (tid < WG_BM) { + const int co = m0 + tid; + if (do_snake && co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + if (tid < WG_BNT) { + const int n = nb + tid; + colok[tid] = (n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mgid * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + constexpr int ATOT = WG_BM * KG; + constexpr int AREG = (ATOT + WG_NT - 1) / WG_NT; + constexpr int CIB = BK / KT; + constexpr int BTOT = NSG * CIB * (WG_BN / 8); + constexpr int BREG = (BTOT + WG_NT - 1) / WG_NT; + int4 ar[AREG]; + int4 br0[BREG], br1[BREG]; + uint32_t edge[BREG]; +#define WG_LOADG(kg0) { const int _ci0 = (kg0) / KT; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; ar[r] = make_int4(0, 0, 0, 0); if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; const int co = m0 + m, kg = (kg0) + v * 8; if (co < Cout && kg < Kg) ar[r] = *reinterpret_cast( w + (size_t)co * Kg + kg); } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; br0[r] = make_int4(0, 0, 0, 0); if constexpr (KT == 3) edge[r] = 0; else br1[r] = make_int4(0, 0, 0, 0); if (BTOT % WG_NT == 0 || idx < BTOT) { const int _g = idx / (CIB * (WG_BN / 8)); const int _r = idx - _g * (CIB * (WG_BN / 8)); const int ci = _ci0 + _r / (WG_BN / 8); const int nc = _r % (WG_BN / 8); if (ci < Cin) { const int4 *src = reinterpret_cast( in + (size_t)ci * SCin + nb + _g * WG_BN + nc * 8); br0[r] = src[0]; if constexpr (KT == 3) { if (nc == WG_BN / 8 - 1) edge[r] = src[1].x; } else { br1[r] = src[1]; } } } } } + WG_LOADG(0) +#define WG_STS(AD,BD) { _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; *reinterpret_cast( (AD) + (m >> 3) * SBO_AE + v * WG_SLOT + (m & 7) * 8) = ar[r]; } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; if (BTOT % WG_NT == 0 || idx < BTOT) { const int _g = idx / (CIB * (WG_BN / 8)); const int _r = idx - _g * (CIB * (WG_BN / 8)); const int cil = _r / (WG_BN / 8), nc = _r % (WG_BN / 8); unsigned int u[8]; u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; if constexpr (KT == 3) { u[4] = __shfl_down_sync(0xffffffffu, u[0], 1, WG_BN / 8); if (nc == WG_BN / 8 - 1) u[4] = edge[r]; u[5] = u[6] = u[7] = 0; } else { u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; } _Pragma("unroll") for (int k = 0; k < KT; ++k) { const int kh = k >> 1; int4 o; if (k & 1) { o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); } else { o.x = u[kh + 0]; o.y = u[kh + 1]; o.z = u[kh + 2]; o.w = u[kh + 3]; } const int row = cil * KT + k; *reinterpret_cast((BD) + _g * B_E + nc * WG_SLOT + (row >> 3) * LBO_BE + (row & 7) * 8) = o; } } } if constexpr (WG_BN == 64) { asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); } } +#define WG_MMA(AD,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t da0 = wg_desc((AD) + mgid * WGROWS * SBO_AE, WG_SLOT * 2, SBO_AE * 2); const uint64_t db0 = wg_desc((BD) + ngid * B_E, LBO_BE * 2, WG_SLOT * 2); _Pragma("unroll") for (int kk = 0; kk < BK; kk += 16) { const uint64_t da = da0 + (kk >> 3) * (WG_SLOT / 8); const uint64_t db = db0 + (kk >> 3) * (LBO_BE / 8); if constexpr (WG_BN == 64) { WGMMA_M64N64K16(acc, da, db); } else { WGMMA_M64N128K16(acc, da, db); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + WG_STS(As0, Bs0) + if (BK < Kg) WG_LOADG(BK) + __syncthreads(); + int cur = 0; + for (int kg0 = 0; kg0 < Kg; kg0 += BK) { + if (cur == 0) { WG_MMA(As0, Bs0) } else { WG_MMA(As1, Bs1) } + if (kg0 + BK < Kg) { + if (STAGES == 2) + asm volatile("wgmma.wait_group.sync.aligned 1;\n" ::: "memory"); + else + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + __syncthreads(); + if (cur == 0) { WG_STS(As1, Bs1) } else { WG_STS(As0, Bs0) } + if (kg0 + 2 * BK < Kg) WG_LOADG(kg0 + 2 * BK) + __syncthreads(); + } else { + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + } + if (STAGES == 2) cur ^= 1; + } +#undef WG_LOADG +#undef WG_STS +#undef WG_MMA + __syncthreads(); + const int r0 = er0; + const int cb = (lane & 3) * 2; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Cs + (r0 + h * 8) * WG_CLD + + ngid * WG_BN + j * 8 + cb) = v; + } + } + __syncthreads(); + constexpr int NC = WG_BNT / 4; + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + for (int m = lr; m < WG_BM; m += WG_NT / NC) { + const int co = m0 + m; + if (co >= Cout) continue; + const int n = nb + lc; + float4 v = *reinterpret_cast(Cs + m * WG_CLD + lc); + __half *op = out + (size_t)co * SCout + n + H; + if (resid != nullptr) { + const __half2 *rp = reinterpret_cast( + resid + (size_t)co * SCout + n + H); + __half2 r0h = rp[0], r1h = rp[1]; + v.x += __half2float(r0h.x); v.y += __half2float(r0h.y); + v.z += __half2float(r1h.x); v.w += __half2float(r1h.y); + } + if (do_snake) { + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + } + if (ok4) { + __half2 o0 = __floats2half2_rn(v.x, v.y); + __half2 o1 = __floats2half2_rn(v.z, v.w); + reinterpret_cast<__half2 *>(op)[0] = o0; + reinterpret_cast<__half2 *>(op)[1] = o1; + } else { + const float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B * H; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int total = Cout * B * pad_len; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} diff --git a/deploy/nemotron_native/kernels/codec/small/wg4_step_widths.inc b/deploy/nemotron_native/kernels/codec/small/wg4_step_widths.inc new file mode 100644 index 0000000..9fe1878 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/small/wg4_step_widths.inc @@ -0,0 +1,459 @@ +template +__device__ __noinline__ void wg4_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int CoutP, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + if constexpr (PDL && !LATE_PDL) PDL_ENTER(); + constexpr int NT = 128 * MSPLIT * NSPLIT; + constexpr int BM = 64 * MSPLIT; + constexpr int BNT = WG_BN * NSPLIT; + constexpr int NCRAW = BNT / 4; + constexpr int EPN = (NT % NCRAW == 0) ? BNT : (((BNT + 127) / 128) * 128); + constexpr int WG_CLD = EPN + 8; + constexpr int CG = CIB / 8; + constexpr int NBR = BNT + KT - 1; + constexpr int NBRP = BC8 ? ((NBR + 7) & ~7) : ((NBR + 31) & ~31); + constexpr int UCOL = NBRP / 32; + constexpr int NUNIT = CG * UCOL; + constexpr int NWARP = NT / 32; + constexpr int BITER = (NUNIT + NWARP - 1) / NWARP; + constexpr int B_E = CIB * NBRP; + constexpr int ST_E = B_E; + constexpr int LBO_B = NBRP * 8; + constexpr int ABUF = ABUF_; + static_assert(!OC8 || PREACT, + "K=7 C8 output is only used for residual-free pair heads"); + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *smem = reinterpret_cast<__half *>(arena); + __half *Hs = smem; + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int wgid = tid >> 7; + const int mg = wgid % MSPLIT, ng = wgid / MSPLIT; + const int warp = (tid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * BM; + const int nb = blockIdx.y * BNT; + const int H = KT - 1; + const int half = Cout >> 1; + const int NCH = (Cin + CIB - 1) / CIB; + for (int i = tid; i < BM; i += NT) { + const int co = m0 + i; + if (co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[i] = a; + row_ri[i] = __frcp_rn(a + 1e-9f); + } + } + for (int i = tid; i < EPN; i += NT) { + const int n = nb + i; + colok[i] = (i < BNT && n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mg * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + const int mtile = blockIdx.x * MSPLIT + mg; + const __half *const apt = w + ((size_t)mtile * NCH * KT) * 1024 + + (size_t)(tid & 127) * 8; + uint32_t areg[ABUF][KT][4]; + const int wi = tid >> 5; + int bdst[BITER], boff[BITER]; + bool blive[BITER]; + const __half *const bbase = in + (size_t)(lane >> 2) * SCin + nb + + 8 * (lane & 3); +#pragma unroll + for (int t = 0; t < BITER; ++t) { + const int u = wi + t * NWARP; + const int gc = (u < NUNIT) ? u / UCOL : 0; + const int cu = (u < NUNIT) ? u - gc * UCOL : 0; + blive[t] = u < NUNIT; + bdst[t] = gc * LBO_B + (cu * 32 + lane) * 8; + boff[t] = gc * 8 * SCin + cu * 32; + } + const size_t SC8i = (size_t)SCin * 8; + constexpr int BQ = B_E / 8; + constexpr int BCREG = (BQ + NT - 1) / NT; + int cdst[BCREG]; + const __half *csrc[BCREG]; + bool clive[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) { + const int q = tid + r * NT; + const int gc = q / NBRP, rem = q - gc * NBRP; + clive[r] = ((BQ % NT == 0) || q < BQ) && gc < CG; + cdst[r] = clive[r] ? gc * LBO_B + rem * 8 : 0; + csrc[r] = in + (clive[r] ? (size_t)gc * SC8i + + (size_t)nb * 8 + rem * 8 : 0); + } + uint32_t br[BITER][4]; + constexpr int CPA_STAGES = 4; + constexpr int BSTAGES = BC8 ? CPA_STAGES : STAGES; + static_assert(STAGES >= CPA_STAGES, "C8 B-panel pipeline allocation"); + static_assert(ABUF == 2 && BSTAGES == 4, "wg4 unrolled pipeline shape"); + uint64_t dslot[BSTAGES]; +#pragma unroll + for (int s = 0; s < BSTAGES; ++s) + dslot[s] = wg_desc(smem + s * ST_E + ng * WG_BN * 8, LBO_B * 2u, 128u); + uint32_t csm0[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) + csm0[r] = static_cast(__cvta_generic_to_shared(smem + cdst[r])); + uint32_t bsm0[BITER]; +#pragma unroll + for (int t = 0; t < BITER; ++t) + bsm0[t] = static_cast(__cvta_generic_to_shared(smem + bdst[t])); + const size_t CSTEP = (size_t)CG * SC8i; + const __half *crol[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) crol[r] = csrc[r] + CPA_STAGES * CSTEP; + const __half *arol = apt + (size_t)ABUF * (KT * 1024); +#define WG4_CPB(c,BD) { const size_t _o = (size_t)(c) * CG * SC8i; _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16((BD) + cdst[r], csrc[r] + _o, true); } +#define WG4_CPB_R(SLOT) { _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16o<(SLOT) * ST_E * 2>(csm0[r], crol[r], true); _Pragma("unroll") for (int r = 0; r < BCREG; ++r) crol[r] += CSTEP; } +#define WG4_LOADB(c) { const size_t _o = (size_t)(c) * CIB * SCin; _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { const int4 _v = *reinterpret_cast( bbase + _o + boff[t]); br[t][0] = _v.x; br[t][1] = _v.y; br[t][2] = _v.z; br[t][3] = _v.w; } } } +#define WG4_XPOSE(u) { const int _q = lane & 3; uint32_t _x, _y; _x = (_q & 1) ? u[0] : u[1]; _y = (_q & 1) ? u[2] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 1); _y = __shfl_xor_sync(0xffffffffu, _y, 1); if (_q & 1) { u[0] = _x; u[2] = _y; } else { u[1] = _x; u[3] = _y; } _x = (_q & 2) ? u[0] : u[2]; _y = (_q & 2) ? u[1] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 2); _y = __shfl_xor_sync(0xffffffffu, _y, 2); if (_q & 2) { u[0] = _x; u[1] = _y; } else { u[2] = _x; u[3] = _y; } } +#define WG4_STOREB(BD) { _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { uint32_t _u[4]; _u[0] = br[t][0]; _u[1] = br[t][1]; _u[2] = br[t][2]; _u[3] = br[t][3]; WG4_XPOSE(_u) const uint32_t _a = static_cast( __cvta_generic_to_shared((BD) + bdst[t])); asm volatile( "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 " "[%0], {%1,%2,%3,%4};\n" ::"r"(_a), "r"(_u[0]), "r"(_u[1]), "r"(_u[2]), "r"(_u[3])); } } } +#define WG4_STOREB_R(SLOT) { _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { uint32_t _u[4]; _u[0] = br[t][0]; _u[1] = br[t][1]; _u[2] = br[t][2]; _u[3] = br[t][3]; WG4_XPOSE(_u) asm volatile( "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 " "[%0+%5], {%1,%2,%3,%4};\n" ::"r"(bsm0[t]), "r"(_u[0]), "r"(_u[1]), "r"(_u[2]), "r"(_u[3]), "n"((SLOT) * ST_E * 2)); } } } +#define WG4_LOADA_R(BUF) { _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { const int4 _v = *reinterpret_cast(arol + kk * 1024); areg[BUF][kk][0] = _v.x; areg[BUF][kk][1] = _v.y; areg[BUF][kk][2] = _v.z; areg[BUF][kk][3] = _v.w; } arol += KT * 1024; } +#define WG4_MMA_R(BUF,SLOT) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { WGMMA_RS_ANY_T(WG_BN, acc, areg[BUF][kk], dslot[SLOT] + kk, 0); } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } +#define WG4_LOADA(BUF,c) { const __half *_p = apt + (size_t)(c) * (KT * 1024); _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { const int4 _v = *reinterpret_cast(_p + kk * 1024); areg[BUF][kk][0] = _v.x; areg[BUF][kk][1] = _v.y; areg[BUF][kk][2] = _v.z; areg[BUF][kk][3] = _v.w; } } +#define WG4_KEEPA(BUF) { _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) asm volatile("" : "+r"(areg[BUF][kk][0]), "+r"(areg[BUF][kk][1]), "+r"(areg[BUF][kk][2]), "+r"(areg[BUF][kk][3])); } +#define WG4_MMA(BUF,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t db0 = wg_desc((BD) + ng * WG_BN * 8, LBO_B * 2u, 128u); _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { WGMMA_RS_ANY_T(WG_BN, acc, areg[BUF][kk], db0 + kk, 0); } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + if constexpr (PDL && LATE_PDL) PDL_ENTER(); + constexpr int WGD = 1; + if constexpr (BC8) { +#pragma unroll 1 + for (int st = 0; st < CPA_STAGES; ++st) { + if (st < NCH) WG4_CPB(st, smem + st * ST_E) + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } + WG4_LOADA(0, 0) + if (ABUF > 1 && 1 < NCH) WG4_LOADA(1, 1) + if (ABUF > 2 && 2 < NCH) WG4_LOADA(2, 2) + } else { + WG4_LOADB(0) + WG4_LOADA(0, 0) + if (ABUF > 1 && 1 < NCH) WG4_LOADA(1, 1) + if (ABUF > 2 && 2 < NCH) WG4_LOADA(2, 2) + WG4_STOREB(smem) + if (1 < NCH) WG4_LOADB(1) + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + __syncthreads(); + } +#define WG4_BODY(BUF,i,SLOT) { if constexpr (BC8) { asm volatile("cp.async.wait_group %0;\n" : : "n"(CPA_STAGES - 3) : "memory"); __syncthreads(); if ((i) >= 2) { if ((i) + 2 < NCH) { WG4_CPB_R(((SLOT) + 2) % CPA_STAGES) } asm volatile("cp.async.commit_group;\n" ::: "memory"); } } else { if (i) __syncthreads(); } WG4_MMA_R(BUF, SLOT) asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(WGD) : "memory"); if constexpr (NSPLIT > 1) __syncthreads(); if constexpr (!BC8) { const int cb = (i) + 1; if (cb < NCH) { WG4_STOREB_R(((SLOT) + 1) % BSTAGES) asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); if (cb + 1 < NCH) WG4_LOADB(cb + 1) } } if ((i) >= 1) { const int nb2 = ((BUF) + ABUF - 1) % ABUF; if ((i) + ABUF - 1 < NCH) { WG4_KEEPA(nb2) WG4_LOADA_R(nb2) } } } + for (int i = 0; i < NCH; i += 4) { + WG4_BODY(0, i, 0) + if (i + 1 < NCH) WG4_BODY(1, i + 1, 1) + if (i + 2 < NCH) WG4_BODY(0, i + 2, 2) + if (i + 3 < NCH) WG4_BODY(1, i + 3, 3) + } + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + if constexpr (BC8) + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); +#undef WG4_CPB +#undef WG4_CPB_R +#undef WG4_STOREB_R +#undef WG4_LOADA_R +#undef WG4_MMA_R +#undef WG4_LOADB +#undef WG4_XPOSE +#undef WG4_STOREB +#undef WG4_LOADA +#undef WG4_KEEPA +#undef WG4_MMA +#undef WG4_BODY +#define TL4_SW(c,r) ((c) * BM + (((r) & ~7) ^ (((c) & 7) << 3)) + ((r) & 7)) + __syncthreads(); + const int cb2 = (lane & 3) * 2; + const bool snake = PREACT && (m0 + er0) < half; + const float sn_a0 = snake ? row_alpha[er0] : 0.f; + const float sn_r0 = snake ? row_ri[er0] : 0.f; + const float sn_a1 = snake ? row_alpha[er0 + 8] : 0.f; + const float sn_r1 = snake ? row_ri[er0 + 8] : 0.f; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + const int row = er0 + h * 8; + if constexpr (PREACT) { + if (snake) { + const float a = h ? sn_a1 : sn_a0; + const float ri = h ? sn_r1 : sn_r0; + const float s0 = __sinf(a * v.x); + const float s1 = __sinf(a * v.y); + v.x = __fmaf_rn(s0 * s0, ri, v.x); + v.y = __fmaf_rn(s1 * s1, ri, v.y); + } else { + v.x = v.x > 0.f ? v.x : 0.01f * v.x; + v.y = v.y > 0.f ? v.y : 0.01f * v.y; + } + if constexpr (OC8) { + const int c0 = ng * WG_BN + j * 8 + cb2; + Hs[TL4_SW(c0, row)] = __float2half(v.x); + Hs[TL4_SW(c0 + 1, row)] = __float2half(v.y); + } else { + *reinterpret_cast<__half2 *>( + Hs + row * WG_CLD + ng * WG_BN + j * 8 + cb2) = + __floats2half2_rn(v.x, v.y); + } + } else { + *reinterpret_cast( + Cs + row * WG_CLD + ng * WG_BN + j * 8 + cb2) = v; + } + } + } + __syncthreads(); + if constexpr (PREACT) { + if constexpr (OC8) { + constexpr int GRP = BM / 8; + constexpr int TOTG = GRP * BNT; + constexpr int GITER = (TOTG + NT - 1) / NT; + const size_t SC8o = (size_t)SCout * 8; +#pragma unroll + for (int t = 0; t < GITER; ++t) { + const int idx = tid + t * NT; + if (TOTG % NT != 0 && idx >= TOTG) continue; + const int g = idx / BNT; + const int j = idx - g * BNT; + if (!colok[j]) continue; + const int m0g = g * 8; + if (m0 + m0g >= Cout) continue; + *reinterpret_cast( + out + (size_t)((m0 + m0g) >> 3) * SC8o + + (size_t)(nb + j + H) * 8) = + *reinterpret_cast(Hs + TL4_SW(j, m0g)); + } + } else { + static_assert((BNT & 3) == 0, "four-column preactivated epilogue"); + constexpr int NC = BNT / 4; + constexpr int TOT = BM * NC; + for (int q = tid; q < TOT; q += NT) { + const int m = q / NC; + const int lc = (q - m * NC) * 4; + const int co = m0 + m; + const bool any4 = (colok[lc] | colok[lc + 1] | + colok[lc + 2] | colok[lc + 3]) != 0; + if (co >= Cout || !any4) continue; + const __half2 *sp = reinterpret_cast( + Hs + m * WG_CLD + lc); + const __half2 v0 = sp[0], v1 = sp[1]; + __half *op = out + (size_t)co * SCout + nb + lc + H; + const bool ok4 = (colok[lc] & colok[lc + 1] & + colok[lc + 2] & colok[lc + 3]) != 0; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = v0; + reinterpret_cast<__half2 *>(op)[1] = v1; + } else { + if (colok[lc + 0]) op[0] = v0.x; + if (colok[lc + 1]) op[1] = v0.y; + if (colok[lc + 2]) op[2] = v1.x; + if (colok[lc + 3]) op[3] = v1.y; + } + } + } + } else { + constexpr bool COMPACT_EPI = + (WG_BN == 136 || WG_BN == 160 || WG_BN == 208); + if constexpr (COMPACT_EPI) { + static_assert((BNT & 3) == 0, "4-column epilogue groups"); + constexpr int NC = BNT / 4; + constexpr int TOT = BM * NC; + constexpr int MITER = (TOT + NT - 1) / NT; + constexpr int UNR = (MITER >= 8) ? 8 : ((MITER >= 4) ? 4 : + ((MITER >= 2) ? 2 : 1)); +#pragma unroll 1 + for (int t0 = 0; t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int q = tid + (t0 + u) * NT; + if (q >= TOT) continue; + const int m = q / NC; + const int lc = (q - m * NC) * 4; + const int co = m0 + m; + const bool any4 = (colok[lc] | colok[lc + 1] | + colok[lc + 2] | colok[lc + 3]) != 0; + if (co >= Cout || !any4) continue; + v[u] = *reinterpret_cast(Cs + m * WG_CLD + lc); + if constexpr (!PREACT) { + const int n = nb + lc; + const __half2 *rp = reinterpret_cast( + resid + (size_t)co * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int q = tid + (t0 + u) * NT; + if (q >= TOT) continue; + const int m = q / NC; + const int lc = (q - m * NC) * 4; + const int co = m0 + m; + const bool any4 = (colok[lc] | colok[lc + 1] | + colok[lc + 2] | colok[lc + 3]) != 0; + if (co >= Cout || !any4) continue; + const int n = nb + lc; + float4 vv = v[u]; + if constexpr (!PREACT) { + vv.x += __half2float(r0[u].x); vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (co < half) { + const float a = row_alpha[m], ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + __half *op = out + (size_t)co * SCout + n + H; + const bool ok4 = (colok[lc] & colok[lc + 1] & + colok[lc + 2] & colok[lc + 3]) != 0; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + } else { + constexpr int NC = EPN / 4; + constexpr int MSTEP = NT / NC; + constexpr int MITER = BM / MSTEP; + constexpr int UNR = (MITER % 8 == 0) ? 8 + : ((MITER % 4 == 0) ? 4 + : ((MITER % 2 == 0) ? 2 : 1)); + static_assert(MITER % UNR == 0, "row batch must divide row count"); + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && + colok[lc + 2] && colok[lc + 3]; + const bool live4 = colok[lc] || colok[lc + 1] || + colok[lc + 2] || colok[lc + 3]; + const int n = nb + lc; +#pragma unroll 1 + for (int t0 = 0; live4 && t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + v[u] = *reinterpret_cast( + Cs + m * WG_CLD + lc); + if constexpr (!PREACT) { + const int co = m0 + m; + const int cc = co < Cout ? co : 0; + const __half2 *rp = reinterpret_cast( + resid + (size_t)cc * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + const int co = m0 + m; + if (co >= Cout) continue; + float4 vv = v[u]; + if constexpr (!PREACT) { + vv.x += __half2float(r0[u].x); + vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); + vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (co < half) { + const float a = row_alpha[m], ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + __half *op = out + (size_t)co * SCout + n + H; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B; + for (int idx = bxy * NT + tid; idx < total; idx += nbx * NT) { + const int b = idx / Cin, ci = idx - b * Cin; + const int x = b * L + native_true_tokens(b,T); + __half *const dst = ns + ((size_t)b * Cin + ci) * H; + if constexpr (BC8) { + const __half *const src = + in + (size_t)(ci >> 3) * ((size_t)SCin * 8) + + (size_t)x * 8 + (ci & 7); +#pragma unroll + for (int j = 0; j < H; ++j) dst[j] = src[j * 8]; + } else { + const __half *const src = in + (size_t)ci * SCin + x; +#pragma unroll + for (int j = 0; j < H; j += 2) + *reinterpret_cast(dst + j) = + *reinterpret_cast(src + j); + } + } + } + if (pad_src != nullptr) { + const int total = Cout * B; + for (int idx = bxy * NT + tid; idx < total; idx += nbx * NT) { + const int b = idx / Cout, co = idx - b * Cout; + const int x = b * L + pad_off; + const __half *const src = + pad_src + ((size_t)b * Cout + co) * pad_len; + if constexpr (OC8) { + __half *const dst = + out + (size_t)(co >> 3) * ((size_t)SCout * 8) + + (size_t)x * 8 + (co & 7); + for (int j = 0; j < pad_len; ++j) dst[j * 8] = src[j]; + } else { + __half *const dst = out + (size_t)co * SCout + x; + for (int j = 0; j < pad_len; ++j) dst[j] = src[j]; + } + } + } + } +} diff --git a/deploy/nemotron_native/kernels/codec/small/wg5_step_widths.inc b/deploy/nemotron_native/kernels/codec/small/wg5_step_widths.inc new file mode 100644 index 0000000..467481f --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/small/wg5_step_widths.inc @@ -0,0 +1,400 @@ +template +__device__ __noinline__ +void wg5_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int KT = 3; + constexpr int BM = 64 * MSPLIT; + constexpr int BNT = WG_BN; + constexpr int NTS = 128 * MSPLIT; + constexpr int WG_NT = NTS * KSPLIT; + constexpr int NCRAW = BNT / 4; + constexpr int EPN = + (WG_NT % NCRAW == 0) ? BNT : (((BNT + 127) / 128) * 128); + constexpr int WG_CLD_R = (OC8 ? WG_BN : EPN) + 8; + constexpr int WG_CLD = (WG_CLD_R % 16 == 0) ? WG_CLD_R + 4 : WG_CLD_R; + constexpr int CG = CIB / 8; + constexpr int CSUB = CIB / 16; + constexpr int AK = KT * CSUB; + constexpr int NBRP = BC8 ? (((BNT + KT - 1) + 7) & ~7) + : (((BNT + KT - 1) + 31) & ~31); + constexpr int UCOL = NBRP / 32; + constexpr int NUNIT = CG * UCOL; + constexpr int NWARP = NTS / 32; + constexpr int BITER = (NUNIT + NWARP - 1) / NWARP; + constexpr int B_E = CIB * NBRP; + constexpr int ST_E = SB * B_E; + constexpr int LBO_B = NBRP * 8; + constexpr int ABUF = 3; + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *smem = reinterpret_cast<__half *>(arena); + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int kgrp = tid / NTS; + const int stid = tid - kgrp * NTS; + const int mg = stid >> 7; + const int warp = (stid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * BM; + const int nb = blockIdx.y * BNT; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + constexpr int NCH = (CIN + CIB - 1) / CIB; + static_assert(NCH % KSPLIT == 0, "wg5 K-split must divide the chunk count"); + if (tid < BM) { + const int co = m0 + tid; + if (co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + for (int i = tid; i < EPN; i += WG_NT) { + const int n = nb + i; + colok[i] = (i < BNT && n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = mg * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (kgrp == 0 && eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (kgrp == 0 && eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + const int mtile = blockIdx.x * MSPLIT + mg; + const __half *const apt = w + ((size_t)mtile * NCH * AK) * 1024 + + (size_t)(stid & 127) * 8; + uint32_t areg[ABUF][AK][4]; + const int wi = stid >> 5; + int bdst[BITER], boff[BITER]; + bool blive[BITER]; + const __half *const bbase = in + (size_t)(lane >> 2) * SCin + nb + + 8 * (lane & 3); +#pragma unroll + for (int t = 0; t < BITER; ++t) { + const int u = wi + t * NWARP; + const int gc = (u < NUNIT) ? u / UCOL : 0; + const int cu = (u < NUNIT) ? u - gc * UCOL : 0; + blive[t] = u < NUNIT; + bdst[t] = gc * LBO_B + (cu * 32 + lane) * 8; + boff[t] = gc * 8 * SCin + cu * 32; + } + const size_t SC8i = (size_t)SCin * 8; + constexpr int BQ = B_E / 8; + constexpr int BCREG = (BQ + NTS - 1) / NTS; + int cdst[BCREG]; + const __half *csrc[BCREG]; + bool clive[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) { + const int q = stid + r * NTS; + const int gc = q / NBRP, rem = q - gc * NBRP; + clive[r] = ((BQ % NTS == 0) || q < BQ) && gc < CG; + cdst[r] = clive[r] ? gc * LBO_B + rem * 8 : 0; + csrc[r] = in + (clive[r] ? (size_t)gc * SC8i + + (size_t)nb * 8 + rem * 8 : 0); + } + uint32_t br[BITER][4]; + __half *const kbB = smem + kgrp * ST_E; + const unsigned barid = 1u + kgrp + (unsigned)(KSPLIT * native_subid()); + uint64_t dbuf[ABUF]; +#pragma unroll + for (int s = 0; s < ABUF; ++s) + dbuf[s] = wg_desc(kbB + s * B_E, LBO_B * 2u, 128u); + uint32_t csm0[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) + csm0[r] = static_cast( + __cvta_generic_to_shared(kbB + cdst[r])); + const size_t CSTEP = (size_t)KSPLIT * CG * SC8i; + constexpr int ASTEP = KSPLIT * AK * 1024; + const __half *crol[BCREG]; +#pragma unroll + for (int r = 0; r < BCREG; ++r) + crol[r] = csrc[r] + (size_t)(ABUF * KSPLIT + kgrp) * CG * SC8i; + const __half *arol = apt + (size_t)(ABUF * KSPLIT + kgrp) * (AK * 1024); + int crolc = ABUF * KSPLIT + kgrp; +#define WG5_BAR() asm volatile("bar.sync %0, %1;\n" ::"r"(barid), "r"(NTS)) +#define WG5_CPB(c,BD) { const size_t _o = (size_t)(c) * CG * SC8i; _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16((BD) + cdst[r], csrc[r] + _o, true); } +#define WG5_LOADA(BUF,c) { const __half *_p = apt + (size_t)(c) * (AK * 1024); _Pragma("unroll") for (int kk = 0; kk < AK; ++kk) { const int4 _v = *reinterpret_cast(_p + kk * 1024); areg[BUF][kk][0] = _v.x; areg[BUF][kk][1] = _v.y; areg[BUF][kk][2] = _v.z; areg[BUF][kk][3] = _v.w; } } +#define WG5_KEEPA(BUF) { _Pragma("unroll") for (int kk = 0; kk < AK; ++kk) asm volatile("" : "+r"(areg[BUF][kk][0]), "+r"(areg[BUF][kk][1]), "+r"(areg[BUF][kk][2]), "+r"(areg[BUF][kk][3])); } +#define WG5_LOADB(c) { const size_t _o = (size_t)(c) * CIB * SCin; _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { const int4 _v = *reinterpret_cast( bbase + _o + boff[t]); br[t][0] = _v.x; br[t][1] = _v.y; br[t][2] = _v.z; br[t][3] = _v.w; } } } +#define WG5_BXPOSE(u) { const int _q = lane & 3; uint32_t _x, _y; _x = (_q & 1) ? u[0] : u[1]; _y = (_q & 1) ? u[2] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 1); _y = __shfl_xor_sync(0xffffffffu, _y, 1); if (_q & 1) { u[0] = _x; u[2] = _y; } else { u[1] = _x; u[3] = _y; } _x = (_q & 2) ? u[0] : u[2]; _y = (_q & 2) ? u[1] : u[3]; _x = __shfl_xor_sync(0xffffffffu, _x, 2); _y = __shfl_xor_sync(0xffffffffu, _y, 2); if (_q & 2) { u[0] = _x; u[1] = _y; } else { u[2] = _x; u[3] = _y; } } +#define WG5_STOREB(BD) { _Pragma("unroll") for (int t = 0; t < BITER; ++t) { if (blive[t]) { uint32_t _u[4]; _u[0] = br[t][0]; _u[1] = br[t][1]; _u[2] = br[t][2]; _u[3] = br[t][3]; WG5_BXPOSE(_u) const uint32_t _a = static_cast( __cvta_generic_to_shared((BD) + bdst[t])); asm volatile( "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 " "[%0], {%1,%2,%3,%4};\n" ::"r"(_a), "r"(_u[0]), "r"(_u[1]), "r"(_u[2]), "r"(_u[3])); } } } +#define WG5_MMA(BUF) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); _Pragma("unroll") for (int h = 0; h < CSUB; ++h) { _Pragma("unroll") for (int kk = 0; kk < KT; ++kk) { WGMMA_RS_ANY_T(WG_BN, acc, areg[BUF][h * KT + kk], dbuf[BUF] + kk + h * (LBO_B / 4), 0); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } +#define WG5_CPB_R(SLOT) { _Pragma("unroll") for (int r = 0; r < BCREG; ++r) if (clive[r]) cpa16o<(SLOT) * B_E * 2>(csm0[r], crol[r], true); } +#define WG5_LOADA_R(BUF) { _Pragma("unroll") for (int kk = 0; kk < AK; ++kk) { const int4 _v = *reinterpret_cast(arol + kk * 1024); areg[BUF][kk][0] = _v.x; areg[BUF][kk][1] = _v.y; areg[BUF][kk][2] = _v.z; areg[BUF][kk][3] = _v.w; } } +#define WG5_ROLL() { crolc += KSPLIT; arol += ASTEP; if constexpr (BC8) { _Pragma("unroll") for (int r = 0; r < BCREG; ++r) crol[r] += CSTEP; } } + if constexpr (PDL) { + (void)0; + if (tid == 0) + (void)0; + } + static_assert(SB == 3 && WGD == 1 && ABUF == 3, "wg5 pipeline shape"); + if constexpr (BC8) { +#pragma unroll 1 + for (int s = 0; s < SB; ++s) { + const int c = s * KSPLIT + kgrp; + if (c < NCH) WG5_CPB(c, kbB + s * B_E) + asm volatile("cp.async.commit_group;\n" ::: "memory"); + } + } else { + WG5_LOADB(kgrp) +#pragma unroll 1 + for (int s = 0; s < SB; ++s) { + const int c = s * KSPLIT + kgrp; + if (c < NCH) WG5_STOREB(kbB + s * B_E) + if (c + KSPLIT < NCH) WG5_LOADB(c + KSPLIT) + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + } + WG5_LOADA(0, kgrp) + if (KSPLIT + kgrp < NCH) WG5_LOADA(1, KSPLIT + kgrp) + if (2 * KSPLIT + kgrp < NCH) WG5_LOADA(2, 2 * KSPLIT + kgrp) + constexpr int NIT = NCH / KSPLIT; + static_assert(NIT % ABUF == 0, "wg5 mainloop unroll shape"); +#define WG5_BODY(BUF,i) { if constexpr (BC8) asm volatile("cp.async.wait_group %0;\n" : : "n"(SB - 2) : "memory"); WG5_BAR(); WG5_MMA(BUF) asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(WGD) : "memory"); if constexpr (NTS > 128) WG5_BAR(); if ((i) >= 1) { const int nb2 = ((BUF) + 2) % ABUF; if (crolc < NCH) { if constexpr (BC8) { WG5_CPB_R(nb2) } else { WG5_STOREB(kbB + nb2 * B_E) asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); if (crolc + KSPLIT < NCH) WG5_LOADB(crolc + KSPLIT) } WG5_KEEPA(nb2) WG5_LOADA_R(nb2) } WG5_ROLL() if constexpr (BC8) asm volatile("cp.async.commit_group;\n" ::: "memory"); } } +#pragma unroll + for (int i = 0; i < NIT; i += 3) { + WG5_BODY(0, i) + WG5_BODY(1, i + 1) + WG5_BODY(2, i + 2) + } + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + if constexpr (BC8) + asm volatile("cp.async.wait_group 0;\n" ::: "memory"); +#undef WG5_BXPOSE +#undef WG5_CPB +#undef WG5_CPB_R +#undef WG5_LOADA_R +#undef WG5_ROLL +#undef WG5_LOADB +#undef WG5_STOREB +#undef WG5_MMA +#undef WG5_LOADA +#undef WG5_KEEPA +#undef WG5_BODY +#undef WG5_BAR + __syncthreads(); + const int cb2 = (lane & 3) * 2; + if constexpr (OC8) { + __half *Ck = smem + kgrp * BM * WG_CLD; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + *reinterpret_cast<__half2 *>(Ck + (er0 + h * 8) * WG_CLD + + j * 8 + cb2) = + __floats2half2_rn(acc[j * 4 + h * 2], + acc[j * 4 + h * 2 + 1]); + } + } + } else { + float *Ck = Cs + kgrp * BM * WG_CLD; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], + acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Ck + (er0 + h * 8) * WG_CLD + + j * 8 + cb2) = v; + } + } + } + __syncthreads(); + if constexpr (OC8) { + constexpr int GRP = BM / 8; + constexpr int TOTG = GRP * WG_BN; + constexpr int GITER = (TOTG + WG_NT - 1) / WG_NT; + const size_t SC8o = (size_t)SCout * 8; +#pragma unroll + for (int t = 0; t < GITER; ++t) { + const int idx = tid + t * WG_NT; + if (TOTG % WG_NT != 0 && idx >= TOTG) continue; + const int j = idx - (idx / WG_BN) * WG_BN, g = idx / WG_BN; + if (!colok[j]) continue; + const int m0g = g * 8; + if (m0 + m0g >= Cout) continue; + const int n = nb + j; + __half rr[8]; + if constexpr (CIN == 1536) { + if constexpr (RC8) { + *reinterpret_cast(rr) = + *reinterpret_cast( + resid + (size_t)((m0 + m0g) >> 3) * SC8o + + (size_t)(n + H) * 8); + } else { +#pragma unroll + for (int h = 0; h < 8; ++h) + rr[h] = resid[(size_t)(m0 + m0g + h) * SCout + n + H]; + } + } + uint32_t pk[4]; + const bool snake = (m0 + m0g) < half; +#pragma unroll + for (int h = 0; h < 8; h += 2) { + float f[2]; +#pragma unroll + for (int e = 0; e < 2; ++e) { + const int m = m0g + h + e; + float a = __half2float(smem[m * WG_CLD + j]); +#pragma unroll + for (int kg = 1; kg < KSPLIT; ++kg) + a += __half2float( + smem[kg * BM * WG_CLD + m * WG_CLD + j]); + if constexpr (CIN == 1536) + a += __half2float(rr[h + e]); + if (snake) { + const float sn = __sinf(row_alpha[m] * a); + a = __fmaf_rn(sn * sn, row_ri[m], a); + } else { + a = a > 0.f ? a : 0.01f * a; + } + f[e] = a; + } + const __half2 p = __floats2half2_rn(f[0], f[1]); + pk[h >> 1] = *reinterpret_cast(&p); + } + *reinterpret_cast( + out + (size_t)((m0 + m0g) >> 3) * SC8o + + (size_t)(n + H) * 8) = make_int4(pk[0], pk[1], pk[2], pk[3]); + } + } else { + constexpr int NC = EPN / 4; + constexpr int MSTEP = WG_NT / NC; + constexpr int MITER = BM / MSTEP; + constexpr int UNR = (MITER % 4 == 0) ? 4 : ((MITER % 2 == 0) ? 2 : 1); + static_assert(MITER % UNR == 0, "row batch must divide the row count"); + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + const bool live4 = colok[lc] || colok[lc + 1] || colok[lc + 2] || + colok[lc + 3]; + const int n = nb + lc; +#pragma unroll 1 + for (int t0 = 0; live4 && t0 < MITER; t0 += UNR) { + float4 v[UNR]; + __half2 r0[UNR], r1[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + v[u] = *reinterpret_cast(Cs + m * WG_CLD + lc); + if constexpr (CIN == 1536) { + const int co = m0 + m; + const int cc = co < Cout ? co : 0; + if constexpr (RC8) { + const __half *rp = + resid + (size_t)(cc >> 3) * ((size_t)SCout * 8) + + (size_t)(n + H) * 8 + (cc & 7); + r0[u].x = rp[0]; r0[u].y = rp[8]; + r1[u].x = rp[16]; r1[u].y = rp[24]; + } else { + const __half2 *rp = reinterpret_cast( + resid + (size_t)cc * SCout + n + H); + r0[u] = rp[0]; + r1[u] = rp[1]; + } + } + } +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const int m = lr + (t0 + u) * MSTEP; + const int co = m0 + m; + if (co >= Cout) continue; + float4 vv = v[u]; +#pragma unroll + for (int g = 1; g < KSPLIT; ++g) { + float4 p = *reinterpret_cast( + Cs + g * BM * WG_CLD + m * WG_CLD + lc); + vv.x += p.x; vv.y += p.y; vv.z += p.z; vv.w += p.w; + } + if constexpr (CIN == 1536) { + vv.x += __half2float(r0[u].x); vv.y += __half2float(r0[u].y); + vv.z += __half2float(r1[u].x); vv.w += __half2float(r1[u].y); + } + float *vp = reinterpret_cast(&vv); + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + __half *op = out + (size_t)co * SCout + n + H; + if (ok4) { + reinterpret_cast<__half2 *>(op)[0] = + __floats2half2_rn(vv.x, vv.y); + reinterpret_cast<__half2 *>(op)[1] = + __floats2half2_rn(vv.z, vv.w); + } else { +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + static_assert(CIN == 768 || CIN == 1536); + constexpr int COUT = (CIN == 768) ? 1536 : 768; + if (ns != nullptr) { + const int total = CIN * B; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int b = idx / CIN, ci = idx - b * CIN; + const int x = b * L + native_true_tokens(b,T); + __half *const dst = ns + ((size_t)b * CIN + ci) * H; + if constexpr (BC8) { + const __half *const src = + in + (size_t)(ci >> 3) * ((size_t)SCin * 8) + + (size_t)x * 8 + (ci & 7); +#pragma unroll + for (int j = 0; j < H; ++j) dst[j] = src[j * 8]; + } else { + const __half *const src = in + (size_t)ci * SCin + x; +#pragma unroll + for (int j = 0; j < H; j += 2) + *reinterpret_cast(dst + j) = + *reinterpret_cast(src + j); + } + } + } + if (pad_src != nullptr) { + const int total = COUT * B; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int b = idx / COUT, co = idx - b * COUT; + const int x = b * L + pad_off; + const __half *const src = + pad_src + ((size_t)b * COUT + co) * pad_len; + if constexpr (OC8) { + __half *const dst = + out + (size_t)(co >> 3) * ((size_t)SCout * 8) + + (size_t)x * 8 + (co & 7); + for (int j = 0; j < pad_len; ++j) dst[j * 8] = src[j]; + } else { + __half *const dst = out + (size_t)co * SCout + x; + for (int j = 0; j < pad_len; ++j) dst[j] = src[j]; + } + } + } + } +} diff --git a/deploy/nemotron_native/kernels/codec/small/wg_kernel.inc b/deploy/nemotron_native/kernels/codec/small/wg_kernel.inc new file mode 100644 index 0000000..2d5c766 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/small/wg_kernel.inc @@ -0,0 +1,182 @@ +__device__ __forceinline__ uint64_t wg_desc(const void *p, uint32_t lbo_bytes, + uint32_t sbo_bytes) { + uint32_t a = static_cast(__cvta_generic_to_shared(p)); + uint64_t d = static_cast((a & 0x3FFFFu) >> 4); + d |= static_cast((lbo_bytes & 0x3FFFFu) >> 4) << 16; + d |= static_cast((sbo_bytes & 0x3FFFFu) >> 4) << 32; + return d; +} +#include "wgmma_macro.inc" +constexpr int WG_SLOT = 72; +template +__device__ __noinline__ void wg_conv_kernel( + const __half *__restrict__ in, const __half *__restrict__ w, + const float *__restrict__ bias, const float *__restrict__ alpha, + const __half *__restrict__ resid, __half *__restrict__ out, + const __half *__restrict__ pad_src, __half *__restrict__ ns, + int Cin, int Cout, int T, int L, int SCin, int SCout, int N, + int do_snake, int pad_off, int pad_len) { + constexpr int WG_CLD = WG_BN + 8; + constexpr int WG_NT = 128 * NWG; + constexpr int WG_BM = 64 * NWG; + constexpr int KG = BK / 8; + constexpr int A_E = (WG_BM / 8) * KG * WG_SLOT; + constexpr int B_E = KG * (WG_BN / 8) * WG_SLOT; + constexpr int SBO_AE = KG * WG_SLOT; + constexpr int LBO_BE = (WG_BN / 8) * WG_SLOT; + char *arena=reinterpret_cast(native_arena()); + char *colok=reinterpret_cast(native_aux()); + float *row_alpha=reinterpret_cast(native_aux()+1024); + float *row_ri=reinterpret_cast(native_aux()+2048); + __half *As0 = reinterpret_cast<__half *>(arena); + __half *Bs0 = As0 + A_E; + __half *As1 = (STAGES == 2) ? Bs0 + B_E : As0; + __half *Bs1 = (STAGES == 2) ? Bs0 + B_E + A_E : Bs0; + float *Cs = reinterpret_cast(arena); + const int tid = threadIdx.x; + const int wgid = tid >> 7; + const int warp = (tid & 127) >> 5, lane = tid & 31; + const int m0 = blockIdx.x * WG_BM; + constexpr int WGROWS = 8; + const int n0 = blockIdx.y * WG_BN; + const int H = KT - 1; + const int Kg = Cin * KT; + const int half = Cout >> 1; + if (tid < WG_BM) { + const int co = m0 + tid; + if (do_snake && co < Cout && co < half) { + const float a = alpha[co]; + row_alpha[tid] = a; + row_ri[tid] = __frcp_rn(a + 1e-9f); + } + } + if (tid < WG_BN) { + const int n = n0 + tid; + colok[tid] = (n < N && (n % L) < T) ? 1 : 0; + } + const int er0 = wgid * 64 + warp * 16 + (lane >> 2); + const int eco0 = m0 + er0, eco1 = m0 + er0 + 8; + const float bias0 = (eco0 < Cout) ? bias[eco0] : 0.f; + const float bias1 = (eco1 < Cout) ? bias[eco1] : 0.f; + float acc[WG_BN / 2]; +#pragma unroll + for (int i = 0; i < WG_BN / 2; ++i) acc[i] = ((i & 2) == 0) ? bias0 : bias1; + constexpr int ATOT = WG_BM * KG; + constexpr int AREG = (ATOT + WG_NT - 1) / WG_NT; + constexpr int CIB = BK / KT; + constexpr int BTOT = CIB * (WG_BN / 8); + constexpr int BREG = (BTOT + WG_NT - 1) / WG_NT; + int4 ar[AREG]; + int4 br0[BREG], br1[BREG]; + uint32_t edge[BREG]; +#define WG_LOADG(kg0) { const int _ci0 = (kg0) / KT; _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; ar[r] = make_int4(0, 0, 0, 0); if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; const int co = m0 + m, kg = (kg0) + v * 8; if (co < Cout && kg < Kg) ar[r] = *reinterpret_cast( w + (size_t)co * Kg + kg); } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; br0[r] = make_int4(0, 0, 0, 0); if constexpr (KT == 3) edge[r] = 0; else br1[r] = make_int4(0, 0, 0, 0); if (BTOT % WG_NT == 0 || idx < BTOT) { const int ci = _ci0 + idx / (WG_BN / 8); const int nc = idx % (WG_BN / 8); if (ci < Cin) { const int4 *src = reinterpret_cast( in + (size_t)ci * SCin + n0 + nc * 8); br0[r] = src[0]; if constexpr (KT == 3) { if (nc == WG_BN / 8 - 1) edge[r] = src[1].x; } else { br1[r] = src[1]; } } } } } + WG_LOADG(0) +#define WG_STS(AD,BD) { _Pragma("unroll") for (int r = 0; r < AREG; ++r) { const int idx = tid + r * WG_NT; if (ATOT % WG_NT == 0 || idx < ATOT) { const int m = idx / KG, v = idx - (idx / KG) * KG; *reinterpret_cast( (AD) + (m >> 3) * SBO_AE + v * WG_SLOT + (m & 7) * 8) = ar[r]; } } _Pragma("unroll") for (int r = 0; r < BREG; ++r) { const int idx = tid + r * WG_NT; if (BTOT % WG_NT == 0 || idx < BTOT) { const int cil = idx / (WG_BN / 8), nc = idx % (WG_BN / 8); unsigned int u[8]; u[0] = br0[r].x; u[1] = br0[r].y; u[2] = br0[r].z; u[3] = br0[r].w; if constexpr (KT == 3) { u[4] = __shfl_down_sync(0xffffffffu, u[0], 1, WG_BN / 8); if (nc == WG_BN / 8 - 1) u[4] = edge[r]; u[5] = u[6] = u[7] = 0; } else { u[4] = br1[r].x; u[5] = br1[r].y; u[6] = br1[r].z; u[7] = br1[r].w; } _Pragma("unroll") for (int k = 0; k < KT; ++k) { const int kh = k >> 1; int4 o; if (k & 1) { o.x = __funnelshift_r(u[kh + 0], u[kh + 1], 16); o.y = __funnelshift_r(u[kh + 1], u[kh + 2], 16); o.z = __funnelshift_r(u[kh + 2], u[kh + 3], 16); o.w = __funnelshift_r(u[kh + 3], u[kh + 4], 16); } else { o.x = u[kh + 0]; o.y = u[kh + 1]; o.z = u[kh + 2]; o.w = u[kh + 3]; } const int row = cil * KT + k; *reinterpret_cast((BD) + nc * WG_SLOT + (row >> 3) * LBO_BE + (row & 7) * 8) = o; } } } if constexpr (WG_BN == 64) { asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); } } +#define WG_MMA(AD,BD) { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); const uint64_t da0 = wg_desc((AD) + wgid * WGROWS * SBO_AE, WG_SLOT * 2, SBO_AE * 2); const uint64_t db0 = wg_desc((BD), LBO_BE * 2, WG_SLOT * 2); _Pragma("unroll") for (int kk = 0; kk < BK; kk += 16) { const uint64_t da = da0 + (kk >> 3) * (WG_SLOT / 8); const uint64_t db = db0 + (kk >> 3) * (LBO_BE / 8); if constexpr (WG_BN == 64) { WGMMA_M64N64K16(acc, da, db); } else { WGMMA_M64N128K16(acc, da, db); } } asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); } + WG_STS(As0, Bs0) + if (BK < Kg) WG_LOADG(BK) + __syncthreads(); + int cur = 0; + for (int kg0 = 0; kg0 < Kg; kg0 += BK) { + if (cur == 0) { WG_MMA(As0, Bs0) } else { WG_MMA(As1, Bs1) } + if (kg0 + BK < Kg) { + if (STAGES == 2) + asm volatile("wgmma.wait_group.sync.aligned 1;\n" ::: "memory"); + else + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + __syncthreads(); + if (cur == 0) { WG_STS(As1, Bs1) } else { WG_STS(As0, Bs0) } + if (kg0 + 2 * BK < Kg) WG_LOADG(kg0 + 2 * BK) + __syncthreads(); + } else { + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); + } + cur ^= 1; + } +#undef WG_LOADG +#undef WG_STS +#undef WG_MMA + __syncthreads(); + const int r0 = er0; + const int cb = (lane & 3) * 2; +#pragma unroll + for (int j = 0; j < WG_BN / 8; ++j) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + float2 v = make_float2(acc[j * 4 + h * 2], acc[j * 4 + h * 2 + 1]); + *reinterpret_cast(Cs + (r0 + h * 8) * WG_CLD + j * 8 + cb) = v; + } + } + __syncthreads(); + constexpr int NC = WG_BN / 4; + const int lc = (tid % NC) * 4; + const int lr = tid / NC; + const bool ok4 = colok[lc] && colok[lc + 1] && colok[lc + 2] && colok[lc + 3]; + for (int m = lr; m < WG_BM; m += WG_NT / NC) { + const int co = m0 + m; + if (co >= Cout) continue; + const int n = n0 + lc; + float4 v = *reinterpret_cast(Cs + m * WG_CLD + lc); + __half *op = out + (size_t)co * SCout + n + H; + if (resid != nullptr) { + const __half2 *rp = reinterpret_cast( + resid + (size_t)co * SCout + n + H); + __half2 r0h = rp[0], r1h = rp[1]; + v.x += __half2float(r0h.x); v.y += __half2float(r0h.y); + v.z += __half2float(r1h.x); v.w += __half2float(r1h.y); + } + if (do_snake) { + if (co < half) { + const float a = row_alpha[m]; + const float ri = row_ri[m]; + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float sn = __sinf(a * vp[e]); + vp[e] = __fmaf_rn(sn * sn, ri, vp[e]); + } + } else { + float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + vp[e] = vp[e] > 0.f ? vp[e] : 0.01f * vp[e]; + } + } + if (ok4) { + __half2 o0 = __floats2half2_rn(v.x, v.y); + __half2 o1 = __floats2half2_rn(v.z, v.w); + reinterpret_cast<__half2 *>(op)[0] = o0; + reinterpret_cast<__half2 *>(op)[1] = o1; + } else { + const float *vp = reinterpret_cast(&v); +#pragma unroll + for (int e = 0; e < 4; ++e) + if (colok[lc + e]) op[e] = __float2half(vp[e]); + } + } + { + const int nbx = gridDim.x * gridDim.y; + const int bxy = blockIdx.x + blockIdx.y * gridDim.x; + const int B = N / L; + if (ns != nullptr) { + const int total = Cin * B * H; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % H, r = idx / H; + const int b = r % B, ci = r / B; + ns[((size_t)b * Cin + ci) * H + j] = + in[(size_t)ci * SCin + b * L + native_true_tokens(b,T) + j]; + } + } + if (pad_src != nullptr) { + const int total = Cout * B * pad_len; + for (int idx = bxy * WG_NT + tid; idx < total; + idx += nbx * WG_NT) { + const int j = idx % pad_len, r = idx / pad_len; + const int b = r % B, co = r / B; + out[(size_t)co * SCout + b * L + pad_off + j] = + pad_src[((size_t)b * Cout + co) * pad_len + j]; + } + } + } +} diff --git a/deploy/nemotron_native/kernels/codec/small/wgmma_macro.inc b/deploy/nemotron_native/kernels/codec/small/wgmma_macro.inc new file mode 100644 index 0000000..db624e9 --- /dev/null +++ b/deploy/nemotron_native/kernels/codec/small/wgmma_macro.inc @@ -0,0 +1,102 @@ +#define WGMMA_M64N8K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " "{" "%0,%1,%2,%3}" ", %4, %5, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N8K16(D,DA,DB) WGMMA_M64N8K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N8K16_ZERO_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " "{" "%0,%1,%2,%3}" ", %4, %5, p, 1, 1, " #TA ", " #TB ";\n}\n" : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) : "l"(DA), "l"(DB), "r"(0)) +#define WGMMA_M64N16K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %10, 0;\n" "wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7}" ", %8, %9, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N16K16(D,DA,DB) WGMMA_M64N16K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N24K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %14, 0;\n" "wgmma.mma_async.sync.aligned.m64n24k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11}" ", %12, %13, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N24K16(D,DA,DB) WGMMA_M64N24K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N32K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %18, 0;\n" "wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15}" ", %16, %17, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N32K16(D,DA,DB) WGMMA_M64N32K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N40K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %22, 0;\n" "wgmma.mma_async.sync.aligned.m64n40k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19}" ", %20, %21, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N40K16(D,DA,DB) WGMMA_M64N40K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N48K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %26, 0;\n" "wgmma.mma_async.sync.aligned.m64n48k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N48K16(D,DA,DB) WGMMA_M64N48K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N56K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %30, 0;\n" "wgmma.mma_async.sync.aligned.m64n56k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N56K16(D,DA,DB) WGMMA_M64N56K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N64K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %34, 0;\n" "wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31}" ", %32, %33, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N64K16(D,DA,DB) WGMMA_M64N64K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N72K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %38, 0;\n" "wgmma.mma_async.sync.aligned.m64n72k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35}" ", %36, %37, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N72K16(D,DA,DB) WGMMA_M64N72K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N80K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %42, 0;\n" "wgmma.mma_async.sync.aligned.m64n80k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39}" ", %40, %41, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N80K16(D,DA,DB) WGMMA_M64N80K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N88K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %46, 0;\n" "wgmma.mma_async.sync.aligned.m64n88k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43}" ", %44, %45, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N88K16(D,DA,DB) WGMMA_M64N88K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N96K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %50, 0;\n" "wgmma.mma_async.sync.aligned.m64n96k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47}" ", %48, %49, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N96K16(D,DA,DB) WGMMA_M64N96K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N104K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %54, 0;\n" "wgmma.mma_async.sync.aligned.m64n104k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51}" ", %52, %53, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N104K16(D,DA,DB) WGMMA_M64N104K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N112K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %58, 0;\n" "wgmma.mma_async.sync.aligned.m64n112k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55}" ", %56, %57, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N112K16(D,DA,DB) WGMMA_M64N112K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N120K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %62, 0;\n" "wgmma.mma_async.sync.aligned.m64n120k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59}" ", %60, %61, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N120K16(D,DA,DB) WGMMA_M64N120K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N128K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %66, 0;\n" "wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63}" ", %64, %65, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N128K16(D,DA,DB) WGMMA_M64N128K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N136K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %70, 0;\n" "wgmma.mma_async.sync.aligned.m64n136k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67}" ", %68, %69, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N136K16(D,DA,DB) WGMMA_M64N136K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N144K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %74, 0;\n" "wgmma.mma_async.sync.aligned.m64n144k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71}" ", %72, %73, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N144K16(D,DA,DB) WGMMA_M64N144K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N152K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %78, 0;\n" "wgmma.mma_async.sync.aligned.m64n152k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75}" ", %76, %77, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N152K16(D,DA,DB) WGMMA_M64N152K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N160K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %82, 0;\n" "wgmma.mma_async.sync.aligned.m64n160k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79}" ", %80, %81, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N160K16(D,DA,DB) WGMMA_M64N160K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N168K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %86, 0;\n" "wgmma.mma_async.sync.aligned.m64n168k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83}" ", %84, %85, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N168K16(D,DA,DB) WGMMA_M64N168K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N176K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %90, 0;\n" "wgmma.mma_async.sync.aligned.m64n176k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87}" ", %88, %89, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N176K16(D,DA,DB) WGMMA_M64N176K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N184K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %94, 0;\n" "wgmma.mma_async.sync.aligned.m64n184k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91}" ", %92, %93, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N184K16(D,DA,DB) WGMMA_M64N184K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N192K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %98, 0;\n" "wgmma.mma_async.sync.aligned.m64n192k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95}" ", %96, %97, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N192K16(D,DA,DB) WGMMA_M64N192K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N200K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %102, 0;\n" "wgmma.mma_async.sync.aligned.m64n200k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99}" ", %100, %101, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N200K16(D,DA,DB) WGMMA_M64N200K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N208K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %106, 0;\n" "wgmma.mma_async.sync.aligned.m64n208k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103}" ", %104, %105, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N208K16(D,DA,DB) WGMMA_M64N208K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N216K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %110, 0;\n" "wgmma.mma_async.sync.aligned.m64n216k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107}" ", %108, %109, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N216K16(D,DA,DB) WGMMA_M64N216K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N224K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %114, 0;\n" "wgmma.mma_async.sync.aligned.m64n224k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111}" ", %112, %113, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N224K16(D,DA,DB) WGMMA_M64N224K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N232K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %118, 0;\n" "wgmma.mma_async.sync.aligned.m64n232k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115}" ", %116, %117, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N232K16(D,DA,DB) WGMMA_M64N232K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N240K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %122, 0;\n" "wgmma.mma_async.sync.aligned.m64n240k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119}" ", %120, %121, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N240K16(D,DA,DB) WGMMA_M64N240K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N248K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %126, 0;\n" "wgmma.mma_async.sync.aligned.m64n248k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123}" ", %124, %125, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N248K16(D,DA,DB) WGMMA_M64N248K16_T(D, DA, DB, 0, 1) +#define WGMMA_M64N256K16_T(D,DA,DB,TA,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %130, 0;\n" "wgmma.mma_async.sync.aligned.m64n256k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123,%124,%125,%126,%127}" ", %128, %129, p, 1, 1, " #TA ", " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]), "+f"(D[124]), "+f"(D[125]), "+f"(D[126]), "+f"(D[127]) : "l"(DA), "l"(DB), "r"(1)) +#define WGMMA_M64N256K16(D,DA,DB) WGMMA_M64N256K16_T(D, DA, DB, 0, 1) +#define WGMMA_ANY_T(NW,D,DA,DB,TA,TB) do { if constexpr (NW == 8) { WGMMA_M64N8K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 16) { WGMMA_M64N16K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 24) { WGMMA_M64N24K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 32) { WGMMA_M64N32K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 40) { WGMMA_M64N40K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 48) { WGMMA_M64N48K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 56) { WGMMA_M64N56K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 64) { WGMMA_M64N64K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 72) { WGMMA_M64N72K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 80) { WGMMA_M64N80K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 88) { WGMMA_M64N88K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 96) { WGMMA_M64N96K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 104) { WGMMA_M64N104K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 112) { WGMMA_M64N112K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 120) { WGMMA_M64N120K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 128) { WGMMA_M64N128K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 136) { WGMMA_M64N136K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 144) { WGMMA_M64N144K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 152) { WGMMA_M64N152K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 160) { WGMMA_M64N160K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 168) { WGMMA_M64N168K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 176) { WGMMA_M64N176K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 184) { WGMMA_M64N184K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 192) { WGMMA_M64N192K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 200) { WGMMA_M64N200K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 208) { WGMMA_M64N208K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 216) { WGMMA_M64N216K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 224) { WGMMA_M64N224K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 232) { WGMMA_M64N232K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 240) { WGMMA_M64N240K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 248) { WGMMA_M64N248K16_T(D, DA, DB, TA, TB); } else if constexpr (NW == 256) { WGMMA_M64N256K16_T(D, DA, DB, TA, TB); } } while (0) +#define WGMMA_ANY(NW,D,DA,DB) WGMMA_ANY_T(NW, D, DA, DB, 0, 1) +#define WGMMA_RS_M64N96K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %53, 0;\n" "wgmma.mma_async.sync.aligned.m64n96k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47}" ", {%48,%49,%50,%51}, %52, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N104K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %57, 0;\n" "wgmma.mma_async.sync.aligned.m64n104k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51}" ", {%52,%53,%54,%55}, %56, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N112K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %61, 0;\n" "wgmma.mma_async.sync.aligned.m64n112k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55}" ", {%56,%57,%58,%59}, %60, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N120K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %65, 0;\n" "wgmma.mma_async.sync.aligned.m64n120k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59}" ", {%60,%61,%62,%63}, %64, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N128K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %69, 0;\n" "wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63}" ", {%64,%65,%66,%67}, %68, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_ANY_T(NW,D,A,DB,TB) do { if constexpr (NW == 96) { WGMMA_RS_M64N96K16_T(D, A, DB, TB); } else if constexpr (NW == 104) { WGMMA_RS_M64N104K16_T(D, A, DB, TB); } else if constexpr (NW == 112) { WGMMA_RS_M64N112K16_T(D, A, DB, TB); } else if constexpr (NW == 120) { WGMMA_RS_M64N120K16_T(D, A, DB, TB); } else if constexpr (NW == 128) { WGMMA_RS_M64N128K16_T(D, A, DB, TB); } } while (0) +#define WGMMA_RS_M64N8K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %9, 0;\n" "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " "{" "%0,%1,%2,%3}" ", {%4,%5,%6,%7}, %8, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N16K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %13, 0;\n" "wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7}" ", {%8,%9,%10,%11}, %12, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N24K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %17, 0;\n" "wgmma.mma_async.sync.aligned.m64n24k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11}" ", {%12,%13,%14,%15}, %16, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N32K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %21, 0;\n" "wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15}" ", {%16,%17,%18,%19}, %20, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N40K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %25, 0;\n" "wgmma.mma_async.sync.aligned.m64n40k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N48K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %29, 0;\n" "wgmma.mma_async.sync.aligned.m64n48k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N56K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %33, 0;\n" "wgmma.mma_async.sync.aligned.m64n56k16.f32.f16.f16 " "{" "%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, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N64K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %37, 0;\n" "wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31}" ", {%32,%33,%34,%35}, %36, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N72K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %41, 0;\n" "wgmma.mma_async.sync.aligned.m64n72k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35}" ", {%36,%37,%38,%39}, %40, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N80K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %45, 0;\n" "wgmma.mma_async.sync.aligned.m64n80k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39}" ", {%40,%41,%42,%43}, %44, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N88K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %49, 0;\n" "wgmma.mma_async.sync.aligned.m64n88k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43}" ", {%44,%45,%46,%47}, %48, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N136K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %73, 0;\n" "wgmma.mma_async.sync.aligned.m64n136k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67}" ", {%68,%69,%70,%71}, %72, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N144K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %77, 0;\n" "wgmma.mma_async.sync.aligned.m64n144k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71}" ", {%72,%73,%74,%75}, %76, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N152K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %81, 0;\n" "wgmma.mma_async.sync.aligned.m64n152k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75}" ", {%76,%77,%78,%79}, %80, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N160K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %85, 0;\n" "wgmma.mma_async.sync.aligned.m64n160k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79}" ", {%80,%81,%82,%83}, %84, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N168K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %89, 0;\n" "wgmma.mma_async.sync.aligned.m64n168k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83}" ", {%84,%85,%86,%87}, %88, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N176K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %93, 0;\n" "wgmma.mma_async.sync.aligned.m64n176k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87}" ", {%88,%89,%90,%91}, %92, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N184K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %97, 0;\n" "wgmma.mma_async.sync.aligned.m64n184k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91}" ", {%92,%93,%94,%95}, %96, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N192K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %101, 0;\n" "wgmma.mma_async.sync.aligned.m64n192k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95}" ", {%96,%97,%98,%99}, %100, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N200K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %105, 0;\n" "wgmma.mma_async.sync.aligned.m64n200k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99}" ", {%100,%101,%102,%103}, %104, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N208K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %109, 0;\n" "wgmma.mma_async.sync.aligned.m64n208k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103}" ", {%104,%105,%106,%107}, %108, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N216K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %113, 0;\n" "wgmma.mma_async.sync.aligned.m64n216k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107}" ", {%108,%109,%110,%111}, %112, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N224K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %117, 0;\n" "wgmma.mma_async.sync.aligned.m64n224k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111}" ", {%112,%113,%114,%115}, %116, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N232K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %121, 0;\n" "wgmma.mma_async.sync.aligned.m64n232k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115}" ", {%116,%117,%118,%119}, %120, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N240K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %125, 0;\n" "wgmma.mma_async.sync.aligned.m64n240k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119}" ", {%120,%121,%122,%123}, %124, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N248K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %129, 0;\n" "wgmma.mma_async.sync.aligned.m64n248k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123}" ", {%124,%125,%126,%127}, %128, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#define WGMMA_RS_M64N256K16_T(D,A,DB,TB) asm volatile( "{\n .reg .pred p;\n setp.ne.b32 p, %133, 0;\n" "wgmma.mma_async.sync.aligned.m64n256k16.f32.f16.f16 " "{" "%0,%1,%2,%3,%4,%5,%6,%7," "%8,%9,%10,%11,%12,%13,%14,%15," "%16,%17,%18,%19,%20,%21,%22,%23," "%24,%25,%26,%27,%28,%29,%30,%31," "%32,%33,%34,%35,%36,%37,%38,%39," "%40,%41,%42,%43,%44,%45,%46,%47," "%48,%49,%50,%51,%52,%53,%54,%55," "%56,%57,%58,%59,%60,%61,%62,%63," "%64,%65,%66,%67,%68,%69,%70,%71," "%72,%73,%74,%75,%76,%77,%78,%79," "%80,%81,%82,%83,%84,%85,%86,%87," "%88,%89,%90,%91,%92,%93,%94,%95," "%96,%97,%98,%99,%100,%101,%102,%103," "%104,%105,%106,%107,%108,%109,%110,%111," "%112,%113,%114,%115,%116,%117,%118,%119," "%120,%121,%122,%123,%124,%125,%126,%127}" ", {%128,%129,%130,%131}, %132, p, 1, 1, " #TB ";\n}\n" : "+f"(D[0]), "+f"(D[1]), "+f"(D[2]), "+f"(D[3]), "+f"(D[4]), "+f"(D[5]), "+f"(D[6]), "+f"(D[7]), "+f"(D[8]), "+f"(D[9]), "+f"(D[10]), "+f"(D[11]), "+f"(D[12]), "+f"(D[13]), "+f"(D[14]), "+f"(D[15]), "+f"(D[16]), "+f"(D[17]), "+f"(D[18]), "+f"(D[19]), "+f"(D[20]), "+f"(D[21]), "+f"(D[22]), "+f"(D[23]), "+f"(D[24]), "+f"(D[25]), "+f"(D[26]), "+f"(D[27]), "+f"(D[28]), "+f"(D[29]), "+f"(D[30]), "+f"(D[31]), "+f"(D[32]), "+f"(D[33]), "+f"(D[34]), "+f"(D[35]), "+f"(D[36]), "+f"(D[37]), "+f"(D[38]), "+f"(D[39]), "+f"(D[40]), "+f"(D[41]), "+f"(D[42]), "+f"(D[43]), "+f"(D[44]), "+f"(D[45]), "+f"(D[46]), "+f"(D[47]), "+f"(D[48]), "+f"(D[49]), "+f"(D[50]), "+f"(D[51]), "+f"(D[52]), "+f"(D[53]), "+f"(D[54]), "+f"(D[55]), "+f"(D[56]), "+f"(D[57]), "+f"(D[58]), "+f"(D[59]), "+f"(D[60]), "+f"(D[61]), "+f"(D[62]), "+f"(D[63]), "+f"(D[64]), "+f"(D[65]), "+f"(D[66]), "+f"(D[67]), "+f"(D[68]), "+f"(D[69]), "+f"(D[70]), "+f"(D[71]), "+f"(D[72]), "+f"(D[73]), "+f"(D[74]), "+f"(D[75]), "+f"(D[76]), "+f"(D[77]), "+f"(D[78]), "+f"(D[79]), "+f"(D[80]), "+f"(D[81]), "+f"(D[82]), "+f"(D[83]), "+f"(D[84]), "+f"(D[85]), "+f"(D[86]), "+f"(D[87]), "+f"(D[88]), "+f"(D[89]), "+f"(D[90]), "+f"(D[91]), "+f"(D[92]), "+f"(D[93]), "+f"(D[94]), "+f"(D[95]), "+f"(D[96]), "+f"(D[97]), "+f"(D[98]), "+f"(D[99]), "+f"(D[100]), "+f"(D[101]), "+f"(D[102]), "+f"(D[103]), "+f"(D[104]), "+f"(D[105]), "+f"(D[106]), "+f"(D[107]), "+f"(D[108]), "+f"(D[109]), "+f"(D[110]), "+f"(D[111]), "+f"(D[112]), "+f"(D[113]), "+f"(D[114]), "+f"(D[115]), "+f"(D[116]), "+f"(D[117]), "+f"(D[118]), "+f"(D[119]), "+f"(D[120]), "+f"(D[121]), "+f"(D[122]), "+f"(D[123]), "+f"(D[124]), "+f"(D[125]), "+f"(D[126]), "+f"(D[127]) : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "l"(DB), "r"(1)) +#undef WGMMA_RS_ANY_T +#define WGMMA_RS_ANY_T(NW,D,A,DB,TB) do { if constexpr (NW == 8) { WGMMA_RS_M64N8K16_T(D, A, DB, TB); } else if constexpr (NW == 16) { WGMMA_RS_M64N16K16_T(D, A, DB, TB); } else if constexpr (NW == 24) { WGMMA_RS_M64N24K16_T(D, A, DB, TB); } else if constexpr (NW == 32) { WGMMA_RS_M64N32K16_T(D, A, DB, TB); } else if constexpr (NW == 40) { WGMMA_RS_M64N40K16_T(D, A, DB, TB); } else if constexpr (NW == 48) { WGMMA_RS_M64N48K16_T(D, A, DB, TB); } else if constexpr (NW == 56) { WGMMA_RS_M64N56K16_T(D, A, DB, TB); } else if constexpr (NW == 64) { WGMMA_RS_M64N64K16_T(D, A, DB, TB); } else if constexpr (NW == 72) { WGMMA_RS_M64N72K16_T(D, A, DB, TB); } else if constexpr (NW == 80) { WGMMA_RS_M64N80K16_T(D, A, DB, TB); } else if constexpr (NW == 88) { WGMMA_RS_M64N88K16_T(D, A, DB, TB); } else if constexpr (NW == 96) { WGMMA_RS_M64N96K16_T(D, A, DB, TB); } else if constexpr (NW == 104) { WGMMA_RS_M64N104K16_T(D, A, DB, TB); } else if constexpr (NW == 112) { WGMMA_RS_M64N112K16_T(D, A, DB, TB); } else if constexpr (NW == 120) { WGMMA_RS_M64N120K16_T(D, A, DB, TB); } else if constexpr (NW == 128) { WGMMA_RS_M64N128K16_T(D, A, DB, TB); } else if constexpr (NW == 136) { WGMMA_RS_M64N136K16_T(D, A, DB, TB); } else if constexpr (NW == 144) { WGMMA_RS_M64N144K16_T(D, A, DB, TB); } else if constexpr (NW == 152) { WGMMA_RS_M64N152K16_T(D, A, DB, TB); } else if constexpr (NW == 160) { WGMMA_RS_M64N160K16_T(D, A, DB, TB); } else if constexpr (NW == 168) { WGMMA_RS_M64N168K16_T(D, A, DB, TB); } else if constexpr (NW == 176) { WGMMA_RS_M64N176K16_T(D, A, DB, TB); } else if constexpr (NW == 184) { WGMMA_RS_M64N184K16_T(D, A, DB, TB); } else if constexpr (NW == 192) { WGMMA_RS_M64N192K16_T(D, A, DB, TB); } else if constexpr (NW == 200) { WGMMA_RS_M64N200K16_T(D, A, DB, TB); } else if constexpr (NW == 208) { WGMMA_RS_M64N208K16_T(D, A, DB, TB); } else if constexpr (NW == 216) { WGMMA_RS_M64N216K16_T(D, A, DB, TB); } else if constexpr (NW == 224) { WGMMA_RS_M64N224K16_T(D, A, DB, TB); } else if constexpr (NW == 232) { WGMMA_RS_M64N232K16_T(D, A, DB, TB); } else if constexpr (NW == 240) { WGMMA_RS_M64N240K16_T(D, A, DB, TB); } else if constexpr (NW == 248) { WGMMA_RS_M64N248K16_T(D, A, DB, TB); } else if constexpr (NW == 256) { WGMMA_RS_M64N256K16_T(D, A, DB, TB); } } while (0) diff --git a/deploy/nemotron_native/kernels/contraction/bf16/mk2_abi.cuh b/deploy/nemotron_native/kernels/contraction/bf16/mk2_abi.cuh new file mode 100644 index 0000000..0e995c0 --- /dev/null +++ b/deploy/nemotron_native/kernels/contraction/bf16/mk2_abi.cuh @@ -0,0 +1,248 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + +#pragma once +#include + +extern "C" { + +struct MK2Predictor; + + + + + + + + + + +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; + + + +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); + + + +MK2Predictor *mk2_create(const char *raw_dir, int batch); + + + + + + + + +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); + + + + + + + + + + + +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); + + + + +void mk2_destroy(MK2Predictor *p); + + + +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); + + + + + + +struct MK2Stage0Bindings { + + + + void *conv_state[14]; + int conv_slot_stride[14]; + + + void *ssm_state[14]; + + int ssm_slot_stride[14]; + + + + + const void *slot_map; + + + + + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + + + + + + const void *seq_len; + + + const void *embedding; + const void *gumbel; + float temperature; + + + + const void *moe_w_inject_debug; + + void *codes_out; + void *phoneme_out; +}; + + + + + + +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + + + + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + + + const void *at_qkv[3]; + const void *at_o[3]; + + + + + + + + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + + const void *phoneme_w; + const void *phoneme_b; + + + + + + + + + + + + const char *pattern; +}; + + + +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); + + + +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); + + + + + + + + +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); + + + + +int mk2_set_native_predictor_active_mask(MK2Predictor *p, const unsigned char *mask); +int mk2_native_startup_skip_abi_version(); + + + + + +unsigned long long mk2_stage0_launches(MK2Predictor *p); + + + + +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); + + + +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); + +} + +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *, int grid); diff --git a/deploy/nemotron_native/kernels/contraction/bf16/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/contraction/bf16/native_prefill_abi.cuh new file mode 100644 index 0000000..02a410b --- /dev/null +++ b/deploy/nemotron_native/kernels/contraction/bf16/native_prefill_abi.cuh @@ -0,0 +1,69 @@ +#pragma once +#include "mk2_abi.cuh" +#include + + + +extern "C" { +struct MK2NativePrefill; + +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; + +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; + + + + + + + + + +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); + + + + + + + + +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); + + + + + + + + + + +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); + + + + +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/contraction/bf16/native_prefill_impl.inc b/deploy/nemotron_native/kernels/contraction/bf16/native_prefill_impl.inc new file mode 100644 index 0000000..ece60e5 --- /dev/null +++ b/deploy/nemotron_native/kernels/contraction/bf16/native_prefill_impl.inc @@ -0,0 +1,216 @@ + +#include "native_prefill_abi.cuh" + +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; + +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; + +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 16 || BATCH == 32 || BATCH == 64)) + throw std::runtime_error("native prefill requires matching batch 1/16/32/64"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + + if (padded > INT_MAX / ((BB_EXPERTS + 1) * D_MODEL)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} + +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + float *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(float))); + + + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(float))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + + + + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} + +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} + +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} + +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} diff --git a/deploy/nemotron_native/kernels/contraction/bf16/predictor_mk2.cu b/deploy/nemotron_native/kernels/contraction/bf16/predictor_mk2.cu new file mode 100644 index 0000000..159d426 --- /dev/null +++ b/deploy/nemotron_native/kernels/contraction/bf16/predictor_mk2.cu @@ -0,0 +1,5642 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace kittens; + +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; + + + +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; + + + + +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif + + + + + + + + + + + + + + + + + + + + + +constexpr int BATCH_PAD = ((BATCH + 15) / 16) * 16; +static_assert(BATCH >= 1 && BATCH <= BATCH_PAD); + +static_assert(BATCH % 16 != 0 || BATCH_PAD == BATCH); +constexpr int CODEBOOK = 1032; + + +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; + + + + + +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; + +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; + +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; + +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; + + + +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; + + + + + + + + + + + + + + + + + + + + + + +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif + + + + +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif + + +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; + +constexpr int OUT_ATOM = 64; +#ifndef KCHUNK +#define KCHUNK 256 +#endif + + + + + +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif + + + + + + + + + + + + + + +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif + +constexpr int THREADS = 256; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; + +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; + + + + +using out_tile_f = st; + +using weight_gl = gl; +using act_gl = gl; + + + +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_ACC_F32 }; + + + +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + + + + + + + + + + + + + bf16 *conv_state; + int conv_slot_stride; + + + float *ssm_state; + int ssm_slot_stride; + const float *r8_initial_ssm = nullptr; + const unsigned int *r8_initial_ready = nullptr; + const int *r8_seq_len = nullptr; + int r8_initial_positions = 0; +}; + +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + + + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; + + + + + + + + + + + + + + + + + + + + + + +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; + + +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + + + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} + +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; + +#ifndef MK2_SPARSE_MOE +#define MK2_SPARSE_MOE 0 +#endif +constexpr int SPARSE_N = 16; +using sparse_act_tile = st_bf; +using sparse_out_tile = st_bf; +using sparse_act_gl = gl; +struct SparseMoe { + bf16 *gathered; + int *rows; + int *counts; + sparse_act_gl gathered_v; + sparse_act_gl hidden_v; +}; + +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + + + + + + const float *moe_w_inject; + + + + + + float *moe_w_audit; + + + + bf16 *dec_snapshot; + bf16 *xbc_act; + + + + const int *slot_map; + + + + + + + + int n_prefill_tokens; + int prefill_slot; + + + + + + + + + + + + + int pf_n_seq; + + + + + + + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + + + + + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + + + + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + + + + + + + + + + + + unsigned long long *launches; + + + + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + + int native_prefill_reset = 0; + const SparseMoe *sparse_moe = nullptr; +}; + + + + + + + + + + + + +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +#define BARRIER_GEN_WORD 32 +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + + + + + + + + + + + + + + auto *wide_state = reinterpret_cast(state); + const unsigned long long arrival = atomicAdd(wide_state, 1ull); + const unsigned int rank = (unsigned int)(arrival % gridDim.x); + if (rank != gridDim.x - 1u) { + const unsigned long long target = arrival - rank + gridDim.x; + volatile unsigned long long *counter = wide_state; + while (*counter < target) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} + +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} + + + + + + + + + + + + + +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / THREADS; + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const bf16 *in = src + (size_t)row * D_MODEL; + float v[PER_THREAD], wv[PER_THREAD]; + float sum = 0.f, sq = 0.f; + + + + + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + v[j] = __bfloat162float(in[threadIdx.x + j * THREADS]); + wv[j] = weight[threadIdx.x + j * THREADS]; + sum += v[j]; + sq += v[j] * v[j]; + } + + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) { scratch[warp] = sum; scratch[WARPS + warp] = sq; } + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f, q = 0.f; + for (int w = 0; w < WARPS; ++w) { s += scratch[w]; q += scratch[WARPS + w]; } + scratch[2 * WARPS] = s / D_MODEL; + scratch[2 * WARPS + 1] = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + } + __syncthreads(); + const float mean = scratch[2 * WARPS], inv = scratch[2 * WARPS + 1]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + + + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * wv[j]); + } + __syncthreads(); + } +} + + + + + + + + + + + + + +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, bf16 *snapshot = nullptr, int snapshot_layer = -1, + float *snapshot_fp32 = nullptr, float *out_fp32 = nullptr) { + constexpr int PER_THREAD = D_MODEL / THREADS; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const size_t base = (size_t)row * D_MODEL; + float v[PER_THREAD]; + bf16 wv[PER_THREAD]; + float sq = 0.f; + + + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + wv[j] = weight[i]; + const bf16 hv = h[base + i]; + if (snapshot) + snapshot[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = hv; + if (snapshot_fp32) + snapshot_fp32[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = + __bfloat162float(hv); + v[j] = __bfloat162float(hv) + __bfloat162float(r[base + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) scratch[warp] = sq; + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f; + for (int w = 0; w < WARPS; ++w) s += scratch[w]; + scratch[WARPS] = rsqrtf(s / D_MODEL + RMS_EPS); + } + __syncthreads(); + const float inv = scratch[WARPS]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + const bf16 normed = __float2bfloat16( + __bfloat162float(wv[j]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + out[base + i] = normed; + if (out_fp32) out_fp32[base + i] = __bfloat162float(normed); + } + __syncthreads(); + } +} + + + + + + + + + + + + + + + + + + + + + + + + + + + +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map) { + + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + const int c = idx - row * MB_CONV_DIM; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + const float raw = __bfloat162float( + zxbcdt[(size_t)row * MB_IN_PROJ + MB_INNER + c]); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + xbc_act[(size_t)row * MB_CONV_DIM + c] = __float2bfloat16(acc); + + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + } +} + + + +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} + +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map) { + + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + + + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + const int pairs = BATCH * MB_HEADS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_HEADS * split); + const int head = (pair / split) % MB_HEADS; + const int first_channel = (pair % split) * channels; + const int group = head / MB_HEADS_PER_GROUP; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + + + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + L.dt_bias[head]); + const float dA = __expf(L.A[head] * dt); + const float Dh = L.D[head]; + const size_t ssm_at = (size_t)slot_map[row] * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE; + bf16 *st = reinterpret_cast(L.ssm_state) + ssm_at; + const bool r8_initial = L.r8_initial_ssm && L.r8_initial_ready && *L.r8_initial_ready == 1 && L.r8_seq_len[row] == L.r8_initial_positions; + const float *st_in = r8_initial ? L.r8_initial_ssm + (size_t)head * MB_HEAD_DIM * MB_STATE : nullptr; + + + + + + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + + + + + + + + + + + + + + constexpr int NSTATE = MB_STATE / 32; + constexpr int SSM_PREFETCH = channels < 8 ? channels : 8; + static_assert(channels % SSM_PREFETCH == 0); + for (int g0 = 0; g0 < channels; g0 += SSM_PREFETCH) { + float sv[SSM_PREFETCH][NSTATE]; + + + + + + + float xvr[SSM_PREFETCH]; + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const size_t prow = (size_t)(first_channel + g0 + gi) * MB_STATE; + xvr[gi] = __bfloat162float( + xa[head * MB_HEAD_DIM + first_channel + g0 + gi]); + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = r8_initial ? st_in[prow + lane + j * 32] : __bfloat162float(st[prow + lane + j * 32]); + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const int p = first_channel + g0 + gi; + const float xv = xvr[gi]; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < NSTATE; ++j) { + const int n = lane + j * 32; + const size_t at = (size_t)p * MB_STATE + n; + const float s = sv[gi][j] * dA + Bv[j] * xv; + st[at] = __float2bfloat16(s); + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } + } +} + + + + + + + + + + +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_CONV_DIM; + const int c = idx - t * MB_CONV_DIM; + + + + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + + + + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v; +#ifdef PF_CONV_ZERO_STATE + + + + + if (src < 0) v = 0.f; +#else + + + + if (src < 0) v = __bfloat162float(state[pos + k]); +#endif + else v = __bfloat162float( + zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]); + acc += __bfloat162float(w[k]) * v; + } + + + +#ifndef PF_CONV_NO_SILU + acc = acc / (1.f + __expf(-acc)); +#endif + xbc_act[(size_t)t * MB_CONV_DIM + c] = __float2bfloat16(acc); + } +} + + + + + +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + + + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + + + + + + + + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; + + +#ifdef PF_CONV_STATE_NO_SHIFT + + + + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + + + + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + + + + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} + + + + + + + + + + +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + if (n_seq > 1) { + + + + + + + + + + + + + const int n_pos = n_tokens / n_seq; + const int units = n_seq * MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int seq = unit / (MB_HEADS * MB_HEAD_DIM); + const int rem = unit - seq * (MB_HEADS * MB_HEAD_DIM); + const int head = rem / MB_HEAD_DIM; + const int p = rem - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p * MB_STATE; + float sreg[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) sreg[j] = st[lane + j * 32]; + + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const bf16 *xa = xbc_act + t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + const bf16 *dtrow = zxbcdt + t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + sreg[j] = sreg[j] * dA + __bfloat162float(Bsec[n]) * dt * xv; + acc += sreg[j] * __bfloat162float(Csec[n]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) st[lane + j * 32] = sreg[j]; + } + return; + } +#if PF_SSM_SPLIT_P + + + + + + + + + + + const int units = MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; unit += warps_total) { + const int head = unit / MB_HEAD_DIM; + const int p_only = unit - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; +#else + for (int head = blockIdx.x * WARPS + warp; head < MB_HEADS; + head += warps_total) { + const int group = head / MB_HEADS_PER_GROUP; +#endif + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *st = L.ssm_state + (size_t)slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE +#if PF_SSM_SPLIT_P + + (size_t)p_only * MB_STATE +#endif + ; + + + + + for (int t = 0; t < n_tokens; ++t) { + const bf16 *xa = xbc_act + (size_t)t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = + zxbcdt + (size_t)t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } +#if PF_SSM_SPLIT_P + { + const int p = p_only; +#else + for (int p = 0; p < MB_HEAD_DIM; ++p) { +#endif + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { +#if PF_SSM_SPLIT_P + + const size_t at = (size_t)lane + j * 32; +#else + const size_t at = (size_t)p * MB_STATE + lane + j * 32; +#endif + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } + } +} + +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + + + + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32], wg[MB_GROUP_SIZE / 32]; + + + + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + wg[j] = __bfloat162float(L.gnorm_w[idx]); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16(held[j] * inv * wg[j]); + } + } +} + + + + + +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + constexpr int Q_PER_WARP = 2; + constexpr int GROUPS_PER_KV = 2; + const int pairs = BATCH * BB_KV_HEADS * GROUPS_PER_KV; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + + for (int pair = blockIdx.x + warp * gridDim.x; pair < pairs; + pair += gridDim.x * WARPS) { + const int row = pair / (BB_KV_HEADS * GROUPS_PER_KV); + const int within_row = pair % (BB_KV_HEADS * GROUPS_PER_KV); + const int kv = within_row / GROUPS_PER_KV; + const int q_group = within_row % GROUPS_PER_KV; + const int q_first = q_group * Q_PER_WARP; + const int q_count = q_group == 0 ? Q_PER_WARP : 1; + + + + + + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; + + +#if !BB_ATTN_SKIP_KV_WRITE + if (q_group == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + + + + + + + + + + + + + + + + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[h][j] = 0.f; + } + float m[Q_PER_WARP] = {-CUDART_INF_F, -CUDART_INF_F}; + float denom[Q_PER_WARP] = {0.f, 0.f}; + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + constexpr int PF = 8; + const int tpb = L.tokens_per_block; + const size_t elems_per_block = + (size_t)tpb * BB_KV_HEADS * BB_HEAD_DIM; + const int *bt_k = L.block_table + (size_t)(row * 2) * L.max_blocks_per_seq; + const int *bt_v = bt_k + L.max_blocks_per_seq; + const int head_base = kv * tpb * BB_HEAD_DIM; + float qr[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const bf16 *q = base + (kv * kv_per_q + q_first + h) * BB_HEAD_DIM; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qr[h][j] = __bfloat162float(q[lane + j * 32]); + } + + for (int page0 = 0; page0 < T; page0 += tpb) { + + + const int k_page = bt_k[page0 / tpb] & 0x7FFFFFFF; + const int v_page = bt_v[page0 / tpb] & 0x7FFFFFFF; + const bf16 *kbase = L.pool_base + + (size_t)k_page * elems_per_block + + head_base; + const bf16 *vbase = L.pool_base + + (size_t)v_page * elems_per_block + + head_base; + const int tend = page0 + tpb < T ? page0 + tpb : T; + for (int t0 = page0; t0 < tend; t0 += PF) { + + + + const int navail = tend - t0 < PF ? tend - t0 : PF; + float kr[PF][PER_LANE], vr[PF][PER_LANE]; + #pragma unroll + for (int u = 0; u < PF; ++u) { + const int tpos = t0 + (u < navail ? u : navail - 1); +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = vbase + (tpos - page0) * BB_HEAD_DIM; +#else + + + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new : kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = cur_token + ? v_new : vbase + (tpos - page0) * BB_HEAD_DIM; +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + kr[u][j] = __bfloat162float(kt[i]); + vr[u][j] = __bfloat162float(vt[i]); + } + } + #pragma unroll + for (int u = 0; u < PF; ++u) { + if (u >= navail) break; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + dot += qr[h][j] * kr[u][j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + + + + + const float m_new = fmaxf(m[h], s); + const float rescale = __expf(m[h] - m_new); + const float w = __expf(s - m_new); + denom[h] = denom[h] * rescale + w; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[h][j] = acc[h][j] * rescale + w * vr[u][j]; + m[h] = m_new; + } + } + } + } + + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const float inv = 1.f / denom[h]; + const int head = kv * kv_per_q + q_first + h; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[h][j] * inv); + } + } + } +} + + + + + +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * BB_KV_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / BB_KV_HEADS; + const int kv = idx - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; + + + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv); + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv); + for (int i = 0; i < BB_HEAD_DIM; ++i) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +} + +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + + + + + + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + + + + + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + for (int s = 0; s <= pos; ++s) { + const bf16 *kt = kv_slot(L, seq, 0, s) + kv_offset(L, s, kv); + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + const bf16 *vt = kv_slot(L, seq, 1, s) + kv_offset(L, s, kv); + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + wt * __bfloat162float(vt[i]); + } + m = m_new; + } + + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} + + + + + + +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH) { + __shared__ float score[BB_EXPERTS]; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + + + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + for (int e = warp; e < BB_EXPERTS; e += WARPS) { + const bf16 *wr = gate_w + (size_t)e * D_MODEL; + float acc = 0.f; + for (int i = lane; i < D_MODEL; i += 32) + acc += __bfloat162float(normed[(size_t)row * D_MODEL + i]) * + __bfloat162float(wr[i]); + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) score[e] = 1.f / (1.f + __expf(-acc)); + } + __syncthreads(); + if (threadIdx.x == 0) { + float remaining[BB_EXPERTS]; + int pick[BB_TOP_K]; + for (int e = 0; e < BB_EXPERTS; ++e) { + moe_w[row * BB_EXPERTS + e] = 0.f; + remaining[e] = score[e] + __bfloat162float(gate_bias[e]); + } + float total = 0.f; + for (int k = 0; k < BB_TOP_K; ++k) { + int best = 0; + for (int e = 1; e < BB_EXPERTS; ++e) + if (remaining[e] > remaining[best]) best = e; + pick[k] = best; + total += score[best]; + remaining[best] = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + for (int k = 0; k < BB_TOP_K; ++k) + moe_w[row * BB_EXPERTS + pick[k]] = score[pick[k]] * norm; + } + __syncthreads(); + } +} + + + +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert <= BB_EXPERTS; ++expert) + value = __fadd_rn(value, parts[base + (size_t)expert * D_MODEL]); + return __float2bfloat16(value); +} + +__device__ __forceinline__ bf16 bb_moe_fold_selected( + const float *parts, const float *weights, int idx) { + const int row = idx / D_MODEL; + const size_t base = (size_t)row * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert < BB_EXPERTS; ++expert) { + const float part = weights[row * BB_EXPERTS + expert] != 0.f + ? parts[base + (size_t)expert * D_MODEL] : 0.f; + value = __fadd_rn(value, part); + } + value = __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL]); + return __float2bfloat16(value); +} + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} + + + +__device__ __forceinline__ void sample_argmax_merge( + float &value, int &rank, float other, int other_rank) { + if (other > value || (other == value && other_rank < rank)) { + value = other; + rank = other_rank; + } +} + +__device__ __forceinline__ int sample_block_argmax(float value, int rank) { + __shared__ float warp_value[WARPS]; + __shared__ int warp_rank[WARPS]; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + if (lane == 0) { warp_value[warp] = value; warp_rank[warp] = rank; } + __syncthreads(); + if (warp == 0) { + value = lane < WARPS ? warp_value[lane] : -CUDART_INF_F; + rank = lane < WARPS ? warp_rank[lane] : INT_MAX; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + } + return rank; +} + +__device__ __forceinline__ int sample_eos_insertion_rank( + const uint64_t *keys, float eos) { + int lo = 0, hi = TOP_K; + while (lo < hi) { + const int mid = (lo + hi) >> 1; + if (eos > ordered_to_float((uint32_t)(keys[mid] >> 32))) hi = mid; + else lo = mid + 1; + } + return lo; +} + +__device__ __forceinline__ void sample_topk_serial( + const uint64_t *keys, float eos, const float *noise, int *out) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float value; + int code; + if (!eos_used && eos > av) { + value = eos; + code = AUDIO_EOS; + eos_used = true; + } else { + value = av; + code = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = value + noise[j]; + if (noised > best_noised) { + best_noised = noised; + best_code = code; + } + } + *out = best_code; +} + +__device__ __forceinline__ void sample_topk_parallel( + const uint64_t *keys, float eos, const float *noise, int *out) { + const int j = threadIdx.x; + const bool sorted_nan = j < TOP_K && + isnan(ordered_to_float((uint32_t)(keys[j] >> 32))); + if (__syncthreads_or(sorted_nan)) { + if (threadIdx.x == 0) sample_topk_serial(keys, eos, noise, out); + return; + } + const int lane = threadIdx.x & 31; + int insertion = lane == 0 ? sample_eos_insertion_rank(keys, eos) : 0; + insertion = __shfl_sync(0xffffffffu, insertion, 0); + float value = -CUDART_INF_F; + int rank = INT_MAX; + if (j < TOP_K) { + const int a = j - (j > insertion); + const float v = j == insertion + ? eos : ordered_to_float((uint32_t)(keys[a] >> 32)); + const float noised = v + noise[j]; + if (noised > -CUDART_INF_F) { value = noised; rank = j; } + } + const int winner = sample_block_argmax(value, rank); + if (threadIdx.x == 0) { + if (winner == INT_MAX) *out = 0; + else if (winner == insertion) *out = AUDIO_EOS; + else { + const int a = winner - (winner > insertion); + *out = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + } + } +} + +constexpr int SORT_ITEMS_PER_THREAD = SORT_N / THREADS; +struct sampler_storage { + uint64_t keys[2][SORT_N]; +}; + + + + + + + + + + + + + + + +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + sampler_storage &storage, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + + + + + + + + + + + uint64_t thread_keys[SORT_ITEMS_PER_THREAD]; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + + + thread_keys[item] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + + + + + + int cross_stage = 0; + for (int k = 2; k <= SORT_N; k <<= 1) { + for (int j = k >> 1; j > 0; j >>= 1) { + if (j >= 32) { + const int bank = cross_stage++ & 1; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + storage.keys[bank][n] = thread_keys[item]; + } + __syncthreads(); + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + const uint64_t other = storage.keys[bank][n ^ j]; + const bool take_hi = ((n & k) == 0) == ((n & j) == 0); + const uint64_t mine = thread_keys[item]; + thread_keys[item] = take_hi + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } else { + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + const uint64_t other = __shfl_xor_sync( + 0xffffffffu, thread_keys[item], j); + const bool take_hi = ((n & k) == 0) == ((n & j) == 0); + const uint64_t mine = thread_keys[item]; + thread_keys[item] = take_hi + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } + } + } + + const int top_bank = cross_stage & 1; + if (threadIdx.x < TOP_K) + storage.keys[top_bank][threadIdx.x] = thread_keys[0]; + __syncthreads(); + + + + + + sample_topk_parallel( + storage.keys[top_bank], eos_v, + gumbel + ((size_t)row * N_POS + pos) * TOP_K, + out_codes + row); + __syncthreads(); + } +} + + + + + + +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} + +__device__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES], act_tile (&a_s)[STAGES], out_tile &o_s, + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + int &tile_io, + const bf16 *bias = nullptr, + + + + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + + + + + + + + + + + + + int n_tok = 1, + + + + + + + + + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0) { + + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + + + + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + const int iters = k_total / KCHUNK; + + + + + + + + + + + + + + + + + + + + + + + + + + int tile_start = (int)blockIdx.x - tile_io % (int)gridDim.x; + if (tile_start < 0) tile_start += (int)gridDim.x; + tile_io += n_tiles; + int tiles_here = 0; + for (int t = tile_start; t < n_tiles; t += gridDim.x) ++tiles_here; + const int iter_base = iter_io; + + + + + + + iter_io += tiles_here * iters * n_tok; + + + + + + + const int my_tiles = (n_tiles - tile_start + gridDim.x - 1) / gridDim.x; + + if (warpgroup::groupid() == 1) { + + if (warpgroup::laneid() == 0) { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) { + const int tile_id = mk2_expert_tile(work, active_mask, compact_tiles); + + for (int tok = 0; tok < n_tok; ++tok) + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + + + + + wait(drained[s], phase); + tma::expect_bytes(filled[s], + sizeof(weight_tile) + sizeof(act_tile)); + + + + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s], w, {tile_id, i}, filled[s]); + const int a_col = out_tiles_per_expert + ? (tile_id / out_tiles_per_expert) * act_k_tiles_per_expert + i + : i; + tma::load_async<2, cache_policy::EVICT_LAST>( + a_s[s], a, {tok, a_col}, filled[s]); + } + } + } + } else { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) + for (int tok = 0; tok < n_tok; ++tok) { + const int tile_id = mk2_expert_tile(work, active_mask, compact_tiles); + + + + + + + + + + + rt_fl<16, BATCH_PAD> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s], a_s[s]); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s]); + } + + + + + + const int base = tile_id * OUT_ATOM; + + + + const int grow0 = tok * BATCH_PAD; + if (epi == EPI_BIAS) { + warpgroup::store(o_sf, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + dest[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(o_sf[int2{col, row}] + + __bfloat162float(bias[base + col])); + } + warpgroup::sync(1); + continue; + } + warpgroup::store(o_s, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + + + + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_s[int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + + + + + + + + + + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_w[grow * BB_EXPERTS + expert]; + + + + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * + D_MODEL + local] = __fmul_rn(wt, v); + break; + } + case EPI_ACC_F32: + + + + + + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + + warpgroup::sync(1); + } + (void)my_tiles; + } +} + +#include "sparse_moe_impl.cuh" + + +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr, + const unsigned char *native_predictor_active_mask = nullptr) { + + + + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + + + + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&w_s)[STAGES] = alloc.allocate(); + act_tile (&a_s)[STAGES] = alloc.allocate(); + out_tile (&o_s) = alloc.allocate(); + out_tile_f (&o_sf) = alloc.allocate(); + + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * WARPS + 2]; + + + __shared__ sampler_storage sort_storage; + + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + + + + init_semaphore(filled[s], 0, 1); + + + + init_semaphore(drained[s], CONSUMER_WARPS, 0); + for (int c = 0; c < CONSUMER_WARPS; ++c) arrive(drained[s]); + } + } + + + int gemm_iter = 0; + int gemm_tile = 0; + __syncthreads(); + + + + + + + + + + + + + + + + + + + + + + + + + + + + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch); + grid_barrier(g.barrier); + + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map); + grid_barrier(g.barrier); + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map); + grid_barrier(g.barrier); + mamba_gated_norm(L, bb->zxbcdt, bb->gated); + grid_barrier(g.barrier); + + + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len); + grid_barrier(g.barrier); + + + + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w); + } + + + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + } + grid_barrier(g.barrier); + + if constexpr (MK2_SPARSE_MOE && BATCH >= 16) { + if (bb->sparse_moe) { + sparse_moe_pack(bb); + grid_barrier(g.barrier); + } + } + + if (bb->sparse_moe) { + sparse_moe_gemm(L.exp_up, bb->sparse_moe->gathered_v, + *bb->sparse_moe, bb->moe_h, bb->moe_w, bb->moe_out, + w_s, a_s, o_s, filled, drained, gemm_iter, gemm_tile); + } else { + + gemm(L.exp_up, bb->normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, nullptr, 1, BATCH, + bb->moe_w, BB_MOE_INTER / OUT_ATOM); + } + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + + + + + + if (bb->sparse_moe) { + sparse_moe_gemm(L.exp_down, bb->sparse_moe->hidden_v, + *bb->sparse_moe, bb->moe_h, bb->moe_w, bb->moe_out, + w_s, a_s, o_s, filled, drained, gemm_iter, gemm_tile); + } else { + gemm(L.exp_down, bb->moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + 1, BATCH, bb->moe_w, D_MODEL / OUT_ATOM); + } + + + + + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_selected(bb->moe_out, bb->moe_w, idx); + grid_barrier(g.barrier); + } + + + + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + + + + + + + + + + + + if (bb && bb->n_prefill_tokens > 0) { + const int n_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + + + + + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = n_total / BATCH; + grid_barrier(g.barrier); + } + + + + + + + + + + + + + + + + + + + + + + + + + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + bb->pf_h[idx] = bb->prompt[idx]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + + + + + + + + + + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = bb->pf_n_seq > 0 ? bb->pf_n_seq : 1; + + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T); + grid_barrier(g.barrier); + + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_gated, + T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + + + + + + + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + + + + + + + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + + + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, T); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * (BB_EXPERTS + 1) * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_moe_out[idx] = 0.f; + grid_barrier(g.barrier); + + gemm(L.exp_up, bb->pf_normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + + gemm(L.exp_down, bb->pf_moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = bb_moe_fold_deterministic(bb->pf_moe_out, idx); + grid_barrier(g.barrier); + } + + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + + + + + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, bb->pf_final, ln_scratch, + n_total); + grid_barrier(g.barrier); + + if (!bb->pf_fuse_decode) return; + + + + + + + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + + + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + + } + + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch); + grid_barrier(g.barrier); + + + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, bb->phoneme_b); + grid_barrier(g.barrier); + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + if (threadIdx.x == 0) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + float best = -CUDART_INF_F; + int arg = 0; + + + for (int n = 0; n < BB_PHONEME; ++n) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + bb->phoneme_out[row] = arg; + } + __syncthreads(); + } + grid_barrier(g.barrier); + } + + if (bb->dec_snapshot) { + + + + + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + + + + + + + __shared__ int native_run_predictor; + if (native_predictor_active_mask && bb && g.native_codes) { + if (threadIdx.x == 0) { + int any_active = 0; + for (int row = 0; row < BATCH; ++row) + any_active |= int(native_predictor_active_mask[row] != 0); + native_run_predictor = any_active; + } + __syncthreads(); + } + const bool skip_native_predictor = native_predictor_active_mask && bb && + g.native_codes && native_run_predictor == 0; + if (skip_native_predictor) { + + + + for (int index = blockIdx.x * THREADS + threadIdx.x; + index < BATCH * N_POS; index += gridDim.x * THREADS) { + g.sampled[index] = 0; + g.native_codes[index] = 0; + } + grid_barrier(g.barrier); + } else { + + + + + + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + + + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + + for (int pos = 0; pos < n_pos; ++pos) { + + for (int layer = 0; layer < n_layers; ++layer) { + +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + + + + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + + + + + + + + + + + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + + + + + + + + constexpr int DPL = D_HEAD / 32; + float qv[DPL], kn[DPL], vn[DPL]; + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + qv[j] = __bfloat162float(q[i]); + kn[j] = __bfloat162float(k_new[i]); + vn[j] = __bfloat162float(v_new[i]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + + + + + + + float sc[N_POS]; + + + + + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += qv[j] * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + { + float acc = 0.f; + #pragma unroll + for (int j = 0; j < DPL; ++j) acc += qv[j] * kn[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[pos] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + + + + + + + + float ov[DPL] = {}; + + + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *vt = g.vcache + cache_base + t * D_HEAD; + const float s = sc[t]; + #pragma unroll + for (int j = 0; j < DPL; ++j) + ov[j] += s * __bfloat162float(vt[lane + j * 32]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + ov[j] += sc[pos] * vn[j]; + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + + lane + j * 32] = __float2bfloat16(ov[j]); + } + } + } +#endif + grid_barrier(g.barrier); + + + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + + +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + + + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + + + + + + + + + + + + + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } + + + + + +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, + g.sampled_pos, sort_storage, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif + + + + + +#ifdef ABLATE_SAMPLE + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { + __syncthreads(); + } else { + grid_barrier(g.barrier); + } +#endif + + + + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + + + if (threadIdx.x == 0) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + g.sampled[row * N_POS + pos] = code; + if (g.native_codes) + g.native_codes[row * N_POS + pos] = (long long)code; + } + } + + + + +#ifdef ABLATE_SAMPLE + grid_barrier(g.barrier); +#else + if (pos + 1 == n_pos) grid_barrier(g.barrier); +#endif + } + + } + + + + if (g.native_seq_len && blockIdx.x == 0 && threadIdx.x < BATCH) + g.native_seq_len[threadIdx.x] += 1; + + + + + + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +} + + + +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) + + + + + + + + +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} + + +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} + +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} + +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} + +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + + + std::string prefill_dir; + + + + + + int prefill_bench = 0; + + + + + int prefill_seq = 1; + + + + int prefill_fuse = 0; + + + + + + int bench_iters = 200; + + + + + + + + + + + + + + + + + + + + + + + + + + + bool poison_cur_kv = false; + int decode_steps = 1; + + + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + + + + + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + + + + + + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 64 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + + + + + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + + + + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + + + + + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + + + + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + + + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + + + + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + + + + + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + + + + + + + const int TPB = 32; + + + + + + + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + + + + + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + + + + + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + + + + + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + + + + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + + + + + + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + + + + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + + + + + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + + + + + + + + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + + + + + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + + + + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + + const int smem = STAGES * (sizeof(weight_tile) + sizeof(act_tile)) + + sizeof(out_tile) + sizeof(out_tile_f) + 4096; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + + + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + + + + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + + + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + + + + + + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + + + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + + + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + + + + + + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + + + + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + + + + + + + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + + + + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + + + + + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + + + + + + + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + + + + + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + + + + + + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + + + + + + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + + + + + + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + + + + + + + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + + + + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + + + + + + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + + + + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + + + + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + + + + + + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + + + + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + + + + + + + + + + + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + + + + + + + + + + + + + + + + + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + + + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + + + + + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + + + + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + + + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + + + + + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + + + + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + + + + + + + + + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + + + + + int self_mismatch = 0; + + + + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + + + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + + + + + + + + + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + + + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + + + + + + + + + + + + + + + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} + +#endif + + + + + + + + + + +#include "mk2_abi.cuh" + +namespace { +thread_local std::string g_last_error; + + + +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + + + + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; + const unsigned char *native_predictor_active_mask = nullptr; + int maximum_grid = 0; + bool grid_reset_pending = false; +}; + +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} + +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * (sizeof(weight_tile) + sizeof(act_tile)) + + sizeof(out_tile) + sizeof(out_tile_f) + 4096), 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = std::min(out->sms * out->blocks_per_sm, 4096); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} + +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + + + + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 64 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + + + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->maximum_grid = info.sms * info.blocks_per_sm; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} + + + +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} + +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + + + + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + + + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} + + + + + +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} + + + + + + + + + + + + + +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + + + + std::vector mb_host; + std::vector at_host; + + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + + + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; + +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; + + + +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; + +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} + +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} + + +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} + +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + + + + + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + + + + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + + + + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + + + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + + + + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + + + + + + + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + + + + + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + + + + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + + if constexpr (MK2_SPARSE_MOE && BATCH >= 16) { + + + bf16 *gathered = nullptr; + int *rows = nullptr, *counts = nullptr; + SparseMoe *device_sparse = nullptr; + mk2_check(owned.alloc(&gathered, + (size_t)BB_EXPERTS * BATCH_PAD * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&rows, (size_t)BB_EXPERTS * BATCH_PAD * sizeof(int))); + mk2_check(owned.alloc(&counts, BB_EXPERTS * sizeof(int))); + SparseMoe sparse{gathered, rows, counts, + sparse_act_gl{gathered, nullptr, nullptr, (size_t)BB_EXPERTS * BATCH_PAD, + (size_t)D_MODEL}, + sparse_act_gl{moe_h, nullptr, nullptr, (size_t)BB_EXPERTS * BATCH_PAD, + (size_t)BB_MOE_INTER}}; + mk2_check(owned.alloc(&device_sparse, sizeof(SparseMoe))); + mk2_check(cudaMemcpy(device_sparse, &sparse, sizeof(sparse), cudaMemcpyHostToDevice)); + hb.sparse_moe = device_sparse; + } + + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + + + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.maximum_grid = pred->maximum_grid; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} + +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + + + + + + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + + + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + + + + + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + + + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + + + + + + + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(s, &capture_status)); + if (st->pred.grid_reset_pending || capture_status != cudaStreamCaptureStatusNone) { + mk2_check(cudaMemsetAsync(st->pred.g.barrier, 0, 64 * sizeof(unsigned int), s)); + st->pred.grid_reset_pending = false; + } + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS, nullptr, + native ? st->pred.native_predictor_active_mask : nullptr); + + + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} + +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, false); +} + +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, true); +} + + + + +extern "C" int mk2_set_native_predictor_active_mask(MK2Predictor *p, + const unsigned char *mask) { + try { + if (!p || reinterpret_cast(p)->kind != 1) + throw std::runtime_error("startup mask requires a full Stage0 handle"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && st->pred.native_predictor_active_mask != mask) + throw std::runtime_error("startup mask pointer cannot change after first Stage0 enqueue"); + if (mask) + mk2_stage0_device_pointer(mask, st->pred.device, "native_predictor_active_mask"); + st->pred.native_predictor_active_mask = mask; + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown startup mask binding failure"; return 1; } +} + +extern "C" int mk2_native_startup_skip_abi_version() { return 1; } + + + + + +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + + + return (n[1] << 32) | (n[0] & 0xffffffffull); +} + + + + + + + +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} + +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } + + +#include "native_prefill_impl.inc" + + + + + + +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *p, int grid) { + if (!p) { g_last_error = "grid setter requires full Stage0 handle"; return 1; } + auto *st = reinterpret_cast(p); + if (st->kind != 1 || grid < 1 || grid > st->maximum_grid) { + g_last_error = "requested Stage0 grid outside independently resident capacity"; + return 1; + } + if (st->grid != grid) st->grid_reset_pending = true; + st->grid = grid; + g_last_error.clear(); + return 0; +} + +extern "C" int mk2_r8_bind_initial_state(MK2Predictor *p, const void *const *initial, const void *ready, const void *seq_len, int positions) { + try { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || st->bindings_initialized || !initial || !ready || !seq_len || positions <= 0) + throw std::runtime_error("invalid immutable initial-state binding"); + for (int i = 0; i < 14; ++i) { + mk2_stage0_device_pointer(initial[i], st->pred.device, "initial state"); + st->mb_host[i].r8_initial_ssm = static_cast(initial[i]); + st->mb_host[i].r8_initial_ready = static_cast(ready); + st->mb_host[i].r8_seq_len = static_cast(seq_len); + st->mb_host[i].r8_initial_positions = positions; + } + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} + +extern "C" int mk2_b32_private_state_bytes() { return 2; } diff --git a/deploy/nemotron_native/kernels/contraction/bf16/sparse_moe_impl.cuh b/deploy/nemotron_native/kernels/contraction/bf16/sparse_moe_impl.cuh new file mode 100644 index 0000000..6ccf766 --- /dev/null +++ b/deploy/nemotron_native/kernels/contraction/bf16/sparse_moe_impl.cuh @@ -0,0 +1,156 @@ + + + + +__device__ void sparse_moe_pack(const Backbone *bb) { + const SparseMoe &s = *bb->sparse_moe; + __shared__ int count; + __shared__ int local_rows[BATCH]; + + + + const int experts_per_wave = min(int(gridDim.x), BB_EXPERTS); + const int first_expert = int(blockIdx.x) % experts_per_wave; + const int rank = int(blockIdx.x) / experts_per_wave; + const int team_size = (int(gridDim.x) + experts_per_wave - 1 - first_expert) / + experts_per_wave; + for (int e = first_expert; e < BB_EXPERTS; e += experts_per_wave) { + if (threadIdx.x == 0) { + int n = 0; + for (int row = 0; row < BATCH; ++row) + if (bb->moe_w[row * BB_EXPERTS + e] != 0.f) + local_rows[n++] = row; + count = n; + if (rank == 0) { + s.counts[e] = n; + for (int i = 0; i < n; ++i) + s.rows[e * BATCH_PAD + i] = local_rows[i]; + } + } + __syncthreads(); + constexpr int chunks = D_MODEL / 8; + const int padded = ((count + SPARSE_N - 1) / SPARSE_N) * SPARSE_N; + + + for (int i = rank * THREADS + threadIdx.x; i < padded * chunks; + i += team_size * THREADS) { + const int row = i / chunks, col = i % chunks; + const int packed = e * BATCH_PAD + row; + auto *dst = reinterpret_cast(s.gathered + + (size_t)packed * D_MODEL + col * 8); + if (row < count) { + const int original = local_rows[row]; + *dst = *reinterpret_cast(bb->normed + + (size_t)original * D_MODEL + col * 8); + } else { + *dst = make_float4(0.f, 0.f, 0.f, 0.f); + } + } + __syncthreads(); + } +} + +template +__device__ void sparse_moe_gemm( + const weight_gl &w, const sparse_act_gl &a, const SparseMoe &s, + bf16 *hidden, const float *routes, float *contributions, + weight_tile (&w_s)[STAGES], act_tile (&a_s)[STAGES], out_tile &o_s, + semaphore *filled, semaphore *drained, int &iter_io, int &tile_io) { + constexpr int width = DOWN ? D_MODEL : BB_MOE_INTER; + constexpr int k_total = DOWN ? BB_MOE_INTER : D_MODEL; + constexpr int out_tiles = width / OUT_ATOM; + constexpr int iters = k_total / KCHUNK; + static_assert(BATCH_PAD >= SPARSE_N && BATCH_PAD % SPARSE_N == 0); + static_assert(sizeof(sparse_act_tile) <= sizeof(act_tile)); + static_assert(sizeof(sparse_out_tile) <= sizeof(out_tile)); + __shared__ int prefix[BB_EXPERTS + 1]; + __shared__ int row_tiles[BB_EXPERTS]; + if (threadIdx.x == 0) { + int total = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + prefix[e] = total; + row_tiles[e] = (s.counts[e] + SPARSE_N - 1) / SPARSE_N; + total += row_tiles[e] * out_tiles; + } + prefix[BB_EXPERTS] = total; + } + __syncthreads(); + const int n_tiles = prefix[BB_EXPERTS]; + + + int start = int(blockIdx.x) - tile_io % int(gridDim.x); + if (start < 0) start += gridDim.x; + tile_io += n_tiles; + const int mine = start < n_tiles ? 1 + (n_tiles - 1 - start) / gridDim.x : 0; + const int iter_base = iter_io; + iter_io += mine * iters; + + if (warpgroup::groupid() == 1) { + if (warpgroup::laneid() == 0) { + int iter = iter_base; + for (int f = start; f < n_tiles; f += gridDim.x) { + int e = 0; + while (prefix[e + 1] <= f) ++e; + const int rem = f - prefix[e], nt = row_tiles[e]; + const int feature_tile = rem / nt; + const int token_tile = rem - feature_tile * nt; + const int weight_row = e * out_tiles + feature_tile; + const int act_row = e * (BATCH_PAD / SPARSE_N) + token_tile; + for (int k = 0; k < iters; ++k, ++iter) { + const int stage = iter % STAGES, phase = (iter / STAGES) & 1; + wait(drained[stage], phase); + tma::expect_bytes(filled[stage], sizeof(weight_tile) + sizeof(sparse_act_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[stage], w, {weight_row, k}, filled[stage]); + + + tma::load_async<2, cache_policy::EVICT_LAST>( + *reinterpret_cast(&a_s[stage]), + a, {act_row, k}, filled[stage]); + } + } + } + } else if (warpgroup::groupid() == 0) { + int iter = iter_base; + for (int f = start; f < n_tiles; f += gridDim.x) { + int e = 0; + while (prefix[e + 1] <= f) ++e; + const int rem = f - prefix[e], nt = row_tiles[e]; + const int feature_tile = rem / nt; + const int token_tile = rem - feature_tile * nt; + rt_fl<16, SPARSE_N> accum; + warp::zero(accum); + for (int k = 0; k < iters; ++k, ++iter) { + const int stage = iter % STAGES, phase = (iter / STAGES) & 1; + wait(filled[stage], phase); + warpgroup::mma_ABt(accum, w_s[stage], + *reinterpret_cast(&a_s[stage])); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[stage]); + } + sparse_out_tile &out = *reinterpret_cast(&o_s); + + warpgroup::store(out, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int i = warpgroup::laneid(); i < OUT_ATOM * SPARSE_N; i += 128) { + const int row = i / OUT_ATOM, col = i % OUT_ATOM; + const int local = token_tile * SPARSE_N + row; + const int packed = e * BATCH_PAD + local; + const int output = feature_tile * OUT_ATOM + col; + const float value = float(out[{col, row}]); + if constexpr (!DOWN) { + + hidden[(size_t)packed * BB_MOE_INTER + output] = + bf16(value / (1.f + __expf(-value))); + } else if (local < s.counts[e]) { + const int original = s.rows[packed]; + const float weight = routes[original * BB_EXPERTS + e]; + contributions[((size_t)original * (BB_EXPERTS + 1) + e) * + D_MODEL + output] = __fmul_rn(weight, value); + } + } + warpgroup::sync(1); + } + } +} diff --git a/deploy/nemotron_native/kernels/contraction/lane_transfer.cu b/deploy/nemotron_native/kernels/contraction/lane_transfer.cu new file mode 100644 index 0000000..4d14925 --- /dev/null +++ b/deploy/nemotron_native/kernels/contraction/lane_transfer.cu @@ -0,0 +1,67 @@ + + +#include +#include "mk2_abi.cuh" +#include + +struct LaneArgs { + const uint16_t *embedding; uint16_t *packed_embedding; + const uint32_t *noise; uint32_t *packed_noise; + const int *rows, *slot; int *packed_slot; + int *seq, *packed_seq; + const int *tables[3]; int *packed_tables[3]; + int64_t *codes, *packed_codes; int *phones, *packed_phones; + int full_width, narrow_width, blocks; +}; +static LaneArgs arguments(const MK2Stage0Bindings &full, const MK2Stage0Bindings &small, + const int *rows, int full_width, int narrow_width) { + LaneArgs a{}; + a.embedding=(const uint16_t*)full.embedding; a.packed_embedding=(uint16_t*)small.embedding; + a.noise=(const uint32_t*)full.gumbel; a.packed_noise=(uint32_t*)small.gumbel; + a.rows=rows; a.slot=(const int*)full.slot_map; a.packed_slot=(int*)small.slot_map; + a.seq=(int*)full.seq_len; a.packed_seq=(int*)small.seq_len; + for(int l=0;l<3;l++){a.tables[l]=(const int*)full.kv_block_table[l];a.packed_tables[l]=(int*)small.kv_block_table[l];} + a.codes=(int64_t*)full.codes_out; a.packed_codes=(int64_t*)small.codes_out; + a.phones=(int*)full.phoneme_out;a.packed_phones=(int*)small.phoneme_out; + a.full_width=full_width;a.narrow_width=narrow_width;a.blocks=full.kv_max_blocks_per_seq; + return a; +} +__global__ void gather_lanes(LaneArgs a) { + int r=blockIdx.x, original=a.rows[r]; + for(int j=threadIdx.x;j<1536;j+=blockDim.x) + a.packed_embedding[r*1536+j]=a.embedding[original*1536+j]; + for(int j=threadIdx.x;j<16*80;j+=blockDim.x) + a.packed_noise[r*16*80+j]=a.noise[original*16*80+j]; + for(int l=0;l<3;l++) for(int j=threadIdx.x;j<2*a.blocks;j+=blockDim.x) + a.packed_tables[l][r*2*a.blocks+j]=a.tables[l][original*2*a.blocks+j]; + if(threadIdx.x==0){a.packed_slot[r]=a.slot[original];a.packed_seq[r]=a.seq[original];} +} +__global__ void scatter_lanes(LaneArgs a) { + int original=blockIdx.x, packed=-1; + for(int r=0;r=0?a.packed_codes[packed*16+threadIdx.x]:0; + if(threadIdx.x==0){ + a.phones[original]=packed>=0?a.packed_phones[packed]:0; + + + if(packed>=0) a.seq[original]=a.packed_seq[packed]; + } +} +static bool valid(const MK2Stage0Bindings *f,const MK2Stage0Bindings *s,const int *rows,int fw,int nw){ + return f&&s&&rows&&((fw==32&&nw==16)||(fw==64&&(nw==16||nw==32)))&& + f->kv_max_blocks_per_seq==s->kv_max_blocks_per_seq&&f->kv_max_blocks_per_seq>0&& + f->kv_tokens_per_block==32&&s->kv_tokens_per_block==32; +} +extern "C" size_t native_lanes_bindings_size(){return sizeof(MK2Stage0Bindings);} +extern "C" int native_lanes_gather(const MK2Stage0Bindings *f,const MK2Stage0Bindings *s, + const int *rows,int fw,int nw,void *stream){ + if(!valid(f,s,rows,fw,nw))return (int)cudaErrorInvalidValue; + gather_lanes<<>>(arguments(*f,*s,rows,fw,nw)); + return (int)cudaGetLastError(); +} +extern "C" int native_lanes_scatter(const MK2Stage0Bindings *f,const MK2Stage0Bindings *s, + const int *rows,int fw,int nw,void *stream){ + if(!valid(f,s,rows,fw,nw))return (int)cudaErrorInvalidValue; + scatter_lanes<<>>(arguments(*f,*s,rows,fw,nw)); + return (int)cudaGetLastError(); +} diff --git a/deploy/nemotron_native/kernels/contraction/mk2_abi.cuh b/deploy/nemotron_native/kernels/contraction/mk2_abi.cuh new file mode 100644 index 0000000..0e995c0 --- /dev/null +++ b/deploy/nemotron_native/kernels/contraction/mk2_abi.cuh @@ -0,0 +1,248 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + +#pragma once +#include + +extern "C" { + +struct MK2Predictor; + + + + + + + + + + +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; + + + +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); + + + +MK2Predictor *mk2_create(const char *raw_dir, int batch); + + + + + + + + +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); + + + + + + + + + + + +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); + + + + +void mk2_destroy(MK2Predictor *p); + + + +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); + + + + + + +struct MK2Stage0Bindings { + + + + void *conv_state[14]; + int conv_slot_stride[14]; + + + void *ssm_state[14]; + + int ssm_slot_stride[14]; + + + + + const void *slot_map; + + + + + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + + + + + + const void *seq_len; + + + const void *embedding; + const void *gumbel; + float temperature; + + + + const void *moe_w_inject_debug; + + void *codes_out; + void *phoneme_out; +}; + + + + + + +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + + + + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + + + const void *at_qkv[3]; + const void *at_o[3]; + + + + + + + + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + + const void *phoneme_w; + const void *phoneme_b; + + + + + + + + + + + + const char *pattern; +}; + + + +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); + + + +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); + + + + + + + + +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); + + + + +int mk2_set_native_predictor_active_mask(MK2Predictor *p, const unsigned char *mask); +int mk2_native_startup_skip_abi_version(); + + + + + +unsigned long long mk2_stage0_launches(MK2Predictor *p); + + + + +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); + + + +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); + +} + +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *, int grid); diff --git a/deploy/nemotron_native/kernels/inflight/progress.cuh b/deploy/nemotron_native/kernels/inflight/progress.cuh new file mode 100644 index 0000000..a1bd113 --- /dev/null +++ b/deploy/nemotron_native/kernels/inflight/progress.cuh @@ -0,0 +1,17 @@ +#pragma once +#ifndef MK2_INFLIGHT +#define MK2_INFLIGHT 0 +#endif +#if MK2_INFLIGHT +#define MK2_PROGRESS_PARAMETER(name) , const unsigned char *name = nullptr +#define MK2_PROGRESS_ARGUMENT(value) , value +#define MK2_PROGRESS_GUARD(mask, row) if ((mask) && !(mask)[row]) continue +#define MK2_PROGRESS_ENABLED(mask, row) (!(mask) || (mask)[row] != 0) +#define MK2_PROGRESS_MEMBER const unsigned char *native_progress_mask = nullptr; +#else +#define MK2_PROGRESS_PARAMETER(name) +#define MK2_PROGRESS_ARGUMENT(value) +#define MK2_PROGRESS_GUARD(mask, row) +#define MK2_PROGRESS_ENABLED(mask, row) true +#define MK2_PROGRESS_MEMBER +#endif diff --git a/deploy/nemotron_native/kernels/inflight/runtime.cuh b/deploy/nemotron_native/kernels/inflight/runtime.cuh new file mode 100644 index 0000000..01faa01 --- /dev/null +++ b/deploy/nemotron_native/kernels/inflight/runtime.cuh @@ -0,0 +1,141 @@ +#pragma once +#if MK2_INFLIGHT +extern "C" int mk2_native_inflight_abi_version() { return 1; } +extern "C" int mk2_native_inflight_batch() { return BATCH; } +extern "C" int mk2_set_native_progress_mask(MK2Predictor *p, + const unsigned char *mask) { + try { + if (!p || reinterpret_cast(p)->kind != 1 || !mask) + throw std::runtime_error("progress mask requires a full Stage0 handle and mask"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && st->pred.native_progress_mask != mask) + throw std::runtime_error("progress mask pointer cannot change after first Stage0 enqueue"); + mk2_stage0_device_pointer(mask, st->pred.device, "native_progress_mask"); + st->pred.native_progress_mask = mask; + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown progress mask binding failure"; return 1; } +} +__global__ void native_inflight_noise( + const unsigned long long *seeds, const int *generations, + const unsigned char *mask, float *out) { + constexpr int groups = N_POS * TOP_K / 4; + const int index = blockIdx.x * blockDim.x + threadIdx.x; + if (index >= BATCH * groups) return; + const int row = index / groups; + if (!mask[row]) return; + const int group = index % groups; + curandStatePhilox4_32_10_t state; + curand_init(seeds[row], (unsigned long long)group, + (unsigned long long)(unsigned int)generations[row] * 4ull, &state); + const float4 uniform = curand_uniform4(&state); + const float maximum = 0x1.fffffep-1f; + const float4 noise = make_float4( + -logf(-logf(fminf(uniform.x, maximum))), + -logf(-logf(fminf(uniform.y, maximum))), + -logf(-logf(fminf(uniform.z, maximum))), + -logf(-logf(fminf(uniform.w, maximum)))); + reinterpret_cast(out)[index] = noise; +} +extern "C" int mk2_inflight_noise( + const unsigned long long *seeds, const int *generations, + const unsigned char *mask, float *out, void *stream) { + try { + int device = -1; + mk2_check(cudaGetDevice(&device)); + mk2_stage0_device_pointer(seeds, device, "request_seeds"); + mk2_stage0_device_pointer(generations, device, "request_generations"); + mk2_stage0_device_pointer(mask, device, "request_progress_mask"); + mk2_stage0_device_pointer(out, device, "request_noise"); + native_inflight_noise<<<(BATCH * N_POS * TOP_K / 4 + 255) / 256, 256, 0, + static_cast(stream)>>>( + seeds, generations, mask, out); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown request noise failure"; return 1; } +} +struct MK2InflightResetLayout { + unsigned long long row_stride, source_row_stride, elements; + int rank, item_bytes, source_compact; + unsigned long long sizes[6], strides[6]; +}; +static_assert(sizeof(MK2InflightResetLayout) == 136); +template +__global__ void native_inflight_reset_rows( + Element *destination, const Element *source, const int *slot_ids, + const int *device_count, int capacity, unsigned long long fill_bits, + MK2InflightResetLayout layout) { + const int count = *device_count; + if (count < 0 || count > capacity) return; + const unsigned long long total = (unsigned long long)count * layout.elements; + for (unsigned long long index = (unsigned long long)blockIdx.x * blockDim.x + threadIdx.x; + index < total; index += (unsigned long long)gridDim.x * blockDim.x) { + const unsigned long long selected = index / layout.elements; + const int row = slot_ids[selected]; + if (row < 0 || row >= capacity) continue; + const unsigned long long inner = index % layout.elements; + unsigned long long remaining = inner, offset = (unsigned long long)row * layout.row_stride; + for (int axis = layout.rank - 1; axis >= 0; --axis) { + offset += remaining % layout.sizes[axis] * layout.strides[axis]; + remaining /= layout.sizes[axis]; + } + destination[offset] = source + ? source[(layout.source_compact ? selected : (unsigned long long)row) * + layout.source_row_stride + inner] + : static_cast(fill_bits); + } +} +extern "C" int mk2_inflight_reset_rows( + void *destination, const void *source, const int *slot_ids, const int *device_count, + int capacity, unsigned long long fill_bits, const MK2InflightResetLayout *layout, + void *stream) { + try { + if (!layout || capacity != BATCH || layout->rank < 1 || layout->rank > 6 || + !layout->row_stride || !layout->elements || + (layout->source_compact != 0 && layout->source_compact != 1)) + throw std::runtime_error("invalid request reset layout"); + unsigned long long elements = 1, span = 0; + for (int axis = 0; axis < layout->rank; ++axis) { + if (!layout->sizes[axis] || !layout->strides[axis] || + layout->sizes[axis] > ULLONG_MAX / elements || + layout->sizes[axis] - 1 > (ULLONG_MAX - span) / layout->strides[axis]) + throw std::runtime_error("invalid request reset shape or strides"); + elements *= layout->sizes[axis]; + span += (layout->sizes[axis] - 1) * layout->strides[axis]; + } + if (elements != layout->elements || span >= layout->row_stride || + elements > ULLONG_MAX / (unsigned long long)capacity || + layout->row_stride > ULLONG_MAX / (unsigned long long)capacity || + (source && layout->source_row_stride && layout->source_row_stride < elements)) + throw std::runtime_error("request reset extent exceeds row storage"); + int device = -1; + mk2_check(cudaGetDevice(&device)); + mk2_stage0_device_pointer(destination, device, "reset_destination"); + if (source) mk2_stage0_device_pointer(source, device, "reset_source"); + mk2_stage0_device_pointer(slot_ids, device, "reset_slot_ids"); + mk2_stage0_device_pointer(device_count, device, "reset_count"); + const unsigned long long work = (unsigned long long)capacity * elements; + const int blocks = (int)std::min(132ull, (work - 1ull) / 256ull + 1ull); + const cudaStream_t cuda_stream = static_cast(stream); +#define MK2_RESET_DISPATCH(bytes, type) \ + case bytes: native_inflight_reset_rows<<>>( \ + static_cast(destination), static_cast(source), slot_ids, \ + device_count, capacity, fill_bits, *layout); break + switch (layout->item_bytes) { + MK2_RESET_DISPATCH(1, unsigned char); + MK2_RESET_DISPATCH(2, unsigned short); + MK2_RESET_DISPATCH(4, unsigned int); + MK2_RESET_DISPATCH(8, unsigned long long); + default: throw std::runtime_error("unsupported request reset element size"); + } +#undef MK2_RESET_DISPATCH + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown request reset failure"; return 1; } +} +#endif diff --git a/deploy/nemotron_native/kernels/prefill/b1/dedup_native_prefill_impl.inc b/deploy/nemotron_native/kernels/prefill/b1/dedup_native_prefill_impl.inc new file mode 100644 index 0000000..cdaa0ea --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b1/dedup_native_prefill_impl.inc @@ -0,0 +1,237 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 16 || BATCH == 32 || BATCH == 64)) + throw std::runtime_error("native prefill requires matching batch 1/16/32/64"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (padded > INT_MAX / ((BB_EXPERTS + 1) * D_MODEL)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + float *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(float))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(float))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_dtda, 2 * MB_HEADS); + { + const size_t estride = padded; + const size_t grows = size_t(BB_EXPERTS) * estride; + mk2_check(owned.alloc(&hb.pf_moe_gath, grows * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_gath, 0, grows * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_hc, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_hc, 0, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_gidx, grows * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_gidx, 0, grows * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ecnt, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ecnt, 0, BB_EXPERTS * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ntiles, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ntiles, 0, BB_EXPERTS * sizeof(int))); + hb.pf_moe_estride = int(estride); + hb.pf_moe_gath_v = act_gl{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_v = act_gl{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + hb.pf_moe_gath_vn = act_gl_n{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_vn = act_gl_n{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + scratch_bytes += grows * (D_MODEL + BB_MOE_INTER) * sizeof(bf16); + } + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(owned.alloc(&hb.native_pf_cohort_check, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(hb.native_pf_cohort_check, 0, 2 * sizeof(unsigned int))); +#if PF_ACTIVE_TOKEN_EXPERTS + mk2_check(owned.alloc(&hb.native_pf_expert_masks, (padded/BATCH_PAD)*sizeof(unsigned int))); +#endif + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} +extern "C" int mk2_native_prefill_effective_sequences(MK2NativePrefill *p, int *out) { + try { + if (!p || !out) throw std::runtime_error("null effective-sequence query"); + auto *st = reinterpret_cast(p->base); + mk2_check(cudaMemcpy(out, st->host_bb.native_pf_cohort_check + 1, + sizeof(int), cudaMemcpyDeviceToHost)); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} diff --git a/deploy/nemotron_native/kernels/prefill/b1/dedup_predictor_mk2.cuh b/deploy/nemotron_native/kernels/prefill/b1/dedup_predictor_mk2.cuh new file mode 100644 index 0000000..a640917 --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b1/dedup_predictor_mk2.cuh @@ -0,0 +1,4588 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef PF_ACTIVE_TOKEN_EXPERTS +#define PF_ACTIVE_TOKEN_EXPERTS 0 +#endif +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +#ifndef MK2_PREFILL_ONLY +#define MK2_PREFILL_ONLY 0 +#endif +#ifndef MK2_BATCH_TILE +#define MK2_BATCH_TILE (((MK2_BATCH + 15) / 16) * 16) +#endif +constexpr int BATCH_PAD = MK2_BATCH_TILE; +static_assert(BATCH >= 1 && BATCH <= BATCH_PAD); +static_assert(BATCH_PAD == 16 || BATCH_PAD == 32 || BATCH_PAD == 64 || BATCH_PAD == 128, + "native prefill tile must be 16, 32, 64 or 128"); +static_assert(MK2_PREFILL_ONLY || BATCH_PAD == ((BATCH + 15) / 16) * 16, + "a wider token tile requires a prefill-only library"); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_PW_MULTI +#define PF_SSM_PW_MULTI 4 +#endif +#ifndef PF_SSM_PW_SINGLE +#define PF_SSM_PW_SINGLE 2 +#endif +static_assert(MB_HEAD_DIM % PF_SSM_PW_MULTI == 0 && + MB_HEAD_DIM % PF_SSM_PW_SINGLE == 0, + "head_dim must split evenly into scan channel groups"); +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef MK2_SMEM_SLACK +#define MK2_SMEM_SLACK 1024 +#endif +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +#ifndef MK2_THREADS +#define MK2_THREADS 256 +#endif +constexpr int THREADS = MK2_THREADS; +constexpr int NORM_THREADS = 256; +constexpr int NORM_WARPS = NORM_THREADS / 32; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +#ifndef MK2_GEMM_CONSUMERS +#define MK2_GEMM_CONSUMERS 1 +#endif +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 +#define MK2_HEAD_STAGE_BYTES 0 +#else +#define MK2_HEAD_STAGE_BYTES sizeof(out_tile_f) +#endif +constexpr int NCONS = MK2_GEMM_CONSUMERS; +static_assert(NCONS >= 1 && NCONS <= 3, "consumers occupy warpgroups 0, 2 and 3"); +static_assert(THREADS >= (NCONS + 1) * 128, "one producer plus NCONS consumers"); +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +constexpr int PAD_N = 32; +using act_tile_n = st_bf; +using out_tile_n = st_bf; +using act_gl_n = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_MOE_ACC_C, EPI_ACC_F32 }; +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + unsigned int *native_pf_expert_masks = nullptr; + float *pf_dtda = nullptr; + bf16 *pf_moe_gath = nullptr; + bf16 *pf_moe_hc = nullptr; + int *pf_moe_gidx = nullptr; + int *pf_moe_ecnt = nullptr; + int *pf_moe_ntiles = nullptr; + int pf_moe_estride = 0; + act_gl pf_moe_gath_v; + act_gl pf_moe_hc_v; + act_gl_n pf_moe_gath_vn; + act_gl_n pf_moe_hc_vn; + unsigned int *native_pf_cohort_check = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + volatile unsigned int *generation = state + 1; + const unsigned int seen = *generation; + if (atomicAdd(state, 1u) == gridDim.x - 1) { + atomicExch(state, 0u); + atomicAdd(state + 1, 1u); + } else { + while (*generation == seen) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const bf16 *in = src + (size_t)row * D_MODEL; + float v[PER_THREAD]; + float sum = 0.f, sq = 0.f; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (threadIdx.x < NORM_THREADS) { + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + v[j] = __bfloat162float(in[threadIdx.x + j * NORM_THREADS]); + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + if (lane == 0) { scratch[warp] = sum; scratch[NORM_WARPS + warp] = sq; } + } + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f, q = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) { s += scratch[w]; q += scratch[NORM_WARPS + w]; } + scratch[2 * NORM_WARPS] = s / D_MODEL; + scratch[2 * NORM_WARPS + 1] = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + } + __syncthreads(); + const float mean = scratch[2 * NORM_WARPS], inv = scratch[2 * NORM_WARPS + 1]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + if (threadIdx.x >= NORM_THREADS) break; + const int i = threadIdx.x + j * NORM_THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * weight[i]); + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, + int idx, unsigned mask); +__device__ __forceinline__ void norm_sub_sync(int sub) { + if (THREADS == NORM_THREADS) { __syncthreads(); return; } + asm volatile("bar.sync %0, %1;" :: "r"(6 + sub), "n"(NORM_THREADS) : "memory"); +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, const float *fold_src = nullptr, + const float *fold_routes = nullptr) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + constexpr int NSUB = THREADS / NORM_THREADS; + if (threadIdx.x >= NSUB * NORM_THREADS) return; + const int sub = threadIdx.x / NORM_THREADS; + const int tid = threadIdx.x - sub * NORM_THREADS; + float *const sc = scratch + sub * (NORM_WARPS + 1); + for (int row = blockIdx.x * NSUB + sub; row < n_rows; + row += gridDim.x * NSUB) { + const size_t base = (size_t)row * D_MODEL; + unsigned fold_mask = (1u << BB_EXPERTS) - 1u; + if (fold_src && fold_routes) { + fold_mask = 0u; + const float *rw = fold_routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) + if (rw[e] != 0.f) fold_mask |= 1u << e; + } + float v[PER_THREAD]; + float sq = 0.f; + const int warp = tid >> 5, lane = tid & 31; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + bf16 hv; + if (fold_src) { + hv = bb_moe_fold_deterministic(fold_src, (int)(base + i), fold_mask); + h[base + i] = hv; + } else { + hv = h[base + i]; + } + v[j] = __bfloat162float(hv) + __bfloat162float(r[base + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + if (lane == 0) sc[warp] = sq; + norm_sub_sync(sub); + if (tid == 0) { + float s = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) s += sc[w]; + sc[NORM_WARPS] = rsqrtf(s / D_MODEL + RMS_EPS); + } + norm_sub_sync(sub); + const float inv = sc[NORM_WARPS]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + out[base + i] = __float2bfloat16( + __bfloat162float(weight[i]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + } + norm_sub_sync(sub); + } +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + const int c = idx - row * MB_CONV_DIM; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + const float raw = __bfloat162float( + zxbcdt[(size_t)row * MB_IN_PROJ + MB_INNER + c]); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + xbc_act[(size_t)row * MB_CONV_DIM + c] = __float2bfloat16(acc); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + const int pairs = BATCH * MB_HEADS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_HEADS * split); + const int head = (pair / split) % MB_HEADS; + const int first_channel = (pair % split) * channels; + const int group = head / MB_HEADS_PER_GROUP; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + L.dt_bias[head]); + const float dA = __expf(L.A[head] * dt); + const float Dh = L.D[head]; + float *st = L.ssm_state + (size_t)slot_map[row] * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + for (int p = first_channel; p < first_channel + channels; ++p) { + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + const size_t at = (size_t)p * MB_STATE + n; + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + static_assert(MB_CONV_K == 4, "packed conv taps assume a 4-tap kernel"); + constexpr int CPT = 4; + static_assert(MB_CONV_DIM % CPT == 0, "channel vector width"); + constexpr int CGROUPS = MB_CONV_DIM / CPT; + const int total = n_tokens * CGROUPS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / CGROUPS; + const int c = (idx - t * CGROUPS) * CPT; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + float wv[CPT][MB_CONV_K]; + { + const float4 *wsrc = reinterpret_cast( + L.conv_w + (size_t)c * MB_CONV_K); + #pragma unroll + for (int j = 0; j < CPT / 2; ++j) { + const float4 w4 = wsrc[j]; + const unsigned uu0 = __float_as_uint(w4.x); + const unsigned uu1 = __float_as_uint(w4.y); + const unsigned uu2 = __float_as_uint(w4.z); + const unsigned uu3 = __float_as_uint(w4.w); + wv[2 * j][0] = __uint_as_float(uu0 << 16); + wv[2 * j][1] = __uint_as_float(uu0 & 0xffff0000u); + wv[2 * j][2] = __uint_as_float(uu1 << 16); + wv[2 * j][3] = __uint_as_float(uu1 & 0xffff0000u); + wv[2 * j + 1][0] = __uint_as_float(uu2 << 16); + wv[2 * j + 1][1] = __uint_as_float(uu2 & 0xffff0000u); + wv[2 * j + 1][2] = __uint_as_float(uu3 << 16); + wv[2 * j + 1][3] = __uint_as_float(uu3 & 0xffff0000u); + } + } + float acc[CPT]; + { + const uint2 b2 = *reinterpret_cast(L.conv_b + c); + acc[0] = __uint_as_float(b2.x << 16); + acc[1] = __uint_as_float(b2.x & 0xffff0000u); + acc[2] = __uint_as_float(b2.y << 16); + acc[3] = __uint_as_float(b2.y & 0xffff0000u); + } + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v[CPT]; + if (src < 0) { + #pragma unroll + for (int q = 0; q < CPT; ++q) + v[q] = __bfloat162float( + state[(size_t)q * (MB_CONV_K - 1) + pos + k]); + } else { + const uint2 p = *reinterpret_cast( + zxbcdt + (size_t)src * MB_IN_PROJ + MB_INNER + c); + v[0] = __uint_as_float(p.x << 16); + v[1] = __uint_as_float(p.x & 0xffff0000u); + v[2] = __uint_as_float(p.y << 16); + v[3] = __uint_as_float(p.y & 0xffff0000u); + } + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] += wv[q][k] * v[q]; + } + uint2 out; + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] = acc[q] / (1.f + __expf(-acc[q])); + out.x = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[0])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[1])) << 16); + out.y = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[2])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[3])) << 16); + *reinterpret_cast(xbc_act + (size_t)t * MB_CONV_DIM + c) = out; + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_dtda_prefill( + const MambaLayer &L, const bf16 *zxbcdt, float *dtda, int n_tokens) { + const int total = n_tokens * MB_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_HEADS; + const int head = idx - t * MB_HEADS; + const float dt = mamba_softplus( + __bfloat162float(zxbcdt[(size_t)t * MB_IN_PROJ + MB_INNER + + MB_CONV_DIM + head]) + L.dt_bias[head]); + dtda[(size_t)idx * 2] = dt; + dtda[(size_t)idx * 2 + 1] = __expf(L.A[head] * dt); + } +} +template +__device__ __forceinline__ void mamba_ssm_prefill_impl( + const MambaLayer &L, const bf16 *xbc_act, const float *dtda, bf16 *gated, + int n_pos, int n_seq, int slot, const int *seq_slot) { + constexpr int NS = MB_STATE / 32; + constexpr int PG = MB_HEAD_DIM / PW; + constexpr bool RS = (PW % 4) == 0; + constexpr int H1 = RS ? PW / 2 : 1; + constexpr int H2 = RS ? PW / 4 : 1; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const bool keep_hi = (lane & 16) != 0; + const bool keep_hi2 = (lane & 8) != 0; + const int warps_total = gridDim.x * WARPS; + const int units = n_seq * MB_HEADS * PG; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int pg = unit % PG; + const int rest = unit / PG; + const int head = rest % MB_HEADS; + const int seq = rest / MB_HEADS; + const int p0 = pg * PW; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Dh = L.D[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p0 * MB_STATE; + float sreg[PW][NS]; + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + sreg[q][j] = st[(size_t)q * MB_STATE + lane + j * 32]; + } + auto fetch = [&](int p, float (&oB)[NS], float (&oC)[NS], + float (&oX)[PW], float2 &od) { + const size_t tt = (size_t)p * n_seq + seq; + const bf16 *xa2 = xbc_act + tt * MB_CONV_DIM; + const bf16 *Bs = xa2 + MB_INNER + group * MB_STATE; + const bf16 *Cs = xa2 + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + od = *reinterpret_cast( + dtda + (tt * MB_HEADS + head) * 2); + #pragma unroll + for (int j = 0; j < NS; ++j) { + oB[j] = __bfloat162float(Bs[lane + j * 32]); + oC[j] = __bfloat162float(Cs[lane + j * 32]); + } + if constexpr (RS) { + const uint2 *xp = reinterpret_cast( + xa2 + head * MB_HEAD_DIM + p0); + #pragma unroll + for (int q = 0; q < PW; q += 4) { + const uint2 w = xp[q >> 2]; + oX[q] = __uint_as_float(w.x << 16); + oX[q + 1] = __uint_as_float(w.x & 0xffff0000u); + oX[q + 2] = __uint_as_float(w.y << 16); + oX[q + 3] = __uint_as_float(w.y & 0xffff0000u); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) + oX[q] = __bfloat162float(xa2[head * MB_HEAD_DIM + p0 + q]); + } + }; + float pB[NS], pC[NS], pX[PW]; + float2 pd; + fetch(0, pB, pC, pX, pd); + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const float dt = pd.x; + const float dA = pd.y; + float Bv[NS], Cv[NS]; + #pragma unroll + for (int j = 0; j < NS; ++j) { + Bv[j] = pB[j] * dt; + Cv[j] = pC[j]; + } + float acc[PW], xvs[PW]; + #pragma unroll + for (int q = 0; q < PW; ++q) xvs[q] = pX[q]; + if (pos + 1 < n_pos) fetch(pos + 1, pB, pC, pX, pd); + #pragma unroll + for (int q = 0; q < PW; ++q) { + const float xv = xvs[q]; + float a = 0.f; + #pragma unroll + for (int j = 0; j < NS; ++j) { + sreg[q][j] = sreg[q][j] * dA + Bv[j] * xv; + a += sreg[q][j] * Cv[j]; + } + acc[q] = a; + } + if constexpr (RS) { + float m1[H1], t1[H1], x1[H1]; + #pragma unroll + for (int q = 0; q < H1; ++q) { + m1[q] = keep_hi ? acc[q + H1] : acc[q]; + t1[q] = keep_hi ? acc[q] : acc[q + H1]; + x1[q] = keep_hi ? xvs[q + H1] : xvs[q]; + } + #pragma unroll + for (int q = 0; q < H1; ++q) + m1[q] += __shfl_xor_sync(0xffffffff, t1[q], 16); + float m2[H2], t2[H2], x2[H2]; + #pragma unroll + for (int q = 0; q < H2; ++q) { + m2[q] = keep_hi2 ? m1[q + H2] : m1[q]; + t2[q] = keep_hi2 ? m1[q] : m1[q + H2]; + x2[q] = keep_hi2 ? x1[q + H2] : x1[q]; + } + #pragma unroll + for (int q = 0; q < H2; ++q) + m2[q] += __shfl_xor_sync(0xffffffff, t2[q], 8); + #pragma unroll + for (int q = 0; q < H2; ++q) { + #pragma unroll + for (int off = 4; off > 0; off >>= 1) + m2[q] += __shfl_xor_sync(0xffffffff, m2[q], off); + } + if ((lane & 7) == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + + p0 + (keep_hi ? H1 : 0) + (keep_hi2 ? H2 : 0); + #pragma unroll + for (int q = 0; q < H2; ++q) + dst[q] = __float2bfloat16(m2[q] + x2[q] * Dh); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc[q] += __shfl_xor_sync(0xffffffff, acc[q], off); + } + if (lane == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + p0; + #pragma unroll + for (int q = 0; q < PW; ++q) + dst[q] = __float2bfloat16(acc[q] + xvs[q] * Dh); + } + } + } + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + st[(size_t)q * MB_STATE + lane + j * 32] = sreg[q][j]; + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, + const float *dtda, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + (void)zxbcdt; + if (n_seq > 1) + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens / n_seq, n_seq, slot, seq_slot); + else + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens, 1, slot, nullptr); +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16( + held[j] * inv * __bfloat162float(L.gnorm_w[idx])); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = BATCH * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / BB_ATT_HEADS; + const int head = pair % BB_ATT_HEADS; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const int kv = head / kv_per_q; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *q = base + head * BB_HEAD_DIM; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (head % kv_per_q == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + for (int tpos = 0; tpos < T; ++tpos) { +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new + : kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#endif + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m, s); + const float rescale = __expf(m - m_new); + const float w = __expf(s - m_new); + denom = denom * rescale + w; +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *vt = kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#else + const bf16 *vt = cur_token + ? v_new + : kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + w * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + constexpr int CHUNKS = BB_HEAD_DIM / 8; + const int total = n_tokens * BB_KV_HEADS * CHUNKS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int c = idx % CHUNKS; + const int rest = idx / CHUNKS; + const int t = rest / BB_KV_HEADS; + const int kv = rest - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM + c * 8; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM + c * 8; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv) + c * 8; + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv) + c * 8; + *reinterpret_cast(kc) = *reinterpret_cast(k_new); + *reinterpret_cast(vc) = *reinterpret_cast(v_new); + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float qv[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qv[j] = __bfloat162float(q[lane + j * 32]); + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + const int *const krow = + L.block_table + ((size_t)seq * 2 + 0) * L.max_blocks_per_seq; + const int *const vrow = + L.block_table + ((size_t)seq * 2 + 1) * L.max_blocks_per_seq; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + const int head_bias = kv * L.tokens_per_block * BB_HEAD_DIM; + int page = -1; + const bf16 *kpage = nullptr, *vpage = nullptr; + float kn[PER_LANE], vn[PER_LANE]; + { + page = 0; + kpage = L.pool_base + (size_t)(krow[0] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + (size_t)(vrow[0] & 0x7FFFFFFF) * elems_per_block; + const int at = head_bias + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + for (int s = 0; s <= pos; ++s) { + float kc[PER_LANE], vc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { kc[j] = kn[j]; vc[j] = vn[j]; } + if (s < pos) { + const int np = (s + 1) / L.tokens_per_block; + if (np != page) { + page = np; + kpage = L.pool_base + + (size_t)(krow[np] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + + (size_t)(vrow[np] & 0x7FFFFFFF) * elems_per_block; + } + const int at = head_bias + + ((s + 1) % L.tokens_per_block) * BB_HEAD_DIM + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) dot += qv[j] * kc[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[j] = acc[j] * rescale + wt * vc[j]; + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH, unsigned int *tile_masks = nullptr, + bf16 *gath = nullptr, int *gidx = nullptr, int *ecnt = nullptr, + int estride = 0, bf16 *row_stage = nullptr) { + __shared__ float score[BB_EXPERTS]; + __shared__ int s_pick[BB_TOP_K], s_slot[BB_TOP_K]; + constexpr int CH = D_MODEL / 8; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + { + const float4 *gsrc = + reinterpret_cast(normed + (size_t)row * D_MODEL); + float4 *gdst = reinterpret_cast(row_stage); + for (int i = threadIdx.x; i < CH; i += THREADS) gdst[i] = gsrc[i]; + } + if ((int)threadIdx.x < BB_EXPERTS) + moe_w[(size_t)row * BB_EXPERTS + threadIdx.x] = 0.f; + __syncthreads(); + for (int e = warp; e < BB_EXPERTS; e += WARPS) { + const bf16 *wr = gate_w + (size_t)e * D_MODEL; + float acc = 0.f; + for (int i = lane; i < D_MODEL; i += 32) + acc += __bfloat162float(row_stage[i]) * + __bfloat162float(wr[i]); + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) score[e] = 1.f / (1.f + __expf(-acc)); + } + __syncthreads(); + if (threadIdx.x == 0) { + float remaining[BB_EXPERTS]; + int pick[BB_TOP_K]; + for (int e = 0; e < BB_EXPERTS; ++e) + remaining[e] = score[e] + __bfloat162float(gate_bias[e]); + float total = 0.f; + for (int k = 0; k < BB_TOP_K; ++k) { + int best = 0; + for (int e = 1; e < BB_EXPERTS; ++e) + if (remaining[e] > remaining[best]) best = e; + pick[k] = best; + total += score[best]; + remaining[best] = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + unsigned int bits = 0u; + for (int k = 0; k < BB_TOP_K; ++k) { + moe_w[row * BB_EXPERTS + pick[k]] = score[pick[k]] * norm; + bits |= 1u << pick[k]; + s_pick[k] = pick[k]; + } + if (tile_masks) atomicOr(&tile_masks[row / BATCH_PAD], bits); + } + __syncthreads(); + if (gath) { + if ((int)threadIdx.x < BB_TOP_K) { + const int k = threadIdx.x; + const int slot = atomicAdd(&ecnt[s_pick[k]], 1); + s_slot[k] = slot; + gidx[(size_t)s_pick[k] * estride + slot] = row; + } + __syncthreads(); + const float4 *src = reinterpret_cast(row_stage); + for (int k = 0; k < BB_TOP_K; ++k) { + float4 *dst = reinterpret_cast( + gath + ((size_t)s_pick[k] * estride + s_slot[k]) * D_MODEL); + for (int i = threadIdx.x; i < CH; i += THREADS) dst[i] = src[i]; + } + } + __syncthreads(); + } +} +__device__ __forceinline__ unsigned moe_row_mask(const float *routes, int row) { + unsigned m = 0u; + const float *rw = routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) if (rw[e] != 0.f) m |= 1u << e; + return m; +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx, + unsigned mask) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + for (unsigned m = mask; m; m &= m - 1u) + value = __fadd_rn(value, parts[base + (size_t)(__ffs(m) - 1) * D_MODEL]); + return __float2bfloat16( + __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL])); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + uint64_t *keys, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int n = threadIdx.x; n < SORT_N; n += THREADS) { + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + keys[n] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + __syncthreads(); + for (int k = 2; k <= SORT_N; k <<= 1) { + if (k >= 64) __syncthreads(); + for (int j = k >> 1; j > 0; j >>= 1) { + for (int i = threadIdx.x; i < SORT_N; i += THREADS) { + const int partner = i ^ j; + if (partner > i) { + const bool descending = ((i & k) == 0); + if ((keys[i] < keys[partner]) == descending) { + const uint64_t t = keys[i]; + keys[i] = keys[partner]; + keys[partner] = t; + } + } + } + if (j >= 32) __syncthreads(); + else __syncwarp(); + } + } + __syncthreads(); + if (threadIdx.x == 0) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float v; + int n; + if (!eos_used && eos_v > av) { + v = eos_v; + n = AUDIO_EOS; + eos_used = true; + } else { + v = av; + n = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = + v + gumbel[((size_t)row * N_POS + pos) * TOP_K + j]; + if (noised > best_noised) { best_noised = noised; best_code = n; } + } + out_codes[row] = best_code; + } + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void prefill_expert_masks(const float *routes, unsigned int *masks, int rows) { + const int lane=threadIdx.x&31, warp=threadIdx.x>>5; + const int nt=(rows+BATCH_PAD-1)/BATCH_PAD; + for (int t=blockIdx.x*WARPS+warp;t(gath + g * D_MODEL + c * 8); + if (local < ecnt[e]) + *dst = *reinterpret_cast( + normed + (size_t)gidx[g] * D_MODEL + c * 8); + else + *dst = make_float4(0.f, 0.f, 0.f, 0.f); + } +} +__device__ __forceinline__ int gemm_eff(int tiles, int reps) { + if (NCONS <= 1) return 1; + const int g = (int)gridDim.x; + const int rounds1 = (tiles * reps + g - 1) / g; + if constexpr (NCONS == 3) { + const int max_eff = reps <= BB_EXPERTS ? 2 : 3; + const int rounds2 = (((tiles + 1) / 2) * reps + g - 1) / g; + if (max_eff == 2) return rounds2 < rounds1 ? 2 : 1; + const int rounds3 = (((tiles + 2) / 3) * reps + g - 1) / g; + if (rounds3 < rounds2 && rounds3 < rounds1) return 3; + return rounds2 < rounds1 ? 2 : 1; + } else { + const int roundsN = (((tiles + NCONS - 1) / NCONS) * reps + g - 1) / g; + return roundsN < rounds1 ? NCONS : 1; + } +} +__device__ __forceinline__ int gemm_groups(int tiles, int reps) { + const int eff = gemm_eff(tiles, reps); + return (tiles + eff - 1) / eff; +} +__device__ __forceinline__ int gemm_consumer_id() { + const int g = warpgroup::groupid(); + if (g == 1) return -1; + const int c = g == 0 ? 0 : g - 1; + return c < NCONS ? c : -1; +} +__device__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][BATCH_PAD]; + __shared__ int moe_gidx_stage[NCONS][BATCH_PAD]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + a_s[s], a, {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + a_s[s], a, {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / BATCH_PAD) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, BATCH_PAD> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], a_s[s]); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * BATCH_PAD; + if (epi == EPI_BIAS) { + warpgroup::store(o_sf, accum); + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + dest[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(o_sf[int2{col, row}] + + __bfloat162float(bias[base + col])); + } + warpgroup::sync(1 + cons); + continue; + } + warpgroup::store(o_s[cons], accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_s[cons][int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__device__ void gemm_narrow( + const weight_gl &w, const act_gl_n &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = PAD_N, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][PAD_N]; + __shared__ int moe_gidx_stage[NCONS][PAD_N]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + PAD_N - 1) / PAD_N; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + PAD_N - 1) / PAD_N; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / PAD_N) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, PAD_N> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], + *reinterpret_cast(&a_s[s])); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * PAD_N; + out_tile_n &o_sn = *reinterpret_cast(&o_s[cons]); + warpgroup::store(o_sn, accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_sn[int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&w_s)[STAGES * NCONS] = + alloc.allocate(); + act_tile (&a_s)[STAGES] = alloc.allocate(); + out_tile (&o_s)[NCONS] = alloc.allocate(); +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 + static_assert(NCONS * sizeof(out_tile) >= sizeof(out_tile_f), + "aliased head staging must fit the output tiles"); + out_tile_f &o_sf = *reinterpret_cast(&o_s[0]); +#else + out_tile_f (&o_sf) = alloc.allocate(); +#endif + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_WARPS + 2]; +#if MK2_PREFILL_ONLY + bf16 *const route_stage = reinterpret_cast(&o_s[0]); + static_assert(sizeof(o_s) >= D_MODEL * sizeof(bf16), + "router row stage must fit the output tiles"); +#else + __shared__ __align__(16) uint64_t sort_keys[SORT_N]; + bf16 *const route_stage = reinterpret_cast(&sort_keys[0]); + static_assert(sizeof(sort_keys) >= D_MODEL * sizeof(bf16), + "router row stage must fit the sampler workspace"); +#endif + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS * NCONS, 0); + for (int c = 0; c < CONSUMER_WARPS * NCONS; ++c) arrive(drained[s]); + } + } + int gemm_iter = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map); + grid_barrier(g.barrier); + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map); + grid_barrier(g.barrier); + mamba_gated_norm(L, bb->zxbcdt, bb->gated); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, nullptr, 1, BATCH, + bb->moe_w, BB_MOE_INTER / OUT_ATOM); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + 1, BATCH, bb->moe_w, D_MODEL / OUT_ATOM); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_deterministic(bb->moe_out, idx, + (1u << BB_EXPERTS) - 1u); + grid_barrier(g.barrier); + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int original_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + bool identical_cohort = false; + if constexpr (BATCH > 1) { + if (bb->native_pf_cohort_check && bb->native_prefill_reset && + bb->pf_n_seq == BATCH && !bb->pf_fuse_decode && + !bb->pf_layer_out && !bb->moe_w_inject && !bb->moe_w_audit) { + if (blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[0] = 0; + grid_barrier(g.barrier); + const auto *bits = reinterpret_cast(bb->prompt); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < original_total * D_MODEL; idx += gridDim.x * THREADS) { + const int canonical = (idx / (BATCH * D_MODEL)) * BATCH * D_MODEL + idx % D_MODEL; + if (bits[idx] != bits[canonical]) atomicOr(bb->native_pf_cohort_check, 1u); + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + if (bb->pf_seq_slot[row] != row || bb->slot_map[row] != row) + atomicOr(bb->native_pf_cohort_check, 1u); + grid_barrier(g.barrier); + identical_cohort = bb->native_pf_cohort_check[0] == 0; + } + } + const int n_total = identical_cohort ? original_total / BATCH : original_total; + if (bb->native_pf_cohort_check && blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[1] = identical_cohort ? 1 : BATCH; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int reset_slots = identical_cohort ? 1 : BATCH; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * reset_slots * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % reset_slots; + const MambaLayer &L = bb->mamba[row_layer / reset_slots]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * reset_slots * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % reset_slots; + const MambaLayer &L = bb->mamba[row_layer / reset_slots]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = original_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + const int src = identical_cohort ? (idx / D_MODEL) * BATCH * D_MODEL + idx % D_MODEL : idx; + bb->pf_h[idx] = bb->prompt[src]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = identical_cohort ? 1 : (bb->pf_n_seq > 0 ? bb->pf_n_seq : 1); + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T, + (!bb->pf_layer_out && layer > 0 && + bb->moe_of_layer[layer - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + if (bb->native_pf_expert_masks) + for (int i = blockIdx.x * THREADS + threadIdx.x; i < n_tok; + i += gridDim.x * THREADS) + bb->native_pf_expert_masks[i] = 0u; + if (bb->pf_moe_ecnt) + for (int e = blockIdx.x * THREADS + threadIdx.x; e < BB_EXPERTS; + e += gridDim.x * THREADS) + bb->pf_moe_ecnt[e] = 0; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + mamba_dtda_prefill(L, bb->pf_zxbcdt, bb->pf_dtda, T); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_dtda, + bb->pf_gated, T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, + T, bb->native_pf_expert_masks, bb->pf_moe_gath, + bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, route_stage); + } + grid_barrier(g.barrier); + int max_ecnt = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + max_ecnt = max(max_ecnt, bb->pf_moe_ecnt[e]); + const bool narrow = BATCH_PAD > PAD_N && max_ecnt <= PAD_N; + const int etile = narrow ? PAD_N : BATCH_PAD; + int nt_sum = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_sum += (bb->pf_moe_ecnt[e] + etile - 1) / etile; + const int up_steps = + gemm_groups(BB_MOE_INTER / OUT_ATOM, nt_sum) * + nt_sum * (D_MODEL / KCHUNK); + const int shup_steps = + gemm_groups(BB_SHARED_INTER / OUT_ATOM, n_tok) * + n_tok * (D_MODEL / KCHUNK); + const int down_steps = + gemm_groups(D_MODEL / OUT_ATOM, nt_sum) * + nt_sum * (BB_MOE_INTER / KCHUNK); + const int shdown_steps = + gemm_groups(D_MODEL / OUT_ATOM, n_tok) * + n_tok * (BB_SHARED_INTER / KCHUNK); + if (narrow) + gemm_narrow(L.exp_up, bb->pf_moe_gath_vn, + BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, PAD_N, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + else + gemm(L.exp_up, bb->pf_moe_gath_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, BATCH_PAD, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + up_steps, up_steps + shup_steps); + grid_barrier(g.barrier); + if (narrow) + gemm_narrow(L.exp_down, bb->pf_moe_hc_vn, + BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, PAD_N, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + else + gemm(L.exp_down, bb->pf_moe_hc_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + down_steps, down_steps + shdown_steps); + grid_barrier(g.barrier); + if (bb->pf_layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * D_MODEL; idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = + bb_moe_fold_deterministic(bb->pf_moe_out, idx, + moe_row_mask(bb->pf_moe_w, idx / D_MODEL)); + grid_barrier(g.barrier); + } + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, + identical_cohort ? bb->pf_normed : bb->pf_final, ln_scratch, + n_total, + (!bb->pf_layer_out && n_backbone_layers > 0 && + bb->moe_of_layer[n_backbone_layers - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + grid_barrier(g.barrier); + if (identical_cohort) { + const int worker = blockIdx.x * THREADS + threadIdx.x; + const int workers = gridDim.x * THREADS; + for (int i = worker; i < original_total * D_MODEL; i += workers) + bb->pf_final[i] = bb->pf_normed[(i / (BATCH * D_MODEL)) * D_MODEL + i % D_MODEL]; + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = worker; i < 14 * (BATCH - 1) * CV; i += workers) { + const int cell = i % CV, pair = i / CV; + const int row = pair % (BATCH - 1) + 1; + const MambaLayer &L = bb->mamba[pair / (BATCH - 1)]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = L.conv_state[cell]; + } + for (int i = worker; i < 14 * (BATCH - 1) * SS; i += workers) { + const int cell = i % SS, pair = i / SS; + const int row = pair % (BATCH - 1) + 1; + const MambaLayer &L = bb->mamba[pair / (BATCH - 1)]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = L.ssm_state[cell]; + } + const int kv_columns = BB_KV_HEADS * BB_HEAD_DIM; + for (int i = worker; i < 3 * 2 * (BATCH - 1) * n_total * kv_columns; i += workers) { + int q = i / kv_columns; + const int col = i % kv_columns; + const int pos = q % n_total; q /= n_total; + const int row = q % (BATCH - 1) + 1; q /= BATCH - 1; + const int kind = q % 2, layer = q / 2; + const AttnLayer &L = bb->attn[layer]; + const int offset = kv_offset(L, pos, col / BB_HEAD_DIM) + col % BB_HEAD_DIM; + kv_slot(L, row, kind, pos)[offset] = kv_slot(L, 0, kind, pos)[offset]; + } + grid_barrier(g.barrier); + } + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } +#if MK2_PREFILL_ONLY + return; +#else + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, bb->phoneme_b); + grid_barrier(g.barrier); + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + if (threadIdx.x == 0) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = 0; n < BB_PHONEME; ++n) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + bb->phoneme_out[row] = arg; + } + __syncthreads(); + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + for (int i = lane; i < D_HEAD; i += 32) { + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + for (int t = 0; t <= pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + for (int i = lane; i < D_HEAD; i += 32) { + float acc = 0.f; + for (int t = 0; t <= pos; ++t) + acc += sc[t] * __bfloat162float( + g.vcache[cache_base + t * D_HEAD + i]); + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + i] = + __float2bfloat16(acc); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, g.sampled_pos, + sort_keys, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { +#ifdef ABLATE_SAMPLE + __syncthreads(); +#endif + } else { + grid_barrier(g.barrier); + } + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x < BATCH && blockIdx.x == 0) { + const int code = g.sampled_pos[threadIdx.x]; + g.sampled[threadIdx.x * N_POS + pos] = code; + if (g.native_codes) + g.native_codes[threadIdx.x * N_POS + pos] = (long long)code; + } + grid_barrier(g.barrier); + } + if (g.native_seq_len && blockIdx.x == 0 && threadIdx.x < BATCH) + g.native_seq_len[threadIdx.x] += 1; + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +#endif +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 2 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK), + 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = MK2_PREFILL_ONLY + ? out->sms * out->blocks_per_sm + : std::min(out->sms * out->blocks_per_sm, D_FFN / OUT_ATOM); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + 0, + nullptr, + nullptr, + nullptr, nullptr, nullptr, nullptr, nullptr, 0, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "dedup_native_prefill_impl.inc" diff --git a/deploy/nemotron_native/kernels/prefill/b1/mk2_abi.cuh b/deploy/nemotron_native/kernels/prefill/b1/mk2_abi.cuh new file mode 100644 index 0000000..95a253b --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b1/mk2_abi.cuh @@ -0,0 +1,82 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} diff --git a/deploy/nemotron_native/kernels/prefill/b1/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/prefill/b1/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b1/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/prefill/b1/prefill.cu b/deploy/nemotron_native/kernels/prefill/b1/prefill.cu new file mode 100644 index 0000000..9854d41 --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b1/prefill.cu @@ -0,0 +1,18 @@ +#define MK2_BATCH 1 +#define MK2_BATCH_TILE 64 +#define MK2_PREFILL_ONLY 1 +#define STAGES 3 +#define KCHUNK 128 +#define MK2_GEMM_CONSUMERS 3 +#define MK2_THREADS 512 +#define MK2_NO_MAIN 1 +#define MK2_DIRECT_INPUT 1 +#define MK2_B1_LOCAL_BARRIERS 1 +#define MK2_SKIP_HIDDEN_SNAPSHOT 1 +#define MK2_ACTIVE_EXPERTS 1 +#define MK2_SSM_SPLITS 1 +#define PF_SSM_PW_MULTI 12 +#define PF_ACTIVE_TOKEN_EXPERTS 1 +#include +#include +#include "dedup_predictor_mk2.cuh" diff --git a/deploy/nemotron_native/kernels/prefill/b16/dedup_native_prefill_impl.inc b/deploy/nemotron_native/kernels/prefill/b16/dedup_native_prefill_impl.inc new file mode 100644 index 0000000..cdaa0ea --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b16/dedup_native_prefill_impl.inc @@ -0,0 +1,237 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 16 || BATCH == 32 || BATCH == 64)) + throw std::runtime_error("native prefill requires matching batch 1/16/32/64"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (padded > INT_MAX / ((BB_EXPERTS + 1) * D_MODEL)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + float *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(float))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(float))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_dtda, 2 * MB_HEADS); + { + const size_t estride = padded; + const size_t grows = size_t(BB_EXPERTS) * estride; + mk2_check(owned.alloc(&hb.pf_moe_gath, grows * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_gath, 0, grows * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_hc, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_hc, 0, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_gidx, grows * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_gidx, 0, grows * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ecnt, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ecnt, 0, BB_EXPERTS * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ntiles, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ntiles, 0, BB_EXPERTS * sizeof(int))); + hb.pf_moe_estride = int(estride); + hb.pf_moe_gath_v = act_gl{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_v = act_gl{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + hb.pf_moe_gath_vn = act_gl_n{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_vn = act_gl_n{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + scratch_bytes += grows * (D_MODEL + BB_MOE_INTER) * sizeof(bf16); + } + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(owned.alloc(&hb.native_pf_cohort_check, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(hb.native_pf_cohort_check, 0, 2 * sizeof(unsigned int))); +#if PF_ACTIVE_TOKEN_EXPERTS + mk2_check(owned.alloc(&hb.native_pf_expert_masks, (padded/BATCH_PAD)*sizeof(unsigned int))); +#endif + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} +extern "C" int mk2_native_prefill_effective_sequences(MK2NativePrefill *p, int *out) { + try { + if (!p || !out) throw std::runtime_error("null effective-sequence query"); + auto *st = reinterpret_cast(p->base); + mk2_check(cudaMemcpy(out, st->host_bb.native_pf_cohort_check + 1, + sizeof(int), cudaMemcpyDeviceToHost)); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} diff --git a/deploy/nemotron_native/kernels/prefill/b16/dedup_predictor_mk2.cuh b/deploy/nemotron_native/kernels/prefill/b16/dedup_predictor_mk2.cuh new file mode 100644 index 0000000..a640917 --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b16/dedup_predictor_mk2.cuh @@ -0,0 +1,4588 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef PF_ACTIVE_TOKEN_EXPERTS +#define PF_ACTIVE_TOKEN_EXPERTS 0 +#endif +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +#ifndef MK2_PREFILL_ONLY +#define MK2_PREFILL_ONLY 0 +#endif +#ifndef MK2_BATCH_TILE +#define MK2_BATCH_TILE (((MK2_BATCH + 15) / 16) * 16) +#endif +constexpr int BATCH_PAD = MK2_BATCH_TILE; +static_assert(BATCH >= 1 && BATCH <= BATCH_PAD); +static_assert(BATCH_PAD == 16 || BATCH_PAD == 32 || BATCH_PAD == 64 || BATCH_PAD == 128, + "native prefill tile must be 16, 32, 64 or 128"); +static_assert(MK2_PREFILL_ONLY || BATCH_PAD == ((BATCH + 15) / 16) * 16, + "a wider token tile requires a prefill-only library"); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_PW_MULTI +#define PF_SSM_PW_MULTI 4 +#endif +#ifndef PF_SSM_PW_SINGLE +#define PF_SSM_PW_SINGLE 2 +#endif +static_assert(MB_HEAD_DIM % PF_SSM_PW_MULTI == 0 && + MB_HEAD_DIM % PF_SSM_PW_SINGLE == 0, + "head_dim must split evenly into scan channel groups"); +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef MK2_SMEM_SLACK +#define MK2_SMEM_SLACK 1024 +#endif +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +#ifndef MK2_THREADS +#define MK2_THREADS 256 +#endif +constexpr int THREADS = MK2_THREADS; +constexpr int NORM_THREADS = 256; +constexpr int NORM_WARPS = NORM_THREADS / 32; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +#ifndef MK2_GEMM_CONSUMERS +#define MK2_GEMM_CONSUMERS 1 +#endif +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 +#define MK2_HEAD_STAGE_BYTES 0 +#else +#define MK2_HEAD_STAGE_BYTES sizeof(out_tile_f) +#endif +constexpr int NCONS = MK2_GEMM_CONSUMERS; +static_assert(NCONS >= 1 && NCONS <= 3, "consumers occupy warpgroups 0, 2 and 3"); +static_assert(THREADS >= (NCONS + 1) * 128, "one producer plus NCONS consumers"); +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +constexpr int PAD_N = 32; +using act_tile_n = st_bf; +using out_tile_n = st_bf; +using act_gl_n = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_MOE_ACC_C, EPI_ACC_F32 }; +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + unsigned int *native_pf_expert_masks = nullptr; + float *pf_dtda = nullptr; + bf16 *pf_moe_gath = nullptr; + bf16 *pf_moe_hc = nullptr; + int *pf_moe_gidx = nullptr; + int *pf_moe_ecnt = nullptr; + int *pf_moe_ntiles = nullptr; + int pf_moe_estride = 0; + act_gl pf_moe_gath_v; + act_gl pf_moe_hc_v; + act_gl_n pf_moe_gath_vn; + act_gl_n pf_moe_hc_vn; + unsigned int *native_pf_cohort_check = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + volatile unsigned int *generation = state + 1; + const unsigned int seen = *generation; + if (atomicAdd(state, 1u) == gridDim.x - 1) { + atomicExch(state, 0u); + atomicAdd(state + 1, 1u); + } else { + while (*generation == seen) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const bf16 *in = src + (size_t)row * D_MODEL; + float v[PER_THREAD]; + float sum = 0.f, sq = 0.f; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (threadIdx.x < NORM_THREADS) { + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + v[j] = __bfloat162float(in[threadIdx.x + j * NORM_THREADS]); + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + if (lane == 0) { scratch[warp] = sum; scratch[NORM_WARPS + warp] = sq; } + } + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f, q = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) { s += scratch[w]; q += scratch[NORM_WARPS + w]; } + scratch[2 * NORM_WARPS] = s / D_MODEL; + scratch[2 * NORM_WARPS + 1] = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + } + __syncthreads(); + const float mean = scratch[2 * NORM_WARPS], inv = scratch[2 * NORM_WARPS + 1]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + if (threadIdx.x >= NORM_THREADS) break; + const int i = threadIdx.x + j * NORM_THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * weight[i]); + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, + int idx, unsigned mask); +__device__ __forceinline__ void norm_sub_sync(int sub) { + if (THREADS == NORM_THREADS) { __syncthreads(); return; } + asm volatile("bar.sync %0, %1;" :: "r"(6 + sub), "n"(NORM_THREADS) : "memory"); +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, const float *fold_src = nullptr, + const float *fold_routes = nullptr) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + constexpr int NSUB = THREADS / NORM_THREADS; + if (threadIdx.x >= NSUB * NORM_THREADS) return; + const int sub = threadIdx.x / NORM_THREADS; + const int tid = threadIdx.x - sub * NORM_THREADS; + float *const sc = scratch + sub * (NORM_WARPS + 1); + for (int row = blockIdx.x * NSUB + sub; row < n_rows; + row += gridDim.x * NSUB) { + const size_t base = (size_t)row * D_MODEL; + unsigned fold_mask = (1u << BB_EXPERTS) - 1u; + if (fold_src && fold_routes) { + fold_mask = 0u; + const float *rw = fold_routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) + if (rw[e] != 0.f) fold_mask |= 1u << e; + } + float v[PER_THREAD]; + float sq = 0.f; + const int warp = tid >> 5, lane = tid & 31; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + bf16 hv; + if (fold_src) { + hv = bb_moe_fold_deterministic(fold_src, (int)(base + i), fold_mask); + h[base + i] = hv; + } else { + hv = h[base + i]; + } + v[j] = __bfloat162float(hv) + __bfloat162float(r[base + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + if (lane == 0) sc[warp] = sq; + norm_sub_sync(sub); + if (tid == 0) { + float s = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) s += sc[w]; + sc[NORM_WARPS] = rsqrtf(s / D_MODEL + RMS_EPS); + } + norm_sub_sync(sub); + const float inv = sc[NORM_WARPS]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + out[base + i] = __float2bfloat16( + __bfloat162float(weight[i]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + } + norm_sub_sync(sub); + } +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + const int c = idx - row * MB_CONV_DIM; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + const float raw = __bfloat162float( + zxbcdt[(size_t)row * MB_IN_PROJ + MB_INNER + c]); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + xbc_act[(size_t)row * MB_CONV_DIM + c] = __float2bfloat16(acc); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + const int pairs = BATCH * MB_HEADS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_HEADS * split); + const int head = (pair / split) % MB_HEADS; + const int first_channel = (pair % split) * channels; + const int group = head / MB_HEADS_PER_GROUP; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + L.dt_bias[head]); + const float dA = __expf(L.A[head] * dt); + const float Dh = L.D[head]; + float *st = L.ssm_state + (size_t)slot_map[row] * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + for (int p = first_channel; p < first_channel + channels; ++p) { + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + const size_t at = (size_t)p * MB_STATE + n; + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + static_assert(MB_CONV_K == 4, "packed conv taps assume a 4-tap kernel"); + constexpr int CPT = 4; + static_assert(MB_CONV_DIM % CPT == 0, "channel vector width"); + constexpr int CGROUPS = MB_CONV_DIM / CPT; + const int total = n_tokens * CGROUPS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / CGROUPS; + const int c = (idx - t * CGROUPS) * CPT; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + float wv[CPT][MB_CONV_K]; + { + const float4 *wsrc = reinterpret_cast( + L.conv_w + (size_t)c * MB_CONV_K); + #pragma unroll + for (int j = 0; j < CPT / 2; ++j) { + const float4 w4 = wsrc[j]; + const unsigned uu0 = __float_as_uint(w4.x); + const unsigned uu1 = __float_as_uint(w4.y); + const unsigned uu2 = __float_as_uint(w4.z); + const unsigned uu3 = __float_as_uint(w4.w); + wv[2 * j][0] = __uint_as_float(uu0 << 16); + wv[2 * j][1] = __uint_as_float(uu0 & 0xffff0000u); + wv[2 * j][2] = __uint_as_float(uu1 << 16); + wv[2 * j][3] = __uint_as_float(uu1 & 0xffff0000u); + wv[2 * j + 1][0] = __uint_as_float(uu2 << 16); + wv[2 * j + 1][1] = __uint_as_float(uu2 & 0xffff0000u); + wv[2 * j + 1][2] = __uint_as_float(uu3 << 16); + wv[2 * j + 1][3] = __uint_as_float(uu3 & 0xffff0000u); + } + } + float acc[CPT]; + { + const uint2 b2 = *reinterpret_cast(L.conv_b + c); + acc[0] = __uint_as_float(b2.x << 16); + acc[1] = __uint_as_float(b2.x & 0xffff0000u); + acc[2] = __uint_as_float(b2.y << 16); + acc[3] = __uint_as_float(b2.y & 0xffff0000u); + } + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v[CPT]; + if (src < 0) { + #pragma unroll + for (int q = 0; q < CPT; ++q) + v[q] = __bfloat162float( + state[(size_t)q * (MB_CONV_K - 1) + pos + k]); + } else { + const uint2 p = *reinterpret_cast( + zxbcdt + (size_t)src * MB_IN_PROJ + MB_INNER + c); + v[0] = __uint_as_float(p.x << 16); + v[1] = __uint_as_float(p.x & 0xffff0000u); + v[2] = __uint_as_float(p.y << 16); + v[3] = __uint_as_float(p.y & 0xffff0000u); + } + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] += wv[q][k] * v[q]; + } + uint2 out; + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] = acc[q] / (1.f + __expf(-acc[q])); + out.x = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[0])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[1])) << 16); + out.y = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[2])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[3])) << 16); + *reinterpret_cast(xbc_act + (size_t)t * MB_CONV_DIM + c) = out; + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_dtda_prefill( + const MambaLayer &L, const bf16 *zxbcdt, float *dtda, int n_tokens) { + const int total = n_tokens * MB_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_HEADS; + const int head = idx - t * MB_HEADS; + const float dt = mamba_softplus( + __bfloat162float(zxbcdt[(size_t)t * MB_IN_PROJ + MB_INNER + + MB_CONV_DIM + head]) + L.dt_bias[head]); + dtda[(size_t)idx * 2] = dt; + dtda[(size_t)idx * 2 + 1] = __expf(L.A[head] * dt); + } +} +template +__device__ __forceinline__ void mamba_ssm_prefill_impl( + const MambaLayer &L, const bf16 *xbc_act, const float *dtda, bf16 *gated, + int n_pos, int n_seq, int slot, const int *seq_slot) { + constexpr int NS = MB_STATE / 32; + constexpr int PG = MB_HEAD_DIM / PW; + constexpr bool RS = (PW % 4) == 0; + constexpr int H1 = RS ? PW / 2 : 1; + constexpr int H2 = RS ? PW / 4 : 1; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const bool keep_hi = (lane & 16) != 0; + const bool keep_hi2 = (lane & 8) != 0; + const int warps_total = gridDim.x * WARPS; + const int units = n_seq * MB_HEADS * PG; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int pg = unit % PG; + const int rest = unit / PG; + const int head = rest % MB_HEADS; + const int seq = rest / MB_HEADS; + const int p0 = pg * PW; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Dh = L.D[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p0 * MB_STATE; + float sreg[PW][NS]; + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + sreg[q][j] = st[(size_t)q * MB_STATE + lane + j * 32]; + } + auto fetch = [&](int p, float (&oB)[NS], float (&oC)[NS], + float (&oX)[PW], float2 &od) { + const size_t tt = (size_t)p * n_seq + seq; + const bf16 *xa2 = xbc_act + tt * MB_CONV_DIM; + const bf16 *Bs = xa2 + MB_INNER + group * MB_STATE; + const bf16 *Cs = xa2 + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + od = *reinterpret_cast( + dtda + (tt * MB_HEADS + head) * 2); + #pragma unroll + for (int j = 0; j < NS; ++j) { + oB[j] = __bfloat162float(Bs[lane + j * 32]); + oC[j] = __bfloat162float(Cs[lane + j * 32]); + } + if constexpr (RS) { + const uint2 *xp = reinterpret_cast( + xa2 + head * MB_HEAD_DIM + p0); + #pragma unroll + for (int q = 0; q < PW; q += 4) { + const uint2 w = xp[q >> 2]; + oX[q] = __uint_as_float(w.x << 16); + oX[q + 1] = __uint_as_float(w.x & 0xffff0000u); + oX[q + 2] = __uint_as_float(w.y << 16); + oX[q + 3] = __uint_as_float(w.y & 0xffff0000u); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) + oX[q] = __bfloat162float(xa2[head * MB_HEAD_DIM + p0 + q]); + } + }; + float pB[NS], pC[NS], pX[PW]; + float2 pd; + fetch(0, pB, pC, pX, pd); + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const float dt = pd.x; + const float dA = pd.y; + float Bv[NS], Cv[NS]; + #pragma unroll + for (int j = 0; j < NS; ++j) { + Bv[j] = pB[j] * dt; + Cv[j] = pC[j]; + } + float acc[PW], xvs[PW]; + #pragma unroll + for (int q = 0; q < PW; ++q) xvs[q] = pX[q]; + if (pos + 1 < n_pos) fetch(pos + 1, pB, pC, pX, pd); + #pragma unroll + for (int q = 0; q < PW; ++q) { + const float xv = xvs[q]; + float a = 0.f; + #pragma unroll + for (int j = 0; j < NS; ++j) { + sreg[q][j] = sreg[q][j] * dA + Bv[j] * xv; + a += sreg[q][j] * Cv[j]; + } + acc[q] = a; + } + if constexpr (RS) { + float m1[H1], t1[H1], x1[H1]; + #pragma unroll + for (int q = 0; q < H1; ++q) { + m1[q] = keep_hi ? acc[q + H1] : acc[q]; + t1[q] = keep_hi ? acc[q] : acc[q + H1]; + x1[q] = keep_hi ? xvs[q + H1] : xvs[q]; + } + #pragma unroll + for (int q = 0; q < H1; ++q) + m1[q] += __shfl_xor_sync(0xffffffff, t1[q], 16); + float m2[H2], t2[H2], x2[H2]; + #pragma unroll + for (int q = 0; q < H2; ++q) { + m2[q] = keep_hi2 ? m1[q + H2] : m1[q]; + t2[q] = keep_hi2 ? m1[q] : m1[q + H2]; + x2[q] = keep_hi2 ? x1[q + H2] : x1[q]; + } + #pragma unroll + for (int q = 0; q < H2; ++q) + m2[q] += __shfl_xor_sync(0xffffffff, t2[q], 8); + #pragma unroll + for (int q = 0; q < H2; ++q) { + #pragma unroll + for (int off = 4; off > 0; off >>= 1) + m2[q] += __shfl_xor_sync(0xffffffff, m2[q], off); + } + if ((lane & 7) == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + + p0 + (keep_hi ? H1 : 0) + (keep_hi2 ? H2 : 0); + #pragma unroll + for (int q = 0; q < H2; ++q) + dst[q] = __float2bfloat16(m2[q] + x2[q] * Dh); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc[q] += __shfl_xor_sync(0xffffffff, acc[q], off); + } + if (lane == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + p0; + #pragma unroll + for (int q = 0; q < PW; ++q) + dst[q] = __float2bfloat16(acc[q] + xvs[q] * Dh); + } + } + } + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + st[(size_t)q * MB_STATE + lane + j * 32] = sreg[q][j]; + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, + const float *dtda, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + (void)zxbcdt; + if (n_seq > 1) + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens / n_seq, n_seq, slot, seq_slot); + else + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens, 1, slot, nullptr); +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16( + held[j] * inv * __bfloat162float(L.gnorm_w[idx])); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = BATCH * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / BB_ATT_HEADS; + const int head = pair % BB_ATT_HEADS; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const int kv = head / kv_per_q; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *q = base + head * BB_HEAD_DIM; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (head % kv_per_q == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + for (int tpos = 0; tpos < T; ++tpos) { +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new + : kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#endif + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m, s); + const float rescale = __expf(m - m_new); + const float w = __expf(s - m_new); + denom = denom * rescale + w; +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *vt = kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#else + const bf16 *vt = cur_token + ? v_new + : kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + w * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + constexpr int CHUNKS = BB_HEAD_DIM / 8; + const int total = n_tokens * BB_KV_HEADS * CHUNKS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int c = idx % CHUNKS; + const int rest = idx / CHUNKS; + const int t = rest / BB_KV_HEADS; + const int kv = rest - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM + c * 8; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM + c * 8; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv) + c * 8; + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv) + c * 8; + *reinterpret_cast(kc) = *reinterpret_cast(k_new); + *reinterpret_cast(vc) = *reinterpret_cast(v_new); + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float qv[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qv[j] = __bfloat162float(q[lane + j * 32]); + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + const int *const krow = + L.block_table + ((size_t)seq * 2 + 0) * L.max_blocks_per_seq; + const int *const vrow = + L.block_table + ((size_t)seq * 2 + 1) * L.max_blocks_per_seq; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + const int head_bias = kv * L.tokens_per_block * BB_HEAD_DIM; + int page = -1; + const bf16 *kpage = nullptr, *vpage = nullptr; + float kn[PER_LANE], vn[PER_LANE]; + { + page = 0; + kpage = L.pool_base + (size_t)(krow[0] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + (size_t)(vrow[0] & 0x7FFFFFFF) * elems_per_block; + const int at = head_bias + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + for (int s = 0; s <= pos; ++s) { + float kc[PER_LANE], vc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { kc[j] = kn[j]; vc[j] = vn[j]; } + if (s < pos) { + const int np = (s + 1) / L.tokens_per_block; + if (np != page) { + page = np; + kpage = L.pool_base + + (size_t)(krow[np] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + + (size_t)(vrow[np] & 0x7FFFFFFF) * elems_per_block; + } + const int at = head_bias + + ((s + 1) % L.tokens_per_block) * BB_HEAD_DIM + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) dot += qv[j] * kc[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[j] = acc[j] * rescale + wt * vc[j]; + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH, unsigned int *tile_masks = nullptr, + bf16 *gath = nullptr, int *gidx = nullptr, int *ecnt = nullptr, + int estride = 0, bf16 *row_stage = nullptr) { + __shared__ float score[BB_EXPERTS]; + __shared__ int s_pick[BB_TOP_K], s_slot[BB_TOP_K]; + constexpr int CH = D_MODEL / 8; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + { + const float4 *gsrc = + reinterpret_cast(normed + (size_t)row * D_MODEL); + float4 *gdst = reinterpret_cast(row_stage); + for (int i = threadIdx.x; i < CH; i += THREADS) gdst[i] = gsrc[i]; + } + if ((int)threadIdx.x < BB_EXPERTS) + moe_w[(size_t)row * BB_EXPERTS + threadIdx.x] = 0.f; + __syncthreads(); + for (int e = warp; e < BB_EXPERTS; e += WARPS) { + const bf16 *wr = gate_w + (size_t)e * D_MODEL; + float acc = 0.f; + for (int i = lane; i < D_MODEL; i += 32) + acc += __bfloat162float(row_stage[i]) * + __bfloat162float(wr[i]); + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) score[e] = 1.f / (1.f + __expf(-acc)); + } + __syncthreads(); + if (threadIdx.x == 0) { + float remaining[BB_EXPERTS]; + int pick[BB_TOP_K]; + for (int e = 0; e < BB_EXPERTS; ++e) + remaining[e] = score[e] + __bfloat162float(gate_bias[e]); + float total = 0.f; + for (int k = 0; k < BB_TOP_K; ++k) { + int best = 0; + for (int e = 1; e < BB_EXPERTS; ++e) + if (remaining[e] > remaining[best]) best = e; + pick[k] = best; + total += score[best]; + remaining[best] = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + unsigned int bits = 0u; + for (int k = 0; k < BB_TOP_K; ++k) { + moe_w[row * BB_EXPERTS + pick[k]] = score[pick[k]] * norm; + bits |= 1u << pick[k]; + s_pick[k] = pick[k]; + } + if (tile_masks) atomicOr(&tile_masks[row / BATCH_PAD], bits); + } + __syncthreads(); + if (gath) { + if ((int)threadIdx.x < BB_TOP_K) { + const int k = threadIdx.x; + const int slot = atomicAdd(&ecnt[s_pick[k]], 1); + s_slot[k] = slot; + gidx[(size_t)s_pick[k] * estride + slot] = row; + } + __syncthreads(); + const float4 *src = reinterpret_cast(row_stage); + for (int k = 0; k < BB_TOP_K; ++k) { + float4 *dst = reinterpret_cast( + gath + ((size_t)s_pick[k] * estride + s_slot[k]) * D_MODEL); + for (int i = threadIdx.x; i < CH; i += THREADS) dst[i] = src[i]; + } + } + __syncthreads(); + } +} +__device__ __forceinline__ unsigned moe_row_mask(const float *routes, int row) { + unsigned m = 0u; + const float *rw = routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) if (rw[e] != 0.f) m |= 1u << e; + return m; +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx, + unsigned mask) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + for (unsigned m = mask; m; m &= m - 1u) + value = __fadd_rn(value, parts[base + (size_t)(__ffs(m) - 1) * D_MODEL]); + return __float2bfloat16( + __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL])); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + uint64_t *keys, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int n = threadIdx.x; n < SORT_N; n += THREADS) { + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + keys[n] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + __syncthreads(); + for (int k = 2; k <= SORT_N; k <<= 1) { + if (k >= 64) __syncthreads(); + for (int j = k >> 1; j > 0; j >>= 1) { + for (int i = threadIdx.x; i < SORT_N; i += THREADS) { + const int partner = i ^ j; + if (partner > i) { + const bool descending = ((i & k) == 0); + if ((keys[i] < keys[partner]) == descending) { + const uint64_t t = keys[i]; + keys[i] = keys[partner]; + keys[partner] = t; + } + } + } + if (j >= 32) __syncthreads(); + else __syncwarp(); + } + } + __syncthreads(); + if (threadIdx.x == 0) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float v; + int n; + if (!eos_used && eos_v > av) { + v = eos_v; + n = AUDIO_EOS; + eos_used = true; + } else { + v = av; + n = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = + v + gumbel[((size_t)row * N_POS + pos) * TOP_K + j]; + if (noised > best_noised) { best_noised = noised; best_code = n; } + } + out_codes[row] = best_code; + } + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void prefill_expert_masks(const float *routes, unsigned int *masks, int rows) { + const int lane=threadIdx.x&31, warp=threadIdx.x>>5; + const int nt=(rows+BATCH_PAD-1)/BATCH_PAD; + for (int t=blockIdx.x*WARPS+warp;t(gath + g * D_MODEL + c * 8); + if (local < ecnt[e]) + *dst = *reinterpret_cast( + normed + (size_t)gidx[g] * D_MODEL + c * 8); + else + *dst = make_float4(0.f, 0.f, 0.f, 0.f); + } +} +__device__ __forceinline__ int gemm_eff(int tiles, int reps) { + if (NCONS <= 1) return 1; + const int g = (int)gridDim.x; + const int rounds1 = (tiles * reps + g - 1) / g; + if constexpr (NCONS == 3) { + const int max_eff = reps <= BB_EXPERTS ? 2 : 3; + const int rounds2 = (((tiles + 1) / 2) * reps + g - 1) / g; + if (max_eff == 2) return rounds2 < rounds1 ? 2 : 1; + const int rounds3 = (((tiles + 2) / 3) * reps + g - 1) / g; + if (rounds3 < rounds2 && rounds3 < rounds1) return 3; + return rounds2 < rounds1 ? 2 : 1; + } else { + const int roundsN = (((tiles + NCONS - 1) / NCONS) * reps + g - 1) / g; + return roundsN < rounds1 ? NCONS : 1; + } +} +__device__ __forceinline__ int gemm_groups(int tiles, int reps) { + const int eff = gemm_eff(tiles, reps); + return (tiles + eff - 1) / eff; +} +__device__ __forceinline__ int gemm_consumer_id() { + const int g = warpgroup::groupid(); + if (g == 1) return -1; + const int c = g == 0 ? 0 : g - 1; + return c < NCONS ? c : -1; +} +__device__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][BATCH_PAD]; + __shared__ int moe_gidx_stage[NCONS][BATCH_PAD]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + a_s[s], a, {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + a_s[s], a, {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / BATCH_PAD) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, BATCH_PAD> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], a_s[s]); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * BATCH_PAD; + if (epi == EPI_BIAS) { + warpgroup::store(o_sf, accum); + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + dest[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(o_sf[int2{col, row}] + + __bfloat162float(bias[base + col])); + } + warpgroup::sync(1 + cons); + continue; + } + warpgroup::store(o_s[cons], accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_s[cons][int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__device__ void gemm_narrow( + const weight_gl &w, const act_gl_n &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = PAD_N, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][PAD_N]; + __shared__ int moe_gidx_stage[NCONS][PAD_N]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + PAD_N - 1) / PAD_N; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + PAD_N - 1) / PAD_N; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / PAD_N) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, PAD_N> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], + *reinterpret_cast(&a_s[s])); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * PAD_N; + out_tile_n &o_sn = *reinterpret_cast(&o_s[cons]); + warpgroup::store(o_sn, accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_sn[int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&w_s)[STAGES * NCONS] = + alloc.allocate(); + act_tile (&a_s)[STAGES] = alloc.allocate(); + out_tile (&o_s)[NCONS] = alloc.allocate(); +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 + static_assert(NCONS * sizeof(out_tile) >= sizeof(out_tile_f), + "aliased head staging must fit the output tiles"); + out_tile_f &o_sf = *reinterpret_cast(&o_s[0]); +#else + out_tile_f (&o_sf) = alloc.allocate(); +#endif + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_WARPS + 2]; +#if MK2_PREFILL_ONLY + bf16 *const route_stage = reinterpret_cast(&o_s[0]); + static_assert(sizeof(o_s) >= D_MODEL * sizeof(bf16), + "router row stage must fit the output tiles"); +#else + __shared__ __align__(16) uint64_t sort_keys[SORT_N]; + bf16 *const route_stage = reinterpret_cast(&sort_keys[0]); + static_assert(sizeof(sort_keys) >= D_MODEL * sizeof(bf16), + "router row stage must fit the sampler workspace"); +#endif + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS * NCONS, 0); + for (int c = 0; c < CONSUMER_WARPS * NCONS; ++c) arrive(drained[s]); + } + } + int gemm_iter = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map); + grid_barrier(g.barrier); + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map); + grid_barrier(g.barrier); + mamba_gated_norm(L, bb->zxbcdt, bb->gated); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, nullptr, 1, BATCH, + bb->moe_w, BB_MOE_INTER / OUT_ATOM); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + 1, BATCH, bb->moe_w, D_MODEL / OUT_ATOM); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_deterministic(bb->moe_out, idx, + (1u << BB_EXPERTS) - 1u); + grid_barrier(g.barrier); + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int original_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + bool identical_cohort = false; + if constexpr (BATCH > 1) { + if (bb->native_pf_cohort_check && bb->native_prefill_reset && + bb->pf_n_seq == BATCH && !bb->pf_fuse_decode && + !bb->pf_layer_out && !bb->moe_w_inject && !bb->moe_w_audit) { + if (blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[0] = 0; + grid_barrier(g.barrier); + const auto *bits = reinterpret_cast(bb->prompt); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < original_total * D_MODEL; idx += gridDim.x * THREADS) { + const int canonical = (idx / (BATCH * D_MODEL)) * BATCH * D_MODEL + idx % D_MODEL; + if (bits[idx] != bits[canonical]) atomicOr(bb->native_pf_cohort_check, 1u); + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + if (bb->pf_seq_slot[row] != row || bb->slot_map[row] != row) + atomicOr(bb->native_pf_cohort_check, 1u); + grid_barrier(g.barrier); + identical_cohort = bb->native_pf_cohort_check[0] == 0; + } + } + const int n_total = identical_cohort ? original_total / BATCH : original_total; + if (bb->native_pf_cohort_check && blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[1] = identical_cohort ? 1 : BATCH; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int reset_slots = identical_cohort ? 1 : BATCH; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * reset_slots * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % reset_slots; + const MambaLayer &L = bb->mamba[row_layer / reset_slots]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * reset_slots * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % reset_slots; + const MambaLayer &L = bb->mamba[row_layer / reset_slots]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = original_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + const int src = identical_cohort ? (idx / D_MODEL) * BATCH * D_MODEL + idx % D_MODEL : idx; + bb->pf_h[idx] = bb->prompt[src]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = identical_cohort ? 1 : (bb->pf_n_seq > 0 ? bb->pf_n_seq : 1); + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T, + (!bb->pf_layer_out && layer > 0 && + bb->moe_of_layer[layer - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + if (bb->native_pf_expert_masks) + for (int i = blockIdx.x * THREADS + threadIdx.x; i < n_tok; + i += gridDim.x * THREADS) + bb->native_pf_expert_masks[i] = 0u; + if (bb->pf_moe_ecnt) + for (int e = blockIdx.x * THREADS + threadIdx.x; e < BB_EXPERTS; + e += gridDim.x * THREADS) + bb->pf_moe_ecnt[e] = 0; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + mamba_dtda_prefill(L, bb->pf_zxbcdt, bb->pf_dtda, T); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_dtda, + bb->pf_gated, T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, + T, bb->native_pf_expert_masks, bb->pf_moe_gath, + bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, route_stage); + } + grid_barrier(g.barrier); + int max_ecnt = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + max_ecnt = max(max_ecnt, bb->pf_moe_ecnt[e]); + const bool narrow = BATCH_PAD > PAD_N && max_ecnt <= PAD_N; + const int etile = narrow ? PAD_N : BATCH_PAD; + int nt_sum = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_sum += (bb->pf_moe_ecnt[e] + etile - 1) / etile; + const int up_steps = + gemm_groups(BB_MOE_INTER / OUT_ATOM, nt_sum) * + nt_sum * (D_MODEL / KCHUNK); + const int shup_steps = + gemm_groups(BB_SHARED_INTER / OUT_ATOM, n_tok) * + n_tok * (D_MODEL / KCHUNK); + const int down_steps = + gemm_groups(D_MODEL / OUT_ATOM, nt_sum) * + nt_sum * (BB_MOE_INTER / KCHUNK); + const int shdown_steps = + gemm_groups(D_MODEL / OUT_ATOM, n_tok) * + n_tok * (BB_SHARED_INTER / KCHUNK); + if (narrow) + gemm_narrow(L.exp_up, bb->pf_moe_gath_vn, + BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, PAD_N, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + else + gemm(L.exp_up, bb->pf_moe_gath_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, BATCH_PAD, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + up_steps, up_steps + shup_steps); + grid_barrier(g.barrier); + if (narrow) + gemm_narrow(L.exp_down, bb->pf_moe_hc_vn, + BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, PAD_N, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + else + gemm(L.exp_down, bb->pf_moe_hc_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + down_steps, down_steps + shdown_steps); + grid_barrier(g.barrier); + if (bb->pf_layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * D_MODEL; idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = + bb_moe_fold_deterministic(bb->pf_moe_out, idx, + moe_row_mask(bb->pf_moe_w, idx / D_MODEL)); + grid_barrier(g.barrier); + } + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, + identical_cohort ? bb->pf_normed : bb->pf_final, ln_scratch, + n_total, + (!bb->pf_layer_out && n_backbone_layers > 0 && + bb->moe_of_layer[n_backbone_layers - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + grid_barrier(g.barrier); + if (identical_cohort) { + const int worker = blockIdx.x * THREADS + threadIdx.x; + const int workers = gridDim.x * THREADS; + for (int i = worker; i < original_total * D_MODEL; i += workers) + bb->pf_final[i] = bb->pf_normed[(i / (BATCH * D_MODEL)) * D_MODEL + i % D_MODEL]; + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = worker; i < 14 * (BATCH - 1) * CV; i += workers) { + const int cell = i % CV, pair = i / CV; + const int row = pair % (BATCH - 1) + 1; + const MambaLayer &L = bb->mamba[pair / (BATCH - 1)]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = L.conv_state[cell]; + } + for (int i = worker; i < 14 * (BATCH - 1) * SS; i += workers) { + const int cell = i % SS, pair = i / SS; + const int row = pair % (BATCH - 1) + 1; + const MambaLayer &L = bb->mamba[pair / (BATCH - 1)]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = L.ssm_state[cell]; + } + const int kv_columns = BB_KV_HEADS * BB_HEAD_DIM; + for (int i = worker; i < 3 * 2 * (BATCH - 1) * n_total * kv_columns; i += workers) { + int q = i / kv_columns; + const int col = i % kv_columns; + const int pos = q % n_total; q /= n_total; + const int row = q % (BATCH - 1) + 1; q /= BATCH - 1; + const int kind = q % 2, layer = q / 2; + const AttnLayer &L = bb->attn[layer]; + const int offset = kv_offset(L, pos, col / BB_HEAD_DIM) + col % BB_HEAD_DIM; + kv_slot(L, row, kind, pos)[offset] = kv_slot(L, 0, kind, pos)[offset]; + } + grid_barrier(g.barrier); + } + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } +#if MK2_PREFILL_ONLY + return; +#else + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, bb->phoneme_b); + grid_barrier(g.barrier); + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + if (threadIdx.x == 0) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = 0; n < BB_PHONEME; ++n) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + bb->phoneme_out[row] = arg; + } + __syncthreads(); + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + for (int i = lane; i < D_HEAD; i += 32) { + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + for (int t = 0; t <= pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + for (int i = lane; i < D_HEAD; i += 32) { + float acc = 0.f; + for (int t = 0; t <= pos; ++t) + acc += sc[t] * __bfloat162float( + g.vcache[cache_base + t * D_HEAD + i]); + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + i] = + __float2bfloat16(acc); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, g.sampled_pos, + sort_keys, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { +#ifdef ABLATE_SAMPLE + __syncthreads(); +#endif + } else { + grid_barrier(g.barrier); + } + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x < BATCH && blockIdx.x == 0) { + const int code = g.sampled_pos[threadIdx.x]; + g.sampled[threadIdx.x * N_POS + pos] = code; + if (g.native_codes) + g.native_codes[threadIdx.x * N_POS + pos] = (long long)code; + } + grid_barrier(g.barrier); + } + if (g.native_seq_len && blockIdx.x == 0 && threadIdx.x < BATCH) + g.native_seq_len[threadIdx.x] += 1; + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +#endif +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 2 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK), + 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = MK2_PREFILL_ONLY + ? out->sms * out->blocks_per_sm + : std::min(out->sms * out->blocks_per_sm, D_FFN / OUT_ATOM); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + 0, + nullptr, + nullptr, + nullptr, nullptr, nullptr, nullptr, nullptr, 0, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "dedup_native_prefill_impl.inc" diff --git a/deploy/nemotron_native/kernels/prefill/b16/mk2_abi.cuh b/deploy/nemotron_native/kernels/prefill/b16/mk2_abi.cuh new file mode 100644 index 0000000..95a253b --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b16/mk2_abi.cuh @@ -0,0 +1,82 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} diff --git a/deploy/nemotron_native/kernels/prefill/b16/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/prefill/b16/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b16/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/prefill/b16/prefill.cu b/deploy/nemotron_native/kernels/prefill/b16/prefill.cu new file mode 100644 index 0000000..d2a7647 --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b16/prefill.cu @@ -0,0 +1,18 @@ +#define MK2_BATCH 16 +#define MK2_BATCH_TILE 64 +#define MK2_PREFILL_ONLY 1 +#define STAGES 3 +#define KCHUNK 128 +#define MK2_GEMM_CONSUMERS 3 +#define MK2_THREADS 512 +#define MK2_NO_MAIN 1 +#define MK2_DIRECT_INPUT 1 +#define MK2_B1_LOCAL_BARRIERS 1 +#define MK2_SKIP_HIDDEN_SNAPSHOT 1 +#define MK2_ACTIVE_EXPERTS 1 +#define MK2_SSM_SPLITS 1 +#define PF_SSM_PW_MULTI 12 +#define PF_ACTIVE_TOKEN_EXPERTS 1 +#include +#include +#include "dedup_predictor_mk2.cuh" diff --git a/deploy/nemotron_native/kernels/prefill/b2/dedup_native_prefill_impl.inc b/deploy/nemotron_native/kernels/prefill/b2/dedup_native_prefill_impl.inc new file mode 100644 index 0000000..f7e9802 --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b2/dedup_native_prefill_impl.inc @@ -0,0 +1,237 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 2 || BATCH == 4 || BATCH == 8 || BATCH == 16 || BATCH == 32 || BATCH == 64)) + throw std::runtime_error("native prefill requires a matching supported physical batch"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (padded > INT_MAX / ((BB_EXPERTS + 1) * D_MODEL)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + float *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(float))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(float))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_dtda, 2 * MB_HEADS); + { + const size_t estride = padded; + const size_t grows = size_t(BB_EXPERTS) * estride; + mk2_check(owned.alloc(&hb.pf_moe_gath, grows * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_gath, 0, grows * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_hc, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_hc, 0, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_gidx, grows * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_gidx, 0, grows * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ecnt, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ecnt, 0, BB_EXPERTS * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ntiles, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ntiles, 0, BB_EXPERTS * sizeof(int))); + hb.pf_moe_estride = int(estride); + hb.pf_moe_gath_v = act_gl{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_v = act_gl{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + hb.pf_moe_gath_vn = act_gl_n{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_vn = act_gl_n{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + scratch_bytes += grows * (D_MODEL + BB_MOE_INTER) * sizeof(bf16); + } + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(owned.alloc(&hb.native_pf_cohort_check, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(hb.native_pf_cohort_check, 0, 2 * sizeof(unsigned int))); +#if PF_ACTIVE_TOKEN_EXPERTS + mk2_check(owned.alloc(&hb.native_pf_expert_masks, (padded/BATCH_PAD)*sizeof(unsigned int))); +#endif + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} +extern "C" int mk2_native_prefill_effective_sequences(MK2NativePrefill *p, int *out) { + try { + if (!p || !out) throw std::runtime_error("null effective-sequence query"); + auto *st = reinterpret_cast(p->base); + mk2_check(cudaMemcpy(out, st->host_bb.native_pf_cohort_check + 1, + sizeof(int), cudaMemcpyDeviceToHost)); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} diff --git a/deploy/nemotron_native/kernels/prefill/b2/dedup_predictor_mk2.cuh b/deploy/nemotron_native/kernels/prefill/b2/dedup_predictor_mk2.cuh new file mode 100644 index 0000000..a640917 --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b2/dedup_predictor_mk2.cuh @@ -0,0 +1,4588 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef PF_ACTIVE_TOKEN_EXPERTS +#define PF_ACTIVE_TOKEN_EXPERTS 0 +#endif +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +#ifndef MK2_PREFILL_ONLY +#define MK2_PREFILL_ONLY 0 +#endif +#ifndef MK2_BATCH_TILE +#define MK2_BATCH_TILE (((MK2_BATCH + 15) / 16) * 16) +#endif +constexpr int BATCH_PAD = MK2_BATCH_TILE; +static_assert(BATCH >= 1 && BATCH <= BATCH_PAD); +static_assert(BATCH_PAD == 16 || BATCH_PAD == 32 || BATCH_PAD == 64 || BATCH_PAD == 128, + "native prefill tile must be 16, 32, 64 or 128"); +static_assert(MK2_PREFILL_ONLY || BATCH_PAD == ((BATCH + 15) / 16) * 16, + "a wider token tile requires a prefill-only library"); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_PW_MULTI +#define PF_SSM_PW_MULTI 4 +#endif +#ifndef PF_SSM_PW_SINGLE +#define PF_SSM_PW_SINGLE 2 +#endif +static_assert(MB_HEAD_DIM % PF_SSM_PW_MULTI == 0 && + MB_HEAD_DIM % PF_SSM_PW_SINGLE == 0, + "head_dim must split evenly into scan channel groups"); +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef MK2_SMEM_SLACK +#define MK2_SMEM_SLACK 1024 +#endif +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +#ifndef MK2_THREADS +#define MK2_THREADS 256 +#endif +constexpr int THREADS = MK2_THREADS; +constexpr int NORM_THREADS = 256; +constexpr int NORM_WARPS = NORM_THREADS / 32; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +#ifndef MK2_GEMM_CONSUMERS +#define MK2_GEMM_CONSUMERS 1 +#endif +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 +#define MK2_HEAD_STAGE_BYTES 0 +#else +#define MK2_HEAD_STAGE_BYTES sizeof(out_tile_f) +#endif +constexpr int NCONS = MK2_GEMM_CONSUMERS; +static_assert(NCONS >= 1 && NCONS <= 3, "consumers occupy warpgroups 0, 2 and 3"); +static_assert(THREADS >= (NCONS + 1) * 128, "one producer plus NCONS consumers"); +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +constexpr int PAD_N = 32; +using act_tile_n = st_bf; +using out_tile_n = st_bf; +using act_gl_n = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_MOE_ACC_C, EPI_ACC_F32 }; +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + unsigned int *native_pf_expert_masks = nullptr; + float *pf_dtda = nullptr; + bf16 *pf_moe_gath = nullptr; + bf16 *pf_moe_hc = nullptr; + int *pf_moe_gidx = nullptr; + int *pf_moe_ecnt = nullptr; + int *pf_moe_ntiles = nullptr; + int pf_moe_estride = 0; + act_gl pf_moe_gath_v; + act_gl pf_moe_hc_v; + act_gl_n pf_moe_gath_vn; + act_gl_n pf_moe_hc_vn; + unsigned int *native_pf_cohort_check = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + volatile unsigned int *generation = state + 1; + const unsigned int seen = *generation; + if (atomicAdd(state, 1u) == gridDim.x - 1) { + atomicExch(state, 0u); + atomicAdd(state + 1, 1u); + } else { + while (*generation == seen) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const bf16 *in = src + (size_t)row * D_MODEL; + float v[PER_THREAD]; + float sum = 0.f, sq = 0.f; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (threadIdx.x < NORM_THREADS) { + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + v[j] = __bfloat162float(in[threadIdx.x + j * NORM_THREADS]); + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + if (lane == 0) { scratch[warp] = sum; scratch[NORM_WARPS + warp] = sq; } + } + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f, q = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) { s += scratch[w]; q += scratch[NORM_WARPS + w]; } + scratch[2 * NORM_WARPS] = s / D_MODEL; + scratch[2 * NORM_WARPS + 1] = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + } + __syncthreads(); + const float mean = scratch[2 * NORM_WARPS], inv = scratch[2 * NORM_WARPS + 1]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + if (threadIdx.x >= NORM_THREADS) break; + const int i = threadIdx.x + j * NORM_THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * weight[i]); + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, + int idx, unsigned mask); +__device__ __forceinline__ void norm_sub_sync(int sub) { + if (THREADS == NORM_THREADS) { __syncthreads(); return; } + asm volatile("bar.sync %0, %1;" :: "r"(6 + sub), "n"(NORM_THREADS) : "memory"); +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, const float *fold_src = nullptr, + const float *fold_routes = nullptr) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + constexpr int NSUB = THREADS / NORM_THREADS; + if (threadIdx.x >= NSUB * NORM_THREADS) return; + const int sub = threadIdx.x / NORM_THREADS; + const int tid = threadIdx.x - sub * NORM_THREADS; + float *const sc = scratch + sub * (NORM_WARPS + 1); + for (int row = blockIdx.x * NSUB + sub; row < n_rows; + row += gridDim.x * NSUB) { + const size_t base = (size_t)row * D_MODEL; + unsigned fold_mask = (1u << BB_EXPERTS) - 1u; + if (fold_src && fold_routes) { + fold_mask = 0u; + const float *rw = fold_routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) + if (rw[e] != 0.f) fold_mask |= 1u << e; + } + float v[PER_THREAD]; + float sq = 0.f; + const int warp = tid >> 5, lane = tid & 31; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + bf16 hv; + if (fold_src) { + hv = bb_moe_fold_deterministic(fold_src, (int)(base + i), fold_mask); + h[base + i] = hv; + } else { + hv = h[base + i]; + } + v[j] = __bfloat162float(hv) + __bfloat162float(r[base + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + if (lane == 0) sc[warp] = sq; + norm_sub_sync(sub); + if (tid == 0) { + float s = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) s += sc[w]; + sc[NORM_WARPS] = rsqrtf(s / D_MODEL + RMS_EPS); + } + norm_sub_sync(sub); + const float inv = sc[NORM_WARPS]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + out[base + i] = __float2bfloat16( + __bfloat162float(weight[i]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + } + norm_sub_sync(sub); + } +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + const int c = idx - row * MB_CONV_DIM; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + const float raw = __bfloat162float( + zxbcdt[(size_t)row * MB_IN_PROJ + MB_INNER + c]); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + xbc_act[(size_t)row * MB_CONV_DIM + c] = __float2bfloat16(acc); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + const int pairs = BATCH * MB_HEADS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_HEADS * split); + const int head = (pair / split) % MB_HEADS; + const int first_channel = (pair % split) * channels; + const int group = head / MB_HEADS_PER_GROUP; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + L.dt_bias[head]); + const float dA = __expf(L.A[head] * dt); + const float Dh = L.D[head]; + float *st = L.ssm_state + (size_t)slot_map[row] * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + for (int p = first_channel; p < first_channel + channels; ++p) { + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + const size_t at = (size_t)p * MB_STATE + n; + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + static_assert(MB_CONV_K == 4, "packed conv taps assume a 4-tap kernel"); + constexpr int CPT = 4; + static_assert(MB_CONV_DIM % CPT == 0, "channel vector width"); + constexpr int CGROUPS = MB_CONV_DIM / CPT; + const int total = n_tokens * CGROUPS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / CGROUPS; + const int c = (idx - t * CGROUPS) * CPT; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + float wv[CPT][MB_CONV_K]; + { + const float4 *wsrc = reinterpret_cast( + L.conv_w + (size_t)c * MB_CONV_K); + #pragma unroll + for (int j = 0; j < CPT / 2; ++j) { + const float4 w4 = wsrc[j]; + const unsigned uu0 = __float_as_uint(w4.x); + const unsigned uu1 = __float_as_uint(w4.y); + const unsigned uu2 = __float_as_uint(w4.z); + const unsigned uu3 = __float_as_uint(w4.w); + wv[2 * j][0] = __uint_as_float(uu0 << 16); + wv[2 * j][1] = __uint_as_float(uu0 & 0xffff0000u); + wv[2 * j][2] = __uint_as_float(uu1 << 16); + wv[2 * j][3] = __uint_as_float(uu1 & 0xffff0000u); + wv[2 * j + 1][0] = __uint_as_float(uu2 << 16); + wv[2 * j + 1][1] = __uint_as_float(uu2 & 0xffff0000u); + wv[2 * j + 1][2] = __uint_as_float(uu3 << 16); + wv[2 * j + 1][3] = __uint_as_float(uu3 & 0xffff0000u); + } + } + float acc[CPT]; + { + const uint2 b2 = *reinterpret_cast(L.conv_b + c); + acc[0] = __uint_as_float(b2.x << 16); + acc[1] = __uint_as_float(b2.x & 0xffff0000u); + acc[2] = __uint_as_float(b2.y << 16); + acc[3] = __uint_as_float(b2.y & 0xffff0000u); + } + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v[CPT]; + if (src < 0) { + #pragma unroll + for (int q = 0; q < CPT; ++q) + v[q] = __bfloat162float( + state[(size_t)q * (MB_CONV_K - 1) + pos + k]); + } else { + const uint2 p = *reinterpret_cast( + zxbcdt + (size_t)src * MB_IN_PROJ + MB_INNER + c); + v[0] = __uint_as_float(p.x << 16); + v[1] = __uint_as_float(p.x & 0xffff0000u); + v[2] = __uint_as_float(p.y << 16); + v[3] = __uint_as_float(p.y & 0xffff0000u); + } + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] += wv[q][k] * v[q]; + } + uint2 out; + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] = acc[q] / (1.f + __expf(-acc[q])); + out.x = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[0])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[1])) << 16); + out.y = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[2])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[3])) << 16); + *reinterpret_cast(xbc_act + (size_t)t * MB_CONV_DIM + c) = out; + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_dtda_prefill( + const MambaLayer &L, const bf16 *zxbcdt, float *dtda, int n_tokens) { + const int total = n_tokens * MB_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_HEADS; + const int head = idx - t * MB_HEADS; + const float dt = mamba_softplus( + __bfloat162float(zxbcdt[(size_t)t * MB_IN_PROJ + MB_INNER + + MB_CONV_DIM + head]) + L.dt_bias[head]); + dtda[(size_t)idx * 2] = dt; + dtda[(size_t)idx * 2 + 1] = __expf(L.A[head] * dt); + } +} +template +__device__ __forceinline__ void mamba_ssm_prefill_impl( + const MambaLayer &L, const bf16 *xbc_act, const float *dtda, bf16 *gated, + int n_pos, int n_seq, int slot, const int *seq_slot) { + constexpr int NS = MB_STATE / 32; + constexpr int PG = MB_HEAD_DIM / PW; + constexpr bool RS = (PW % 4) == 0; + constexpr int H1 = RS ? PW / 2 : 1; + constexpr int H2 = RS ? PW / 4 : 1; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const bool keep_hi = (lane & 16) != 0; + const bool keep_hi2 = (lane & 8) != 0; + const int warps_total = gridDim.x * WARPS; + const int units = n_seq * MB_HEADS * PG; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int pg = unit % PG; + const int rest = unit / PG; + const int head = rest % MB_HEADS; + const int seq = rest / MB_HEADS; + const int p0 = pg * PW; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Dh = L.D[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p0 * MB_STATE; + float sreg[PW][NS]; + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + sreg[q][j] = st[(size_t)q * MB_STATE + lane + j * 32]; + } + auto fetch = [&](int p, float (&oB)[NS], float (&oC)[NS], + float (&oX)[PW], float2 &od) { + const size_t tt = (size_t)p * n_seq + seq; + const bf16 *xa2 = xbc_act + tt * MB_CONV_DIM; + const bf16 *Bs = xa2 + MB_INNER + group * MB_STATE; + const bf16 *Cs = xa2 + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + od = *reinterpret_cast( + dtda + (tt * MB_HEADS + head) * 2); + #pragma unroll + for (int j = 0; j < NS; ++j) { + oB[j] = __bfloat162float(Bs[lane + j * 32]); + oC[j] = __bfloat162float(Cs[lane + j * 32]); + } + if constexpr (RS) { + const uint2 *xp = reinterpret_cast( + xa2 + head * MB_HEAD_DIM + p0); + #pragma unroll + for (int q = 0; q < PW; q += 4) { + const uint2 w = xp[q >> 2]; + oX[q] = __uint_as_float(w.x << 16); + oX[q + 1] = __uint_as_float(w.x & 0xffff0000u); + oX[q + 2] = __uint_as_float(w.y << 16); + oX[q + 3] = __uint_as_float(w.y & 0xffff0000u); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) + oX[q] = __bfloat162float(xa2[head * MB_HEAD_DIM + p0 + q]); + } + }; + float pB[NS], pC[NS], pX[PW]; + float2 pd; + fetch(0, pB, pC, pX, pd); + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const float dt = pd.x; + const float dA = pd.y; + float Bv[NS], Cv[NS]; + #pragma unroll + for (int j = 0; j < NS; ++j) { + Bv[j] = pB[j] * dt; + Cv[j] = pC[j]; + } + float acc[PW], xvs[PW]; + #pragma unroll + for (int q = 0; q < PW; ++q) xvs[q] = pX[q]; + if (pos + 1 < n_pos) fetch(pos + 1, pB, pC, pX, pd); + #pragma unroll + for (int q = 0; q < PW; ++q) { + const float xv = xvs[q]; + float a = 0.f; + #pragma unroll + for (int j = 0; j < NS; ++j) { + sreg[q][j] = sreg[q][j] * dA + Bv[j] * xv; + a += sreg[q][j] * Cv[j]; + } + acc[q] = a; + } + if constexpr (RS) { + float m1[H1], t1[H1], x1[H1]; + #pragma unroll + for (int q = 0; q < H1; ++q) { + m1[q] = keep_hi ? acc[q + H1] : acc[q]; + t1[q] = keep_hi ? acc[q] : acc[q + H1]; + x1[q] = keep_hi ? xvs[q + H1] : xvs[q]; + } + #pragma unroll + for (int q = 0; q < H1; ++q) + m1[q] += __shfl_xor_sync(0xffffffff, t1[q], 16); + float m2[H2], t2[H2], x2[H2]; + #pragma unroll + for (int q = 0; q < H2; ++q) { + m2[q] = keep_hi2 ? m1[q + H2] : m1[q]; + t2[q] = keep_hi2 ? m1[q] : m1[q + H2]; + x2[q] = keep_hi2 ? x1[q + H2] : x1[q]; + } + #pragma unroll + for (int q = 0; q < H2; ++q) + m2[q] += __shfl_xor_sync(0xffffffff, t2[q], 8); + #pragma unroll + for (int q = 0; q < H2; ++q) { + #pragma unroll + for (int off = 4; off > 0; off >>= 1) + m2[q] += __shfl_xor_sync(0xffffffff, m2[q], off); + } + if ((lane & 7) == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + + p0 + (keep_hi ? H1 : 0) + (keep_hi2 ? H2 : 0); + #pragma unroll + for (int q = 0; q < H2; ++q) + dst[q] = __float2bfloat16(m2[q] + x2[q] * Dh); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc[q] += __shfl_xor_sync(0xffffffff, acc[q], off); + } + if (lane == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + p0; + #pragma unroll + for (int q = 0; q < PW; ++q) + dst[q] = __float2bfloat16(acc[q] + xvs[q] * Dh); + } + } + } + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + st[(size_t)q * MB_STATE + lane + j * 32] = sreg[q][j]; + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, + const float *dtda, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + (void)zxbcdt; + if (n_seq > 1) + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens / n_seq, n_seq, slot, seq_slot); + else + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens, 1, slot, nullptr); +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16( + held[j] * inv * __bfloat162float(L.gnorm_w[idx])); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = BATCH * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / BB_ATT_HEADS; + const int head = pair % BB_ATT_HEADS; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const int kv = head / kv_per_q; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *q = base + head * BB_HEAD_DIM; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (head % kv_per_q == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + for (int tpos = 0; tpos < T; ++tpos) { +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new + : kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#endif + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m, s); + const float rescale = __expf(m - m_new); + const float w = __expf(s - m_new); + denom = denom * rescale + w; +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *vt = kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#else + const bf16 *vt = cur_token + ? v_new + : kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + w * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + constexpr int CHUNKS = BB_HEAD_DIM / 8; + const int total = n_tokens * BB_KV_HEADS * CHUNKS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int c = idx % CHUNKS; + const int rest = idx / CHUNKS; + const int t = rest / BB_KV_HEADS; + const int kv = rest - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM + c * 8; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM + c * 8; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv) + c * 8; + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv) + c * 8; + *reinterpret_cast(kc) = *reinterpret_cast(k_new); + *reinterpret_cast(vc) = *reinterpret_cast(v_new); + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float qv[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qv[j] = __bfloat162float(q[lane + j * 32]); + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + const int *const krow = + L.block_table + ((size_t)seq * 2 + 0) * L.max_blocks_per_seq; + const int *const vrow = + L.block_table + ((size_t)seq * 2 + 1) * L.max_blocks_per_seq; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + const int head_bias = kv * L.tokens_per_block * BB_HEAD_DIM; + int page = -1; + const bf16 *kpage = nullptr, *vpage = nullptr; + float kn[PER_LANE], vn[PER_LANE]; + { + page = 0; + kpage = L.pool_base + (size_t)(krow[0] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + (size_t)(vrow[0] & 0x7FFFFFFF) * elems_per_block; + const int at = head_bias + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + for (int s = 0; s <= pos; ++s) { + float kc[PER_LANE], vc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { kc[j] = kn[j]; vc[j] = vn[j]; } + if (s < pos) { + const int np = (s + 1) / L.tokens_per_block; + if (np != page) { + page = np; + kpage = L.pool_base + + (size_t)(krow[np] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + + (size_t)(vrow[np] & 0x7FFFFFFF) * elems_per_block; + } + const int at = head_bias + + ((s + 1) % L.tokens_per_block) * BB_HEAD_DIM + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) dot += qv[j] * kc[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[j] = acc[j] * rescale + wt * vc[j]; + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH, unsigned int *tile_masks = nullptr, + bf16 *gath = nullptr, int *gidx = nullptr, int *ecnt = nullptr, + int estride = 0, bf16 *row_stage = nullptr) { + __shared__ float score[BB_EXPERTS]; + __shared__ int s_pick[BB_TOP_K], s_slot[BB_TOP_K]; + constexpr int CH = D_MODEL / 8; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + { + const float4 *gsrc = + reinterpret_cast(normed + (size_t)row * D_MODEL); + float4 *gdst = reinterpret_cast(row_stage); + for (int i = threadIdx.x; i < CH; i += THREADS) gdst[i] = gsrc[i]; + } + if ((int)threadIdx.x < BB_EXPERTS) + moe_w[(size_t)row * BB_EXPERTS + threadIdx.x] = 0.f; + __syncthreads(); + for (int e = warp; e < BB_EXPERTS; e += WARPS) { + const bf16 *wr = gate_w + (size_t)e * D_MODEL; + float acc = 0.f; + for (int i = lane; i < D_MODEL; i += 32) + acc += __bfloat162float(row_stage[i]) * + __bfloat162float(wr[i]); + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) score[e] = 1.f / (1.f + __expf(-acc)); + } + __syncthreads(); + if (threadIdx.x == 0) { + float remaining[BB_EXPERTS]; + int pick[BB_TOP_K]; + for (int e = 0; e < BB_EXPERTS; ++e) + remaining[e] = score[e] + __bfloat162float(gate_bias[e]); + float total = 0.f; + for (int k = 0; k < BB_TOP_K; ++k) { + int best = 0; + for (int e = 1; e < BB_EXPERTS; ++e) + if (remaining[e] > remaining[best]) best = e; + pick[k] = best; + total += score[best]; + remaining[best] = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + unsigned int bits = 0u; + for (int k = 0; k < BB_TOP_K; ++k) { + moe_w[row * BB_EXPERTS + pick[k]] = score[pick[k]] * norm; + bits |= 1u << pick[k]; + s_pick[k] = pick[k]; + } + if (tile_masks) atomicOr(&tile_masks[row / BATCH_PAD], bits); + } + __syncthreads(); + if (gath) { + if ((int)threadIdx.x < BB_TOP_K) { + const int k = threadIdx.x; + const int slot = atomicAdd(&ecnt[s_pick[k]], 1); + s_slot[k] = slot; + gidx[(size_t)s_pick[k] * estride + slot] = row; + } + __syncthreads(); + const float4 *src = reinterpret_cast(row_stage); + for (int k = 0; k < BB_TOP_K; ++k) { + float4 *dst = reinterpret_cast( + gath + ((size_t)s_pick[k] * estride + s_slot[k]) * D_MODEL); + for (int i = threadIdx.x; i < CH; i += THREADS) dst[i] = src[i]; + } + } + __syncthreads(); + } +} +__device__ __forceinline__ unsigned moe_row_mask(const float *routes, int row) { + unsigned m = 0u; + const float *rw = routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) if (rw[e] != 0.f) m |= 1u << e; + return m; +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx, + unsigned mask) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + for (unsigned m = mask; m; m &= m - 1u) + value = __fadd_rn(value, parts[base + (size_t)(__ffs(m) - 1) * D_MODEL]); + return __float2bfloat16( + __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL])); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + uint64_t *keys, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int n = threadIdx.x; n < SORT_N; n += THREADS) { + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + keys[n] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + __syncthreads(); + for (int k = 2; k <= SORT_N; k <<= 1) { + if (k >= 64) __syncthreads(); + for (int j = k >> 1; j > 0; j >>= 1) { + for (int i = threadIdx.x; i < SORT_N; i += THREADS) { + const int partner = i ^ j; + if (partner > i) { + const bool descending = ((i & k) == 0); + if ((keys[i] < keys[partner]) == descending) { + const uint64_t t = keys[i]; + keys[i] = keys[partner]; + keys[partner] = t; + } + } + } + if (j >= 32) __syncthreads(); + else __syncwarp(); + } + } + __syncthreads(); + if (threadIdx.x == 0) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float v; + int n; + if (!eos_used && eos_v > av) { + v = eos_v; + n = AUDIO_EOS; + eos_used = true; + } else { + v = av; + n = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = + v + gumbel[((size_t)row * N_POS + pos) * TOP_K + j]; + if (noised > best_noised) { best_noised = noised; best_code = n; } + } + out_codes[row] = best_code; + } + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void prefill_expert_masks(const float *routes, unsigned int *masks, int rows) { + const int lane=threadIdx.x&31, warp=threadIdx.x>>5; + const int nt=(rows+BATCH_PAD-1)/BATCH_PAD; + for (int t=blockIdx.x*WARPS+warp;t(gath + g * D_MODEL + c * 8); + if (local < ecnt[e]) + *dst = *reinterpret_cast( + normed + (size_t)gidx[g] * D_MODEL + c * 8); + else + *dst = make_float4(0.f, 0.f, 0.f, 0.f); + } +} +__device__ __forceinline__ int gemm_eff(int tiles, int reps) { + if (NCONS <= 1) return 1; + const int g = (int)gridDim.x; + const int rounds1 = (tiles * reps + g - 1) / g; + if constexpr (NCONS == 3) { + const int max_eff = reps <= BB_EXPERTS ? 2 : 3; + const int rounds2 = (((tiles + 1) / 2) * reps + g - 1) / g; + if (max_eff == 2) return rounds2 < rounds1 ? 2 : 1; + const int rounds3 = (((tiles + 2) / 3) * reps + g - 1) / g; + if (rounds3 < rounds2 && rounds3 < rounds1) return 3; + return rounds2 < rounds1 ? 2 : 1; + } else { + const int roundsN = (((tiles + NCONS - 1) / NCONS) * reps + g - 1) / g; + return roundsN < rounds1 ? NCONS : 1; + } +} +__device__ __forceinline__ int gemm_groups(int tiles, int reps) { + const int eff = gemm_eff(tiles, reps); + return (tiles + eff - 1) / eff; +} +__device__ __forceinline__ int gemm_consumer_id() { + const int g = warpgroup::groupid(); + if (g == 1) return -1; + const int c = g == 0 ? 0 : g - 1; + return c < NCONS ? c : -1; +} +__device__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][BATCH_PAD]; + __shared__ int moe_gidx_stage[NCONS][BATCH_PAD]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + a_s[s], a, {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + a_s[s], a, {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / BATCH_PAD) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, BATCH_PAD> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], a_s[s]); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * BATCH_PAD; + if (epi == EPI_BIAS) { + warpgroup::store(o_sf, accum); + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + dest[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(o_sf[int2{col, row}] + + __bfloat162float(bias[base + col])); + } + warpgroup::sync(1 + cons); + continue; + } + warpgroup::store(o_s[cons], accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_s[cons][int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__device__ void gemm_narrow( + const weight_gl &w, const act_gl_n &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = PAD_N, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][PAD_N]; + __shared__ int moe_gidx_stage[NCONS][PAD_N]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + PAD_N - 1) / PAD_N; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + PAD_N - 1) / PAD_N; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / PAD_N) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, PAD_N> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], + *reinterpret_cast(&a_s[s])); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * PAD_N; + out_tile_n &o_sn = *reinterpret_cast(&o_s[cons]); + warpgroup::store(o_sn, accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_sn[int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&w_s)[STAGES * NCONS] = + alloc.allocate(); + act_tile (&a_s)[STAGES] = alloc.allocate(); + out_tile (&o_s)[NCONS] = alloc.allocate(); +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 + static_assert(NCONS * sizeof(out_tile) >= sizeof(out_tile_f), + "aliased head staging must fit the output tiles"); + out_tile_f &o_sf = *reinterpret_cast(&o_s[0]); +#else + out_tile_f (&o_sf) = alloc.allocate(); +#endif + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_WARPS + 2]; +#if MK2_PREFILL_ONLY + bf16 *const route_stage = reinterpret_cast(&o_s[0]); + static_assert(sizeof(o_s) >= D_MODEL * sizeof(bf16), + "router row stage must fit the output tiles"); +#else + __shared__ __align__(16) uint64_t sort_keys[SORT_N]; + bf16 *const route_stage = reinterpret_cast(&sort_keys[0]); + static_assert(sizeof(sort_keys) >= D_MODEL * sizeof(bf16), + "router row stage must fit the sampler workspace"); +#endif + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS * NCONS, 0); + for (int c = 0; c < CONSUMER_WARPS * NCONS; ++c) arrive(drained[s]); + } + } + int gemm_iter = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map); + grid_barrier(g.barrier); + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map); + grid_barrier(g.barrier); + mamba_gated_norm(L, bb->zxbcdt, bb->gated); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, nullptr, 1, BATCH, + bb->moe_w, BB_MOE_INTER / OUT_ATOM); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + 1, BATCH, bb->moe_w, D_MODEL / OUT_ATOM); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_deterministic(bb->moe_out, idx, + (1u << BB_EXPERTS) - 1u); + grid_barrier(g.barrier); + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int original_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + bool identical_cohort = false; + if constexpr (BATCH > 1) { + if (bb->native_pf_cohort_check && bb->native_prefill_reset && + bb->pf_n_seq == BATCH && !bb->pf_fuse_decode && + !bb->pf_layer_out && !bb->moe_w_inject && !bb->moe_w_audit) { + if (blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[0] = 0; + grid_barrier(g.barrier); + const auto *bits = reinterpret_cast(bb->prompt); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < original_total * D_MODEL; idx += gridDim.x * THREADS) { + const int canonical = (idx / (BATCH * D_MODEL)) * BATCH * D_MODEL + idx % D_MODEL; + if (bits[idx] != bits[canonical]) atomicOr(bb->native_pf_cohort_check, 1u); + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + if (bb->pf_seq_slot[row] != row || bb->slot_map[row] != row) + atomicOr(bb->native_pf_cohort_check, 1u); + grid_barrier(g.barrier); + identical_cohort = bb->native_pf_cohort_check[0] == 0; + } + } + const int n_total = identical_cohort ? original_total / BATCH : original_total; + if (bb->native_pf_cohort_check && blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[1] = identical_cohort ? 1 : BATCH; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int reset_slots = identical_cohort ? 1 : BATCH; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * reset_slots * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % reset_slots; + const MambaLayer &L = bb->mamba[row_layer / reset_slots]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * reset_slots * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % reset_slots; + const MambaLayer &L = bb->mamba[row_layer / reset_slots]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = original_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + const int src = identical_cohort ? (idx / D_MODEL) * BATCH * D_MODEL + idx % D_MODEL : idx; + bb->pf_h[idx] = bb->prompt[src]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = identical_cohort ? 1 : (bb->pf_n_seq > 0 ? bb->pf_n_seq : 1); + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T, + (!bb->pf_layer_out && layer > 0 && + bb->moe_of_layer[layer - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + if (bb->native_pf_expert_masks) + for (int i = blockIdx.x * THREADS + threadIdx.x; i < n_tok; + i += gridDim.x * THREADS) + bb->native_pf_expert_masks[i] = 0u; + if (bb->pf_moe_ecnt) + for (int e = blockIdx.x * THREADS + threadIdx.x; e < BB_EXPERTS; + e += gridDim.x * THREADS) + bb->pf_moe_ecnt[e] = 0; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + mamba_dtda_prefill(L, bb->pf_zxbcdt, bb->pf_dtda, T); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_dtda, + bb->pf_gated, T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, + T, bb->native_pf_expert_masks, bb->pf_moe_gath, + bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, route_stage); + } + grid_barrier(g.barrier); + int max_ecnt = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + max_ecnt = max(max_ecnt, bb->pf_moe_ecnt[e]); + const bool narrow = BATCH_PAD > PAD_N && max_ecnt <= PAD_N; + const int etile = narrow ? PAD_N : BATCH_PAD; + int nt_sum = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_sum += (bb->pf_moe_ecnt[e] + etile - 1) / etile; + const int up_steps = + gemm_groups(BB_MOE_INTER / OUT_ATOM, nt_sum) * + nt_sum * (D_MODEL / KCHUNK); + const int shup_steps = + gemm_groups(BB_SHARED_INTER / OUT_ATOM, n_tok) * + n_tok * (D_MODEL / KCHUNK); + const int down_steps = + gemm_groups(D_MODEL / OUT_ATOM, nt_sum) * + nt_sum * (BB_MOE_INTER / KCHUNK); + const int shdown_steps = + gemm_groups(D_MODEL / OUT_ATOM, n_tok) * + n_tok * (BB_SHARED_INTER / KCHUNK); + if (narrow) + gemm_narrow(L.exp_up, bb->pf_moe_gath_vn, + BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, PAD_N, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + else + gemm(L.exp_up, bb->pf_moe_gath_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, BATCH_PAD, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + up_steps, up_steps + shup_steps); + grid_barrier(g.barrier); + if (narrow) + gemm_narrow(L.exp_down, bb->pf_moe_hc_vn, + BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, PAD_N, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + else + gemm(L.exp_down, bb->pf_moe_hc_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + down_steps, down_steps + shdown_steps); + grid_barrier(g.barrier); + if (bb->pf_layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * D_MODEL; idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = + bb_moe_fold_deterministic(bb->pf_moe_out, idx, + moe_row_mask(bb->pf_moe_w, idx / D_MODEL)); + grid_barrier(g.barrier); + } + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, + identical_cohort ? bb->pf_normed : bb->pf_final, ln_scratch, + n_total, + (!bb->pf_layer_out && n_backbone_layers > 0 && + bb->moe_of_layer[n_backbone_layers - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + grid_barrier(g.barrier); + if (identical_cohort) { + const int worker = blockIdx.x * THREADS + threadIdx.x; + const int workers = gridDim.x * THREADS; + for (int i = worker; i < original_total * D_MODEL; i += workers) + bb->pf_final[i] = bb->pf_normed[(i / (BATCH * D_MODEL)) * D_MODEL + i % D_MODEL]; + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = worker; i < 14 * (BATCH - 1) * CV; i += workers) { + const int cell = i % CV, pair = i / CV; + const int row = pair % (BATCH - 1) + 1; + const MambaLayer &L = bb->mamba[pair / (BATCH - 1)]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = L.conv_state[cell]; + } + for (int i = worker; i < 14 * (BATCH - 1) * SS; i += workers) { + const int cell = i % SS, pair = i / SS; + const int row = pair % (BATCH - 1) + 1; + const MambaLayer &L = bb->mamba[pair / (BATCH - 1)]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = L.ssm_state[cell]; + } + const int kv_columns = BB_KV_HEADS * BB_HEAD_DIM; + for (int i = worker; i < 3 * 2 * (BATCH - 1) * n_total * kv_columns; i += workers) { + int q = i / kv_columns; + const int col = i % kv_columns; + const int pos = q % n_total; q /= n_total; + const int row = q % (BATCH - 1) + 1; q /= BATCH - 1; + const int kind = q % 2, layer = q / 2; + const AttnLayer &L = bb->attn[layer]; + const int offset = kv_offset(L, pos, col / BB_HEAD_DIM) + col % BB_HEAD_DIM; + kv_slot(L, row, kind, pos)[offset] = kv_slot(L, 0, kind, pos)[offset]; + } + grid_barrier(g.barrier); + } + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } +#if MK2_PREFILL_ONLY + return; +#else + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, bb->phoneme_b); + grid_barrier(g.barrier); + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + if (threadIdx.x == 0) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = 0; n < BB_PHONEME; ++n) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + bb->phoneme_out[row] = arg; + } + __syncthreads(); + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + for (int i = lane; i < D_HEAD; i += 32) { + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + for (int t = 0; t <= pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + for (int i = lane; i < D_HEAD; i += 32) { + float acc = 0.f; + for (int t = 0; t <= pos; ++t) + acc += sc[t] * __bfloat162float( + g.vcache[cache_base + t * D_HEAD + i]); + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + i] = + __float2bfloat16(acc); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, g.sampled_pos, + sort_keys, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { +#ifdef ABLATE_SAMPLE + __syncthreads(); +#endif + } else { + grid_barrier(g.barrier); + } + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x < BATCH && blockIdx.x == 0) { + const int code = g.sampled_pos[threadIdx.x]; + g.sampled[threadIdx.x * N_POS + pos] = code; + if (g.native_codes) + g.native_codes[threadIdx.x * N_POS + pos] = (long long)code; + } + grid_barrier(g.barrier); + } + if (g.native_seq_len && blockIdx.x == 0 && threadIdx.x < BATCH) + g.native_seq_len[threadIdx.x] += 1; + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +#endif +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 2 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK), + 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = MK2_PREFILL_ONLY + ? out->sms * out->blocks_per_sm + : std::min(out->sms * out->blocks_per_sm, D_FFN / OUT_ATOM); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + 0, + nullptr, + nullptr, + nullptr, nullptr, nullptr, nullptr, nullptr, 0, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "dedup_native_prefill_impl.inc" diff --git a/deploy/nemotron_native/kernels/prefill/b2/mk2_abi.cuh b/deploy/nemotron_native/kernels/prefill/b2/mk2_abi.cuh new file mode 100644 index 0000000..95a253b --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b2/mk2_abi.cuh @@ -0,0 +1,82 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} diff --git a/deploy/nemotron_native/kernels/prefill/b2/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/prefill/b2/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b2/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/prefill/b2/prefill.cu b/deploy/nemotron_native/kernels/prefill/b2/prefill.cu new file mode 100644 index 0000000..c514331 --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b2/prefill.cu @@ -0,0 +1,18 @@ +#define MK2_BATCH 2 +#define MK2_BATCH_TILE 64 +#define MK2_PREFILL_ONLY 1 +#define STAGES 3 +#define KCHUNK 128 +#define MK2_GEMM_CONSUMERS 3 +#define MK2_THREADS 512 +#define MK2_NO_MAIN 1 +#define MK2_DIRECT_INPUT 1 +#define MK2_B1_LOCAL_BARRIERS 1 +#define MK2_SKIP_HIDDEN_SNAPSHOT 1 +#define MK2_ACTIVE_EXPERTS 1 +#define MK2_SSM_SPLITS 1 +#define PF_SSM_PW_MULTI 12 +#define PF_ACTIVE_TOKEN_EXPERTS 1 +#include +#include +#include "dedup_predictor_mk2.cuh" diff --git a/deploy/nemotron_native/kernels/prefill/b4/dedup_native_prefill_impl.inc b/deploy/nemotron_native/kernels/prefill/b4/dedup_native_prefill_impl.inc new file mode 100644 index 0000000..f7e9802 --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b4/dedup_native_prefill_impl.inc @@ -0,0 +1,237 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 2 || BATCH == 4 || BATCH == 8 || BATCH == 16 || BATCH == 32 || BATCH == 64)) + throw std::runtime_error("native prefill requires a matching supported physical batch"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (padded > INT_MAX / ((BB_EXPERTS + 1) * D_MODEL)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + float *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(float))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(float))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_dtda, 2 * MB_HEADS); + { + const size_t estride = padded; + const size_t grows = size_t(BB_EXPERTS) * estride; + mk2_check(owned.alloc(&hb.pf_moe_gath, grows * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_gath, 0, grows * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_hc, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_hc, 0, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_gidx, grows * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_gidx, 0, grows * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ecnt, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ecnt, 0, BB_EXPERTS * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ntiles, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ntiles, 0, BB_EXPERTS * sizeof(int))); + hb.pf_moe_estride = int(estride); + hb.pf_moe_gath_v = act_gl{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_v = act_gl{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + hb.pf_moe_gath_vn = act_gl_n{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_vn = act_gl_n{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + scratch_bytes += grows * (D_MODEL + BB_MOE_INTER) * sizeof(bf16); + } + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(owned.alloc(&hb.native_pf_cohort_check, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(hb.native_pf_cohort_check, 0, 2 * sizeof(unsigned int))); +#if PF_ACTIVE_TOKEN_EXPERTS + mk2_check(owned.alloc(&hb.native_pf_expert_masks, (padded/BATCH_PAD)*sizeof(unsigned int))); +#endif + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} +extern "C" int mk2_native_prefill_effective_sequences(MK2NativePrefill *p, int *out) { + try { + if (!p || !out) throw std::runtime_error("null effective-sequence query"); + auto *st = reinterpret_cast(p->base); + mk2_check(cudaMemcpy(out, st->host_bb.native_pf_cohort_check + 1, + sizeof(int), cudaMemcpyDeviceToHost)); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} diff --git a/deploy/nemotron_native/kernels/prefill/b4/dedup_predictor_mk2.cuh b/deploy/nemotron_native/kernels/prefill/b4/dedup_predictor_mk2.cuh new file mode 100644 index 0000000..a640917 --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b4/dedup_predictor_mk2.cuh @@ -0,0 +1,4588 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef PF_ACTIVE_TOKEN_EXPERTS +#define PF_ACTIVE_TOKEN_EXPERTS 0 +#endif +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +#ifndef MK2_PREFILL_ONLY +#define MK2_PREFILL_ONLY 0 +#endif +#ifndef MK2_BATCH_TILE +#define MK2_BATCH_TILE (((MK2_BATCH + 15) / 16) * 16) +#endif +constexpr int BATCH_PAD = MK2_BATCH_TILE; +static_assert(BATCH >= 1 && BATCH <= BATCH_PAD); +static_assert(BATCH_PAD == 16 || BATCH_PAD == 32 || BATCH_PAD == 64 || BATCH_PAD == 128, + "native prefill tile must be 16, 32, 64 or 128"); +static_assert(MK2_PREFILL_ONLY || BATCH_PAD == ((BATCH + 15) / 16) * 16, + "a wider token tile requires a prefill-only library"); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_PW_MULTI +#define PF_SSM_PW_MULTI 4 +#endif +#ifndef PF_SSM_PW_SINGLE +#define PF_SSM_PW_SINGLE 2 +#endif +static_assert(MB_HEAD_DIM % PF_SSM_PW_MULTI == 0 && + MB_HEAD_DIM % PF_SSM_PW_SINGLE == 0, + "head_dim must split evenly into scan channel groups"); +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef MK2_SMEM_SLACK +#define MK2_SMEM_SLACK 1024 +#endif +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +#ifndef MK2_THREADS +#define MK2_THREADS 256 +#endif +constexpr int THREADS = MK2_THREADS; +constexpr int NORM_THREADS = 256; +constexpr int NORM_WARPS = NORM_THREADS / 32; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +#ifndef MK2_GEMM_CONSUMERS +#define MK2_GEMM_CONSUMERS 1 +#endif +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 +#define MK2_HEAD_STAGE_BYTES 0 +#else +#define MK2_HEAD_STAGE_BYTES sizeof(out_tile_f) +#endif +constexpr int NCONS = MK2_GEMM_CONSUMERS; +static_assert(NCONS >= 1 && NCONS <= 3, "consumers occupy warpgroups 0, 2 and 3"); +static_assert(THREADS >= (NCONS + 1) * 128, "one producer plus NCONS consumers"); +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +constexpr int PAD_N = 32; +using act_tile_n = st_bf; +using out_tile_n = st_bf; +using act_gl_n = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_MOE_ACC_C, EPI_ACC_F32 }; +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + unsigned int *native_pf_expert_masks = nullptr; + float *pf_dtda = nullptr; + bf16 *pf_moe_gath = nullptr; + bf16 *pf_moe_hc = nullptr; + int *pf_moe_gidx = nullptr; + int *pf_moe_ecnt = nullptr; + int *pf_moe_ntiles = nullptr; + int pf_moe_estride = 0; + act_gl pf_moe_gath_v; + act_gl pf_moe_hc_v; + act_gl_n pf_moe_gath_vn; + act_gl_n pf_moe_hc_vn; + unsigned int *native_pf_cohort_check = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + volatile unsigned int *generation = state + 1; + const unsigned int seen = *generation; + if (atomicAdd(state, 1u) == gridDim.x - 1) { + atomicExch(state, 0u); + atomicAdd(state + 1, 1u); + } else { + while (*generation == seen) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const bf16 *in = src + (size_t)row * D_MODEL; + float v[PER_THREAD]; + float sum = 0.f, sq = 0.f; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (threadIdx.x < NORM_THREADS) { + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + v[j] = __bfloat162float(in[threadIdx.x + j * NORM_THREADS]); + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + if (lane == 0) { scratch[warp] = sum; scratch[NORM_WARPS + warp] = sq; } + } + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f, q = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) { s += scratch[w]; q += scratch[NORM_WARPS + w]; } + scratch[2 * NORM_WARPS] = s / D_MODEL; + scratch[2 * NORM_WARPS + 1] = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + } + __syncthreads(); + const float mean = scratch[2 * NORM_WARPS], inv = scratch[2 * NORM_WARPS + 1]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + if (threadIdx.x >= NORM_THREADS) break; + const int i = threadIdx.x + j * NORM_THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * weight[i]); + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, + int idx, unsigned mask); +__device__ __forceinline__ void norm_sub_sync(int sub) { + if (THREADS == NORM_THREADS) { __syncthreads(); return; } + asm volatile("bar.sync %0, %1;" :: "r"(6 + sub), "n"(NORM_THREADS) : "memory"); +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, const float *fold_src = nullptr, + const float *fold_routes = nullptr) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + constexpr int NSUB = THREADS / NORM_THREADS; + if (threadIdx.x >= NSUB * NORM_THREADS) return; + const int sub = threadIdx.x / NORM_THREADS; + const int tid = threadIdx.x - sub * NORM_THREADS; + float *const sc = scratch + sub * (NORM_WARPS + 1); + for (int row = blockIdx.x * NSUB + sub; row < n_rows; + row += gridDim.x * NSUB) { + const size_t base = (size_t)row * D_MODEL; + unsigned fold_mask = (1u << BB_EXPERTS) - 1u; + if (fold_src && fold_routes) { + fold_mask = 0u; + const float *rw = fold_routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) + if (rw[e] != 0.f) fold_mask |= 1u << e; + } + float v[PER_THREAD]; + float sq = 0.f; + const int warp = tid >> 5, lane = tid & 31; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + bf16 hv; + if (fold_src) { + hv = bb_moe_fold_deterministic(fold_src, (int)(base + i), fold_mask); + h[base + i] = hv; + } else { + hv = h[base + i]; + } + v[j] = __bfloat162float(hv) + __bfloat162float(r[base + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + if (lane == 0) sc[warp] = sq; + norm_sub_sync(sub); + if (tid == 0) { + float s = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) s += sc[w]; + sc[NORM_WARPS] = rsqrtf(s / D_MODEL + RMS_EPS); + } + norm_sub_sync(sub); + const float inv = sc[NORM_WARPS]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + out[base + i] = __float2bfloat16( + __bfloat162float(weight[i]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + } + norm_sub_sync(sub); + } +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + const int c = idx - row * MB_CONV_DIM; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + const float raw = __bfloat162float( + zxbcdt[(size_t)row * MB_IN_PROJ + MB_INNER + c]); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + xbc_act[(size_t)row * MB_CONV_DIM + c] = __float2bfloat16(acc); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + const int pairs = BATCH * MB_HEADS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_HEADS * split); + const int head = (pair / split) % MB_HEADS; + const int first_channel = (pair % split) * channels; + const int group = head / MB_HEADS_PER_GROUP; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + L.dt_bias[head]); + const float dA = __expf(L.A[head] * dt); + const float Dh = L.D[head]; + float *st = L.ssm_state + (size_t)slot_map[row] * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + for (int p = first_channel; p < first_channel + channels; ++p) { + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + const size_t at = (size_t)p * MB_STATE + n; + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + static_assert(MB_CONV_K == 4, "packed conv taps assume a 4-tap kernel"); + constexpr int CPT = 4; + static_assert(MB_CONV_DIM % CPT == 0, "channel vector width"); + constexpr int CGROUPS = MB_CONV_DIM / CPT; + const int total = n_tokens * CGROUPS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / CGROUPS; + const int c = (idx - t * CGROUPS) * CPT; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + float wv[CPT][MB_CONV_K]; + { + const float4 *wsrc = reinterpret_cast( + L.conv_w + (size_t)c * MB_CONV_K); + #pragma unroll + for (int j = 0; j < CPT / 2; ++j) { + const float4 w4 = wsrc[j]; + const unsigned uu0 = __float_as_uint(w4.x); + const unsigned uu1 = __float_as_uint(w4.y); + const unsigned uu2 = __float_as_uint(w4.z); + const unsigned uu3 = __float_as_uint(w4.w); + wv[2 * j][0] = __uint_as_float(uu0 << 16); + wv[2 * j][1] = __uint_as_float(uu0 & 0xffff0000u); + wv[2 * j][2] = __uint_as_float(uu1 << 16); + wv[2 * j][3] = __uint_as_float(uu1 & 0xffff0000u); + wv[2 * j + 1][0] = __uint_as_float(uu2 << 16); + wv[2 * j + 1][1] = __uint_as_float(uu2 & 0xffff0000u); + wv[2 * j + 1][2] = __uint_as_float(uu3 << 16); + wv[2 * j + 1][3] = __uint_as_float(uu3 & 0xffff0000u); + } + } + float acc[CPT]; + { + const uint2 b2 = *reinterpret_cast(L.conv_b + c); + acc[0] = __uint_as_float(b2.x << 16); + acc[1] = __uint_as_float(b2.x & 0xffff0000u); + acc[2] = __uint_as_float(b2.y << 16); + acc[3] = __uint_as_float(b2.y & 0xffff0000u); + } + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v[CPT]; + if (src < 0) { + #pragma unroll + for (int q = 0; q < CPT; ++q) + v[q] = __bfloat162float( + state[(size_t)q * (MB_CONV_K - 1) + pos + k]); + } else { + const uint2 p = *reinterpret_cast( + zxbcdt + (size_t)src * MB_IN_PROJ + MB_INNER + c); + v[0] = __uint_as_float(p.x << 16); + v[1] = __uint_as_float(p.x & 0xffff0000u); + v[2] = __uint_as_float(p.y << 16); + v[3] = __uint_as_float(p.y & 0xffff0000u); + } + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] += wv[q][k] * v[q]; + } + uint2 out; + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] = acc[q] / (1.f + __expf(-acc[q])); + out.x = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[0])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[1])) << 16); + out.y = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[2])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[3])) << 16); + *reinterpret_cast(xbc_act + (size_t)t * MB_CONV_DIM + c) = out; + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_dtda_prefill( + const MambaLayer &L, const bf16 *zxbcdt, float *dtda, int n_tokens) { + const int total = n_tokens * MB_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_HEADS; + const int head = idx - t * MB_HEADS; + const float dt = mamba_softplus( + __bfloat162float(zxbcdt[(size_t)t * MB_IN_PROJ + MB_INNER + + MB_CONV_DIM + head]) + L.dt_bias[head]); + dtda[(size_t)idx * 2] = dt; + dtda[(size_t)idx * 2 + 1] = __expf(L.A[head] * dt); + } +} +template +__device__ __forceinline__ void mamba_ssm_prefill_impl( + const MambaLayer &L, const bf16 *xbc_act, const float *dtda, bf16 *gated, + int n_pos, int n_seq, int slot, const int *seq_slot) { + constexpr int NS = MB_STATE / 32; + constexpr int PG = MB_HEAD_DIM / PW; + constexpr bool RS = (PW % 4) == 0; + constexpr int H1 = RS ? PW / 2 : 1; + constexpr int H2 = RS ? PW / 4 : 1; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const bool keep_hi = (lane & 16) != 0; + const bool keep_hi2 = (lane & 8) != 0; + const int warps_total = gridDim.x * WARPS; + const int units = n_seq * MB_HEADS * PG; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int pg = unit % PG; + const int rest = unit / PG; + const int head = rest % MB_HEADS; + const int seq = rest / MB_HEADS; + const int p0 = pg * PW; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Dh = L.D[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p0 * MB_STATE; + float sreg[PW][NS]; + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + sreg[q][j] = st[(size_t)q * MB_STATE + lane + j * 32]; + } + auto fetch = [&](int p, float (&oB)[NS], float (&oC)[NS], + float (&oX)[PW], float2 &od) { + const size_t tt = (size_t)p * n_seq + seq; + const bf16 *xa2 = xbc_act + tt * MB_CONV_DIM; + const bf16 *Bs = xa2 + MB_INNER + group * MB_STATE; + const bf16 *Cs = xa2 + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + od = *reinterpret_cast( + dtda + (tt * MB_HEADS + head) * 2); + #pragma unroll + for (int j = 0; j < NS; ++j) { + oB[j] = __bfloat162float(Bs[lane + j * 32]); + oC[j] = __bfloat162float(Cs[lane + j * 32]); + } + if constexpr (RS) { + const uint2 *xp = reinterpret_cast( + xa2 + head * MB_HEAD_DIM + p0); + #pragma unroll + for (int q = 0; q < PW; q += 4) { + const uint2 w = xp[q >> 2]; + oX[q] = __uint_as_float(w.x << 16); + oX[q + 1] = __uint_as_float(w.x & 0xffff0000u); + oX[q + 2] = __uint_as_float(w.y << 16); + oX[q + 3] = __uint_as_float(w.y & 0xffff0000u); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) + oX[q] = __bfloat162float(xa2[head * MB_HEAD_DIM + p0 + q]); + } + }; + float pB[NS], pC[NS], pX[PW]; + float2 pd; + fetch(0, pB, pC, pX, pd); + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const float dt = pd.x; + const float dA = pd.y; + float Bv[NS], Cv[NS]; + #pragma unroll + for (int j = 0; j < NS; ++j) { + Bv[j] = pB[j] * dt; + Cv[j] = pC[j]; + } + float acc[PW], xvs[PW]; + #pragma unroll + for (int q = 0; q < PW; ++q) xvs[q] = pX[q]; + if (pos + 1 < n_pos) fetch(pos + 1, pB, pC, pX, pd); + #pragma unroll + for (int q = 0; q < PW; ++q) { + const float xv = xvs[q]; + float a = 0.f; + #pragma unroll + for (int j = 0; j < NS; ++j) { + sreg[q][j] = sreg[q][j] * dA + Bv[j] * xv; + a += sreg[q][j] * Cv[j]; + } + acc[q] = a; + } + if constexpr (RS) { + float m1[H1], t1[H1], x1[H1]; + #pragma unroll + for (int q = 0; q < H1; ++q) { + m1[q] = keep_hi ? acc[q + H1] : acc[q]; + t1[q] = keep_hi ? acc[q] : acc[q + H1]; + x1[q] = keep_hi ? xvs[q + H1] : xvs[q]; + } + #pragma unroll + for (int q = 0; q < H1; ++q) + m1[q] += __shfl_xor_sync(0xffffffff, t1[q], 16); + float m2[H2], t2[H2], x2[H2]; + #pragma unroll + for (int q = 0; q < H2; ++q) { + m2[q] = keep_hi2 ? m1[q + H2] : m1[q]; + t2[q] = keep_hi2 ? m1[q] : m1[q + H2]; + x2[q] = keep_hi2 ? x1[q + H2] : x1[q]; + } + #pragma unroll + for (int q = 0; q < H2; ++q) + m2[q] += __shfl_xor_sync(0xffffffff, t2[q], 8); + #pragma unroll + for (int q = 0; q < H2; ++q) { + #pragma unroll + for (int off = 4; off > 0; off >>= 1) + m2[q] += __shfl_xor_sync(0xffffffff, m2[q], off); + } + if ((lane & 7) == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + + p0 + (keep_hi ? H1 : 0) + (keep_hi2 ? H2 : 0); + #pragma unroll + for (int q = 0; q < H2; ++q) + dst[q] = __float2bfloat16(m2[q] + x2[q] * Dh); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc[q] += __shfl_xor_sync(0xffffffff, acc[q], off); + } + if (lane == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + p0; + #pragma unroll + for (int q = 0; q < PW; ++q) + dst[q] = __float2bfloat16(acc[q] + xvs[q] * Dh); + } + } + } + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + st[(size_t)q * MB_STATE + lane + j * 32] = sreg[q][j]; + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, + const float *dtda, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + (void)zxbcdt; + if (n_seq > 1) + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens / n_seq, n_seq, slot, seq_slot); + else + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens, 1, slot, nullptr); +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16( + held[j] * inv * __bfloat162float(L.gnorm_w[idx])); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = BATCH * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / BB_ATT_HEADS; + const int head = pair % BB_ATT_HEADS; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const int kv = head / kv_per_q; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *q = base + head * BB_HEAD_DIM; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (head % kv_per_q == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + for (int tpos = 0; tpos < T; ++tpos) { +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new + : kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#endif + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m, s); + const float rescale = __expf(m - m_new); + const float w = __expf(s - m_new); + denom = denom * rescale + w; +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *vt = kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#else + const bf16 *vt = cur_token + ? v_new + : kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + w * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + constexpr int CHUNKS = BB_HEAD_DIM / 8; + const int total = n_tokens * BB_KV_HEADS * CHUNKS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int c = idx % CHUNKS; + const int rest = idx / CHUNKS; + const int t = rest / BB_KV_HEADS; + const int kv = rest - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM + c * 8; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM + c * 8; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv) + c * 8; + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv) + c * 8; + *reinterpret_cast(kc) = *reinterpret_cast(k_new); + *reinterpret_cast(vc) = *reinterpret_cast(v_new); + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float qv[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qv[j] = __bfloat162float(q[lane + j * 32]); + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + const int *const krow = + L.block_table + ((size_t)seq * 2 + 0) * L.max_blocks_per_seq; + const int *const vrow = + L.block_table + ((size_t)seq * 2 + 1) * L.max_blocks_per_seq; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + const int head_bias = kv * L.tokens_per_block * BB_HEAD_DIM; + int page = -1; + const bf16 *kpage = nullptr, *vpage = nullptr; + float kn[PER_LANE], vn[PER_LANE]; + { + page = 0; + kpage = L.pool_base + (size_t)(krow[0] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + (size_t)(vrow[0] & 0x7FFFFFFF) * elems_per_block; + const int at = head_bias + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + for (int s = 0; s <= pos; ++s) { + float kc[PER_LANE], vc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { kc[j] = kn[j]; vc[j] = vn[j]; } + if (s < pos) { + const int np = (s + 1) / L.tokens_per_block; + if (np != page) { + page = np; + kpage = L.pool_base + + (size_t)(krow[np] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + + (size_t)(vrow[np] & 0x7FFFFFFF) * elems_per_block; + } + const int at = head_bias + + ((s + 1) % L.tokens_per_block) * BB_HEAD_DIM + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) dot += qv[j] * kc[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[j] = acc[j] * rescale + wt * vc[j]; + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH, unsigned int *tile_masks = nullptr, + bf16 *gath = nullptr, int *gidx = nullptr, int *ecnt = nullptr, + int estride = 0, bf16 *row_stage = nullptr) { + __shared__ float score[BB_EXPERTS]; + __shared__ int s_pick[BB_TOP_K], s_slot[BB_TOP_K]; + constexpr int CH = D_MODEL / 8; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + { + const float4 *gsrc = + reinterpret_cast(normed + (size_t)row * D_MODEL); + float4 *gdst = reinterpret_cast(row_stage); + for (int i = threadIdx.x; i < CH; i += THREADS) gdst[i] = gsrc[i]; + } + if ((int)threadIdx.x < BB_EXPERTS) + moe_w[(size_t)row * BB_EXPERTS + threadIdx.x] = 0.f; + __syncthreads(); + for (int e = warp; e < BB_EXPERTS; e += WARPS) { + const bf16 *wr = gate_w + (size_t)e * D_MODEL; + float acc = 0.f; + for (int i = lane; i < D_MODEL; i += 32) + acc += __bfloat162float(row_stage[i]) * + __bfloat162float(wr[i]); + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) score[e] = 1.f / (1.f + __expf(-acc)); + } + __syncthreads(); + if (threadIdx.x == 0) { + float remaining[BB_EXPERTS]; + int pick[BB_TOP_K]; + for (int e = 0; e < BB_EXPERTS; ++e) + remaining[e] = score[e] + __bfloat162float(gate_bias[e]); + float total = 0.f; + for (int k = 0; k < BB_TOP_K; ++k) { + int best = 0; + for (int e = 1; e < BB_EXPERTS; ++e) + if (remaining[e] > remaining[best]) best = e; + pick[k] = best; + total += score[best]; + remaining[best] = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + unsigned int bits = 0u; + for (int k = 0; k < BB_TOP_K; ++k) { + moe_w[row * BB_EXPERTS + pick[k]] = score[pick[k]] * norm; + bits |= 1u << pick[k]; + s_pick[k] = pick[k]; + } + if (tile_masks) atomicOr(&tile_masks[row / BATCH_PAD], bits); + } + __syncthreads(); + if (gath) { + if ((int)threadIdx.x < BB_TOP_K) { + const int k = threadIdx.x; + const int slot = atomicAdd(&ecnt[s_pick[k]], 1); + s_slot[k] = slot; + gidx[(size_t)s_pick[k] * estride + slot] = row; + } + __syncthreads(); + const float4 *src = reinterpret_cast(row_stage); + for (int k = 0; k < BB_TOP_K; ++k) { + float4 *dst = reinterpret_cast( + gath + ((size_t)s_pick[k] * estride + s_slot[k]) * D_MODEL); + for (int i = threadIdx.x; i < CH; i += THREADS) dst[i] = src[i]; + } + } + __syncthreads(); + } +} +__device__ __forceinline__ unsigned moe_row_mask(const float *routes, int row) { + unsigned m = 0u; + const float *rw = routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) if (rw[e] != 0.f) m |= 1u << e; + return m; +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx, + unsigned mask) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + for (unsigned m = mask; m; m &= m - 1u) + value = __fadd_rn(value, parts[base + (size_t)(__ffs(m) - 1) * D_MODEL]); + return __float2bfloat16( + __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL])); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + uint64_t *keys, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int n = threadIdx.x; n < SORT_N; n += THREADS) { + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + keys[n] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + __syncthreads(); + for (int k = 2; k <= SORT_N; k <<= 1) { + if (k >= 64) __syncthreads(); + for (int j = k >> 1; j > 0; j >>= 1) { + for (int i = threadIdx.x; i < SORT_N; i += THREADS) { + const int partner = i ^ j; + if (partner > i) { + const bool descending = ((i & k) == 0); + if ((keys[i] < keys[partner]) == descending) { + const uint64_t t = keys[i]; + keys[i] = keys[partner]; + keys[partner] = t; + } + } + } + if (j >= 32) __syncthreads(); + else __syncwarp(); + } + } + __syncthreads(); + if (threadIdx.x == 0) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float v; + int n; + if (!eos_used && eos_v > av) { + v = eos_v; + n = AUDIO_EOS; + eos_used = true; + } else { + v = av; + n = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = + v + gumbel[((size_t)row * N_POS + pos) * TOP_K + j]; + if (noised > best_noised) { best_noised = noised; best_code = n; } + } + out_codes[row] = best_code; + } + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void prefill_expert_masks(const float *routes, unsigned int *masks, int rows) { + const int lane=threadIdx.x&31, warp=threadIdx.x>>5; + const int nt=(rows+BATCH_PAD-1)/BATCH_PAD; + for (int t=blockIdx.x*WARPS+warp;t(gath + g * D_MODEL + c * 8); + if (local < ecnt[e]) + *dst = *reinterpret_cast( + normed + (size_t)gidx[g] * D_MODEL + c * 8); + else + *dst = make_float4(0.f, 0.f, 0.f, 0.f); + } +} +__device__ __forceinline__ int gemm_eff(int tiles, int reps) { + if (NCONS <= 1) return 1; + const int g = (int)gridDim.x; + const int rounds1 = (tiles * reps + g - 1) / g; + if constexpr (NCONS == 3) { + const int max_eff = reps <= BB_EXPERTS ? 2 : 3; + const int rounds2 = (((tiles + 1) / 2) * reps + g - 1) / g; + if (max_eff == 2) return rounds2 < rounds1 ? 2 : 1; + const int rounds3 = (((tiles + 2) / 3) * reps + g - 1) / g; + if (rounds3 < rounds2 && rounds3 < rounds1) return 3; + return rounds2 < rounds1 ? 2 : 1; + } else { + const int roundsN = (((tiles + NCONS - 1) / NCONS) * reps + g - 1) / g; + return roundsN < rounds1 ? NCONS : 1; + } +} +__device__ __forceinline__ int gemm_groups(int tiles, int reps) { + const int eff = gemm_eff(tiles, reps); + return (tiles + eff - 1) / eff; +} +__device__ __forceinline__ int gemm_consumer_id() { + const int g = warpgroup::groupid(); + if (g == 1) return -1; + const int c = g == 0 ? 0 : g - 1; + return c < NCONS ? c : -1; +} +__device__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][BATCH_PAD]; + __shared__ int moe_gidx_stage[NCONS][BATCH_PAD]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + a_s[s], a, {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + a_s[s], a, {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / BATCH_PAD) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, BATCH_PAD> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], a_s[s]); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * BATCH_PAD; + if (epi == EPI_BIAS) { + warpgroup::store(o_sf, accum); + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + dest[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(o_sf[int2{col, row}] + + __bfloat162float(bias[base + col])); + } + warpgroup::sync(1 + cons); + continue; + } + warpgroup::store(o_s[cons], accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_s[cons][int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__device__ void gemm_narrow( + const weight_gl &w, const act_gl_n &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = PAD_N, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][PAD_N]; + __shared__ int moe_gidx_stage[NCONS][PAD_N]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + PAD_N - 1) / PAD_N; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + PAD_N - 1) / PAD_N; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / PAD_N) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, PAD_N> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], + *reinterpret_cast(&a_s[s])); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * PAD_N; + out_tile_n &o_sn = *reinterpret_cast(&o_s[cons]); + warpgroup::store(o_sn, accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_sn[int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&w_s)[STAGES * NCONS] = + alloc.allocate(); + act_tile (&a_s)[STAGES] = alloc.allocate(); + out_tile (&o_s)[NCONS] = alloc.allocate(); +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 + static_assert(NCONS * sizeof(out_tile) >= sizeof(out_tile_f), + "aliased head staging must fit the output tiles"); + out_tile_f &o_sf = *reinterpret_cast(&o_s[0]); +#else + out_tile_f (&o_sf) = alloc.allocate(); +#endif + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_WARPS + 2]; +#if MK2_PREFILL_ONLY + bf16 *const route_stage = reinterpret_cast(&o_s[0]); + static_assert(sizeof(o_s) >= D_MODEL * sizeof(bf16), + "router row stage must fit the output tiles"); +#else + __shared__ __align__(16) uint64_t sort_keys[SORT_N]; + bf16 *const route_stage = reinterpret_cast(&sort_keys[0]); + static_assert(sizeof(sort_keys) >= D_MODEL * sizeof(bf16), + "router row stage must fit the sampler workspace"); +#endif + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS * NCONS, 0); + for (int c = 0; c < CONSUMER_WARPS * NCONS; ++c) arrive(drained[s]); + } + } + int gemm_iter = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map); + grid_barrier(g.barrier); + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map); + grid_barrier(g.barrier); + mamba_gated_norm(L, bb->zxbcdt, bb->gated); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, nullptr, 1, BATCH, + bb->moe_w, BB_MOE_INTER / OUT_ATOM); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + 1, BATCH, bb->moe_w, D_MODEL / OUT_ATOM); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_deterministic(bb->moe_out, idx, + (1u << BB_EXPERTS) - 1u); + grid_barrier(g.barrier); + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int original_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + bool identical_cohort = false; + if constexpr (BATCH > 1) { + if (bb->native_pf_cohort_check && bb->native_prefill_reset && + bb->pf_n_seq == BATCH && !bb->pf_fuse_decode && + !bb->pf_layer_out && !bb->moe_w_inject && !bb->moe_w_audit) { + if (blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[0] = 0; + grid_barrier(g.barrier); + const auto *bits = reinterpret_cast(bb->prompt); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < original_total * D_MODEL; idx += gridDim.x * THREADS) { + const int canonical = (idx / (BATCH * D_MODEL)) * BATCH * D_MODEL + idx % D_MODEL; + if (bits[idx] != bits[canonical]) atomicOr(bb->native_pf_cohort_check, 1u); + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + if (bb->pf_seq_slot[row] != row || bb->slot_map[row] != row) + atomicOr(bb->native_pf_cohort_check, 1u); + grid_barrier(g.barrier); + identical_cohort = bb->native_pf_cohort_check[0] == 0; + } + } + const int n_total = identical_cohort ? original_total / BATCH : original_total; + if (bb->native_pf_cohort_check && blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[1] = identical_cohort ? 1 : BATCH; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int reset_slots = identical_cohort ? 1 : BATCH; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * reset_slots * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % reset_slots; + const MambaLayer &L = bb->mamba[row_layer / reset_slots]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * reset_slots * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % reset_slots; + const MambaLayer &L = bb->mamba[row_layer / reset_slots]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = original_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + const int src = identical_cohort ? (idx / D_MODEL) * BATCH * D_MODEL + idx % D_MODEL : idx; + bb->pf_h[idx] = bb->prompt[src]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = identical_cohort ? 1 : (bb->pf_n_seq > 0 ? bb->pf_n_seq : 1); + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T, + (!bb->pf_layer_out && layer > 0 && + bb->moe_of_layer[layer - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + if (bb->native_pf_expert_masks) + for (int i = blockIdx.x * THREADS + threadIdx.x; i < n_tok; + i += gridDim.x * THREADS) + bb->native_pf_expert_masks[i] = 0u; + if (bb->pf_moe_ecnt) + for (int e = blockIdx.x * THREADS + threadIdx.x; e < BB_EXPERTS; + e += gridDim.x * THREADS) + bb->pf_moe_ecnt[e] = 0; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + mamba_dtda_prefill(L, bb->pf_zxbcdt, bb->pf_dtda, T); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_dtda, + bb->pf_gated, T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, + T, bb->native_pf_expert_masks, bb->pf_moe_gath, + bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, route_stage); + } + grid_barrier(g.barrier); + int max_ecnt = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + max_ecnt = max(max_ecnt, bb->pf_moe_ecnt[e]); + const bool narrow = BATCH_PAD > PAD_N && max_ecnt <= PAD_N; + const int etile = narrow ? PAD_N : BATCH_PAD; + int nt_sum = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_sum += (bb->pf_moe_ecnt[e] + etile - 1) / etile; + const int up_steps = + gemm_groups(BB_MOE_INTER / OUT_ATOM, nt_sum) * + nt_sum * (D_MODEL / KCHUNK); + const int shup_steps = + gemm_groups(BB_SHARED_INTER / OUT_ATOM, n_tok) * + n_tok * (D_MODEL / KCHUNK); + const int down_steps = + gemm_groups(D_MODEL / OUT_ATOM, nt_sum) * + nt_sum * (BB_MOE_INTER / KCHUNK); + const int shdown_steps = + gemm_groups(D_MODEL / OUT_ATOM, n_tok) * + n_tok * (BB_SHARED_INTER / KCHUNK); + if (narrow) + gemm_narrow(L.exp_up, bb->pf_moe_gath_vn, + BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, PAD_N, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + else + gemm(L.exp_up, bb->pf_moe_gath_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, BATCH_PAD, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + up_steps, up_steps + shup_steps); + grid_barrier(g.barrier); + if (narrow) + gemm_narrow(L.exp_down, bb->pf_moe_hc_vn, + BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, PAD_N, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + else + gemm(L.exp_down, bb->pf_moe_hc_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + down_steps, down_steps + shdown_steps); + grid_barrier(g.barrier); + if (bb->pf_layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * D_MODEL; idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = + bb_moe_fold_deterministic(bb->pf_moe_out, idx, + moe_row_mask(bb->pf_moe_w, idx / D_MODEL)); + grid_barrier(g.barrier); + } + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, + identical_cohort ? bb->pf_normed : bb->pf_final, ln_scratch, + n_total, + (!bb->pf_layer_out && n_backbone_layers > 0 && + bb->moe_of_layer[n_backbone_layers - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + grid_barrier(g.barrier); + if (identical_cohort) { + const int worker = blockIdx.x * THREADS + threadIdx.x; + const int workers = gridDim.x * THREADS; + for (int i = worker; i < original_total * D_MODEL; i += workers) + bb->pf_final[i] = bb->pf_normed[(i / (BATCH * D_MODEL)) * D_MODEL + i % D_MODEL]; + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = worker; i < 14 * (BATCH - 1) * CV; i += workers) { + const int cell = i % CV, pair = i / CV; + const int row = pair % (BATCH - 1) + 1; + const MambaLayer &L = bb->mamba[pair / (BATCH - 1)]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = L.conv_state[cell]; + } + for (int i = worker; i < 14 * (BATCH - 1) * SS; i += workers) { + const int cell = i % SS, pair = i / SS; + const int row = pair % (BATCH - 1) + 1; + const MambaLayer &L = bb->mamba[pair / (BATCH - 1)]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = L.ssm_state[cell]; + } + const int kv_columns = BB_KV_HEADS * BB_HEAD_DIM; + for (int i = worker; i < 3 * 2 * (BATCH - 1) * n_total * kv_columns; i += workers) { + int q = i / kv_columns; + const int col = i % kv_columns; + const int pos = q % n_total; q /= n_total; + const int row = q % (BATCH - 1) + 1; q /= BATCH - 1; + const int kind = q % 2, layer = q / 2; + const AttnLayer &L = bb->attn[layer]; + const int offset = kv_offset(L, pos, col / BB_HEAD_DIM) + col % BB_HEAD_DIM; + kv_slot(L, row, kind, pos)[offset] = kv_slot(L, 0, kind, pos)[offset]; + } + grid_barrier(g.barrier); + } + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } +#if MK2_PREFILL_ONLY + return; +#else + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, bb->phoneme_b); + grid_barrier(g.barrier); + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + if (threadIdx.x == 0) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = 0; n < BB_PHONEME; ++n) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + bb->phoneme_out[row] = arg; + } + __syncthreads(); + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + for (int i = lane; i < D_HEAD; i += 32) { + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + for (int t = 0; t <= pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + for (int i = lane; i < D_HEAD; i += 32) { + float acc = 0.f; + for (int t = 0; t <= pos; ++t) + acc += sc[t] * __bfloat162float( + g.vcache[cache_base + t * D_HEAD + i]); + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + i] = + __float2bfloat16(acc); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, g.sampled_pos, + sort_keys, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { +#ifdef ABLATE_SAMPLE + __syncthreads(); +#endif + } else { + grid_barrier(g.barrier); + } + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x < BATCH && blockIdx.x == 0) { + const int code = g.sampled_pos[threadIdx.x]; + g.sampled[threadIdx.x * N_POS + pos] = code; + if (g.native_codes) + g.native_codes[threadIdx.x * N_POS + pos] = (long long)code; + } + grid_barrier(g.barrier); + } + if (g.native_seq_len && blockIdx.x == 0 && threadIdx.x < BATCH) + g.native_seq_len[threadIdx.x] += 1; + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +#endif +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 2 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK), + 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = MK2_PREFILL_ONLY + ? out->sms * out->blocks_per_sm + : std::min(out->sms * out->blocks_per_sm, D_FFN / OUT_ATOM); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + 0, + nullptr, + nullptr, + nullptr, nullptr, nullptr, nullptr, nullptr, 0, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "dedup_native_prefill_impl.inc" diff --git a/deploy/nemotron_native/kernels/prefill/b4/mk2_abi.cuh b/deploy/nemotron_native/kernels/prefill/b4/mk2_abi.cuh new file mode 100644 index 0000000..95a253b --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b4/mk2_abi.cuh @@ -0,0 +1,82 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} diff --git a/deploy/nemotron_native/kernels/prefill/b4/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/prefill/b4/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b4/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/prefill/b4/prefill.cu b/deploy/nemotron_native/kernels/prefill/b4/prefill.cu new file mode 100644 index 0000000..1398c0d --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b4/prefill.cu @@ -0,0 +1,18 @@ +#define MK2_BATCH 4 +#define MK2_BATCH_TILE 64 +#define MK2_PREFILL_ONLY 1 +#define STAGES 3 +#define KCHUNK 128 +#define MK2_GEMM_CONSUMERS 3 +#define MK2_THREADS 512 +#define MK2_NO_MAIN 1 +#define MK2_DIRECT_INPUT 1 +#define MK2_B1_LOCAL_BARRIERS 1 +#define MK2_SKIP_HIDDEN_SNAPSHOT 1 +#define MK2_ACTIVE_EXPERTS 1 +#define MK2_SSM_SPLITS 1 +#define PF_SSM_PW_MULTI 12 +#define PF_ACTIVE_TOKEN_EXPERTS 1 +#include +#include +#include "dedup_predictor_mk2.cuh" diff --git a/deploy/nemotron_native/kernels/prefill/b64/dedup_native_prefill_impl.inc b/deploy/nemotron_native/kernels/prefill/b64/dedup_native_prefill_impl.inc new file mode 100644 index 0000000..cdaa0ea --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b64/dedup_native_prefill_impl.inc @@ -0,0 +1,237 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 16 || BATCH == 32 || BATCH == 64)) + throw std::runtime_error("native prefill requires matching batch 1/16/32/64"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (padded > INT_MAX / ((BB_EXPERTS + 1) * D_MODEL)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + float *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(float))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(float))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_dtda, 2 * MB_HEADS); + { + const size_t estride = padded; + const size_t grows = size_t(BB_EXPERTS) * estride; + mk2_check(owned.alloc(&hb.pf_moe_gath, grows * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_gath, 0, grows * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_hc, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_hc, 0, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_gidx, grows * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_gidx, 0, grows * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ecnt, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ecnt, 0, BB_EXPERTS * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ntiles, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ntiles, 0, BB_EXPERTS * sizeof(int))); + hb.pf_moe_estride = int(estride); + hb.pf_moe_gath_v = act_gl{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_v = act_gl{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + hb.pf_moe_gath_vn = act_gl_n{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_vn = act_gl_n{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + scratch_bytes += grows * (D_MODEL + BB_MOE_INTER) * sizeof(bf16); + } + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(owned.alloc(&hb.native_pf_cohort_check, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(hb.native_pf_cohort_check, 0, 2 * sizeof(unsigned int))); +#if PF_ACTIVE_TOKEN_EXPERTS + mk2_check(owned.alloc(&hb.native_pf_expert_masks, (padded/BATCH_PAD)*sizeof(unsigned int))); +#endif + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} +extern "C" int mk2_native_prefill_effective_sequences(MK2NativePrefill *p, int *out) { + try { + if (!p || !out) throw std::runtime_error("null effective-sequence query"); + auto *st = reinterpret_cast(p->base); + mk2_check(cudaMemcpy(out, st->host_bb.native_pf_cohort_check + 1, + sizeof(int), cudaMemcpyDeviceToHost)); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} diff --git a/deploy/nemotron_native/kernels/prefill/b64/dedup_predictor_mk2.cuh b/deploy/nemotron_native/kernels/prefill/b64/dedup_predictor_mk2.cuh new file mode 100644 index 0000000..a640917 --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b64/dedup_predictor_mk2.cuh @@ -0,0 +1,4588 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef PF_ACTIVE_TOKEN_EXPERTS +#define PF_ACTIVE_TOKEN_EXPERTS 0 +#endif +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +#ifndef MK2_PREFILL_ONLY +#define MK2_PREFILL_ONLY 0 +#endif +#ifndef MK2_BATCH_TILE +#define MK2_BATCH_TILE (((MK2_BATCH + 15) / 16) * 16) +#endif +constexpr int BATCH_PAD = MK2_BATCH_TILE; +static_assert(BATCH >= 1 && BATCH <= BATCH_PAD); +static_assert(BATCH_PAD == 16 || BATCH_PAD == 32 || BATCH_PAD == 64 || BATCH_PAD == 128, + "native prefill tile must be 16, 32, 64 or 128"); +static_assert(MK2_PREFILL_ONLY || BATCH_PAD == ((BATCH + 15) / 16) * 16, + "a wider token tile requires a prefill-only library"); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_PW_MULTI +#define PF_SSM_PW_MULTI 4 +#endif +#ifndef PF_SSM_PW_SINGLE +#define PF_SSM_PW_SINGLE 2 +#endif +static_assert(MB_HEAD_DIM % PF_SSM_PW_MULTI == 0 && + MB_HEAD_DIM % PF_SSM_PW_SINGLE == 0, + "head_dim must split evenly into scan channel groups"); +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef MK2_SMEM_SLACK +#define MK2_SMEM_SLACK 1024 +#endif +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +#ifndef MK2_THREADS +#define MK2_THREADS 256 +#endif +constexpr int THREADS = MK2_THREADS; +constexpr int NORM_THREADS = 256; +constexpr int NORM_WARPS = NORM_THREADS / 32; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +#ifndef MK2_GEMM_CONSUMERS +#define MK2_GEMM_CONSUMERS 1 +#endif +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 +#define MK2_HEAD_STAGE_BYTES 0 +#else +#define MK2_HEAD_STAGE_BYTES sizeof(out_tile_f) +#endif +constexpr int NCONS = MK2_GEMM_CONSUMERS; +static_assert(NCONS >= 1 && NCONS <= 3, "consumers occupy warpgroups 0, 2 and 3"); +static_assert(THREADS >= (NCONS + 1) * 128, "one producer plus NCONS consumers"); +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +constexpr int PAD_N = 32; +using act_tile_n = st_bf; +using out_tile_n = st_bf; +using act_gl_n = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_MOE_ACC_C, EPI_ACC_F32 }; +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + unsigned int *native_pf_expert_masks = nullptr; + float *pf_dtda = nullptr; + bf16 *pf_moe_gath = nullptr; + bf16 *pf_moe_hc = nullptr; + int *pf_moe_gidx = nullptr; + int *pf_moe_ecnt = nullptr; + int *pf_moe_ntiles = nullptr; + int pf_moe_estride = 0; + act_gl pf_moe_gath_v; + act_gl pf_moe_hc_v; + act_gl_n pf_moe_gath_vn; + act_gl_n pf_moe_hc_vn; + unsigned int *native_pf_cohort_check = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + volatile unsigned int *generation = state + 1; + const unsigned int seen = *generation; + if (atomicAdd(state, 1u) == gridDim.x - 1) { + atomicExch(state, 0u); + atomicAdd(state + 1, 1u); + } else { + while (*generation == seen) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const bf16 *in = src + (size_t)row * D_MODEL; + float v[PER_THREAD]; + float sum = 0.f, sq = 0.f; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (threadIdx.x < NORM_THREADS) { + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + v[j] = __bfloat162float(in[threadIdx.x + j * NORM_THREADS]); + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + if (lane == 0) { scratch[warp] = sum; scratch[NORM_WARPS + warp] = sq; } + } + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f, q = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) { s += scratch[w]; q += scratch[NORM_WARPS + w]; } + scratch[2 * NORM_WARPS] = s / D_MODEL; + scratch[2 * NORM_WARPS + 1] = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + } + __syncthreads(); + const float mean = scratch[2 * NORM_WARPS], inv = scratch[2 * NORM_WARPS + 1]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + if (threadIdx.x >= NORM_THREADS) break; + const int i = threadIdx.x + j * NORM_THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * weight[i]); + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, + int idx, unsigned mask); +__device__ __forceinline__ void norm_sub_sync(int sub) { + if (THREADS == NORM_THREADS) { __syncthreads(); return; } + asm volatile("bar.sync %0, %1;" :: "r"(6 + sub), "n"(NORM_THREADS) : "memory"); +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, const float *fold_src = nullptr, + const float *fold_routes = nullptr) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + constexpr int NSUB = THREADS / NORM_THREADS; + if (threadIdx.x >= NSUB * NORM_THREADS) return; + const int sub = threadIdx.x / NORM_THREADS; + const int tid = threadIdx.x - sub * NORM_THREADS; + float *const sc = scratch + sub * (NORM_WARPS + 1); + for (int row = blockIdx.x * NSUB + sub; row < n_rows; + row += gridDim.x * NSUB) { + const size_t base = (size_t)row * D_MODEL; + unsigned fold_mask = (1u << BB_EXPERTS) - 1u; + if (fold_src && fold_routes) { + fold_mask = 0u; + const float *rw = fold_routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) + if (rw[e] != 0.f) fold_mask |= 1u << e; + } + float v[PER_THREAD]; + float sq = 0.f; + const int warp = tid >> 5, lane = tid & 31; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + bf16 hv; + if (fold_src) { + hv = bb_moe_fold_deterministic(fold_src, (int)(base + i), fold_mask); + h[base + i] = hv; + } else { + hv = h[base + i]; + } + v[j] = __bfloat162float(hv) + __bfloat162float(r[base + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + if (lane == 0) sc[warp] = sq; + norm_sub_sync(sub); + if (tid == 0) { + float s = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) s += sc[w]; + sc[NORM_WARPS] = rsqrtf(s / D_MODEL + RMS_EPS); + } + norm_sub_sync(sub); + const float inv = sc[NORM_WARPS]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + out[base + i] = __float2bfloat16( + __bfloat162float(weight[i]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + } + norm_sub_sync(sub); + } +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + const int c = idx - row * MB_CONV_DIM; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + const float raw = __bfloat162float( + zxbcdt[(size_t)row * MB_IN_PROJ + MB_INNER + c]); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + xbc_act[(size_t)row * MB_CONV_DIM + c] = __float2bfloat16(acc); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + const int pairs = BATCH * MB_HEADS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_HEADS * split); + const int head = (pair / split) % MB_HEADS; + const int first_channel = (pair % split) * channels; + const int group = head / MB_HEADS_PER_GROUP; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + L.dt_bias[head]); + const float dA = __expf(L.A[head] * dt); + const float Dh = L.D[head]; + float *st = L.ssm_state + (size_t)slot_map[row] * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + for (int p = first_channel; p < first_channel + channels; ++p) { + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + const size_t at = (size_t)p * MB_STATE + n; + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + static_assert(MB_CONV_K == 4, "packed conv taps assume a 4-tap kernel"); + constexpr int CPT = 4; + static_assert(MB_CONV_DIM % CPT == 0, "channel vector width"); + constexpr int CGROUPS = MB_CONV_DIM / CPT; + const int total = n_tokens * CGROUPS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / CGROUPS; + const int c = (idx - t * CGROUPS) * CPT; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + float wv[CPT][MB_CONV_K]; + { + const float4 *wsrc = reinterpret_cast( + L.conv_w + (size_t)c * MB_CONV_K); + #pragma unroll + for (int j = 0; j < CPT / 2; ++j) { + const float4 w4 = wsrc[j]; + const unsigned uu0 = __float_as_uint(w4.x); + const unsigned uu1 = __float_as_uint(w4.y); + const unsigned uu2 = __float_as_uint(w4.z); + const unsigned uu3 = __float_as_uint(w4.w); + wv[2 * j][0] = __uint_as_float(uu0 << 16); + wv[2 * j][1] = __uint_as_float(uu0 & 0xffff0000u); + wv[2 * j][2] = __uint_as_float(uu1 << 16); + wv[2 * j][3] = __uint_as_float(uu1 & 0xffff0000u); + wv[2 * j + 1][0] = __uint_as_float(uu2 << 16); + wv[2 * j + 1][1] = __uint_as_float(uu2 & 0xffff0000u); + wv[2 * j + 1][2] = __uint_as_float(uu3 << 16); + wv[2 * j + 1][3] = __uint_as_float(uu3 & 0xffff0000u); + } + } + float acc[CPT]; + { + const uint2 b2 = *reinterpret_cast(L.conv_b + c); + acc[0] = __uint_as_float(b2.x << 16); + acc[1] = __uint_as_float(b2.x & 0xffff0000u); + acc[2] = __uint_as_float(b2.y << 16); + acc[3] = __uint_as_float(b2.y & 0xffff0000u); + } + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v[CPT]; + if (src < 0) { + #pragma unroll + for (int q = 0; q < CPT; ++q) + v[q] = __bfloat162float( + state[(size_t)q * (MB_CONV_K - 1) + pos + k]); + } else { + const uint2 p = *reinterpret_cast( + zxbcdt + (size_t)src * MB_IN_PROJ + MB_INNER + c); + v[0] = __uint_as_float(p.x << 16); + v[1] = __uint_as_float(p.x & 0xffff0000u); + v[2] = __uint_as_float(p.y << 16); + v[3] = __uint_as_float(p.y & 0xffff0000u); + } + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] += wv[q][k] * v[q]; + } + uint2 out; + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] = acc[q] / (1.f + __expf(-acc[q])); + out.x = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[0])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[1])) << 16); + out.y = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[2])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[3])) << 16); + *reinterpret_cast(xbc_act + (size_t)t * MB_CONV_DIM + c) = out; + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_dtda_prefill( + const MambaLayer &L, const bf16 *zxbcdt, float *dtda, int n_tokens) { + const int total = n_tokens * MB_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_HEADS; + const int head = idx - t * MB_HEADS; + const float dt = mamba_softplus( + __bfloat162float(zxbcdt[(size_t)t * MB_IN_PROJ + MB_INNER + + MB_CONV_DIM + head]) + L.dt_bias[head]); + dtda[(size_t)idx * 2] = dt; + dtda[(size_t)idx * 2 + 1] = __expf(L.A[head] * dt); + } +} +template +__device__ __forceinline__ void mamba_ssm_prefill_impl( + const MambaLayer &L, const bf16 *xbc_act, const float *dtda, bf16 *gated, + int n_pos, int n_seq, int slot, const int *seq_slot) { + constexpr int NS = MB_STATE / 32; + constexpr int PG = MB_HEAD_DIM / PW; + constexpr bool RS = (PW % 4) == 0; + constexpr int H1 = RS ? PW / 2 : 1; + constexpr int H2 = RS ? PW / 4 : 1; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const bool keep_hi = (lane & 16) != 0; + const bool keep_hi2 = (lane & 8) != 0; + const int warps_total = gridDim.x * WARPS; + const int units = n_seq * MB_HEADS * PG; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int pg = unit % PG; + const int rest = unit / PG; + const int head = rest % MB_HEADS; + const int seq = rest / MB_HEADS; + const int p0 = pg * PW; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Dh = L.D[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p0 * MB_STATE; + float sreg[PW][NS]; + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + sreg[q][j] = st[(size_t)q * MB_STATE + lane + j * 32]; + } + auto fetch = [&](int p, float (&oB)[NS], float (&oC)[NS], + float (&oX)[PW], float2 &od) { + const size_t tt = (size_t)p * n_seq + seq; + const bf16 *xa2 = xbc_act + tt * MB_CONV_DIM; + const bf16 *Bs = xa2 + MB_INNER + group * MB_STATE; + const bf16 *Cs = xa2 + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + od = *reinterpret_cast( + dtda + (tt * MB_HEADS + head) * 2); + #pragma unroll + for (int j = 0; j < NS; ++j) { + oB[j] = __bfloat162float(Bs[lane + j * 32]); + oC[j] = __bfloat162float(Cs[lane + j * 32]); + } + if constexpr (RS) { + const uint2 *xp = reinterpret_cast( + xa2 + head * MB_HEAD_DIM + p0); + #pragma unroll + for (int q = 0; q < PW; q += 4) { + const uint2 w = xp[q >> 2]; + oX[q] = __uint_as_float(w.x << 16); + oX[q + 1] = __uint_as_float(w.x & 0xffff0000u); + oX[q + 2] = __uint_as_float(w.y << 16); + oX[q + 3] = __uint_as_float(w.y & 0xffff0000u); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) + oX[q] = __bfloat162float(xa2[head * MB_HEAD_DIM + p0 + q]); + } + }; + float pB[NS], pC[NS], pX[PW]; + float2 pd; + fetch(0, pB, pC, pX, pd); + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const float dt = pd.x; + const float dA = pd.y; + float Bv[NS], Cv[NS]; + #pragma unroll + for (int j = 0; j < NS; ++j) { + Bv[j] = pB[j] * dt; + Cv[j] = pC[j]; + } + float acc[PW], xvs[PW]; + #pragma unroll + for (int q = 0; q < PW; ++q) xvs[q] = pX[q]; + if (pos + 1 < n_pos) fetch(pos + 1, pB, pC, pX, pd); + #pragma unroll + for (int q = 0; q < PW; ++q) { + const float xv = xvs[q]; + float a = 0.f; + #pragma unroll + for (int j = 0; j < NS; ++j) { + sreg[q][j] = sreg[q][j] * dA + Bv[j] * xv; + a += sreg[q][j] * Cv[j]; + } + acc[q] = a; + } + if constexpr (RS) { + float m1[H1], t1[H1], x1[H1]; + #pragma unroll + for (int q = 0; q < H1; ++q) { + m1[q] = keep_hi ? acc[q + H1] : acc[q]; + t1[q] = keep_hi ? acc[q] : acc[q + H1]; + x1[q] = keep_hi ? xvs[q + H1] : xvs[q]; + } + #pragma unroll + for (int q = 0; q < H1; ++q) + m1[q] += __shfl_xor_sync(0xffffffff, t1[q], 16); + float m2[H2], t2[H2], x2[H2]; + #pragma unroll + for (int q = 0; q < H2; ++q) { + m2[q] = keep_hi2 ? m1[q + H2] : m1[q]; + t2[q] = keep_hi2 ? m1[q] : m1[q + H2]; + x2[q] = keep_hi2 ? x1[q + H2] : x1[q]; + } + #pragma unroll + for (int q = 0; q < H2; ++q) + m2[q] += __shfl_xor_sync(0xffffffff, t2[q], 8); + #pragma unroll + for (int q = 0; q < H2; ++q) { + #pragma unroll + for (int off = 4; off > 0; off >>= 1) + m2[q] += __shfl_xor_sync(0xffffffff, m2[q], off); + } + if ((lane & 7) == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + + p0 + (keep_hi ? H1 : 0) + (keep_hi2 ? H2 : 0); + #pragma unroll + for (int q = 0; q < H2; ++q) + dst[q] = __float2bfloat16(m2[q] + x2[q] * Dh); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc[q] += __shfl_xor_sync(0xffffffff, acc[q], off); + } + if (lane == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + p0; + #pragma unroll + for (int q = 0; q < PW; ++q) + dst[q] = __float2bfloat16(acc[q] + xvs[q] * Dh); + } + } + } + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + st[(size_t)q * MB_STATE + lane + j * 32] = sreg[q][j]; + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, + const float *dtda, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + (void)zxbcdt; + if (n_seq > 1) + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens / n_seq, n_seq, slot, seq_slot); + else + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens, 1, slot, nullptr); +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16( + held[j] * inv * __bfloat162float(L.gnorm_w[idx])); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = BATCH * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / BB_ATT_HEADS; + const int head = pair % BB_ATT_HEADS; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const int kv = head / kv_per_q; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *q = base + head * BB_HEAD_DIM; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (head % kv_per_q == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + for (int tpos = 0; tpos < T; ++tpos) { +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new + : kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#endif + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m, s); + const float rescale = __expf(m - m_new); + const float w = __expf(s - m_new); + denom = denom * rescale + w; +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *vt = kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#else + const bf16 *vt = cur_token + ? v_new + : kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + w * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + constexpr int CHUNKS = BB_HEAD_DIM / 8; + const int total = n_tokens * BB_KV_HEADS * CHUNKS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int c = idx % CHUNKS; + const int rest = idx / CHUNKS; + const int t = rest / BB_KV_HEADS; + const int kv = rest - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM + c * 8; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM + c * 8; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv) + c * 8; + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv) + c * 8; + *reinterpret_cast(kc) = *reinterpret_cast(k_new); + *reinterpret_cast(vc) = *reinterpret_cast(v_new); + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float qv[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qv[j] = __bfloat162float(q[lane + j * 32]); + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + const int *const krow = + L.block_table + ((size_t)seq * 2 + 0) * L.max_blocks_per_seq; + const int *const vrow = + L.block_table + ((size_t)seq * 2 + 1) * L.max_blocks_per_seq; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + const int head_bias = kv * L.tokens_per_block * BB_HEAD_DIM; + int page = -1; + const bf16 *kpage = nullptr, *vpage = nullptr; + float kn[PER_LANE], vn[PER_LANE]; + { + page = 0; + kpage = L.pool_base + (size_t)(krow[0] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + (size_t)(vrow[0] & 0x7FFFFFFF) * elems_per_block; + const int at = head_bias + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + for (int s = 0; s <= pos; ++s) { + float kc[PER_LANE], vc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { kc[j] = kn[j]; vc[j] = vn[j]; } + if (s < pos) { + const int np = (s + 1) / L.tokens_per_block; + if (np != page) { + page = np; + kpage = L.pool_base + + (size_t)(krow[np] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + + (size_t)(vrow[np] & 0x7FFFFFFF) * elems_per_block; + } + const int at = head_bias + + ((s + 1) % L.tokens_per_block) * BB_HEAD_DIM + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) dot += qv[j] * kc[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[j] = acc[j] * rescale + wt * vc[j]; + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH, unsigned int *tile_masks = nullptr, + bf16 *gath = nullptr, int *gidx = nullptr, int *ecnt = nullptr, + int estride = 0, bf16 *row_stage = nullptr) { + __shared__ float score[BB_EXPERTS]; + __shared__ int s_pick[BB_TOP_K], s_slot[BB_TOP_K]; + constexpr int CH = D_MODEL / 8; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + { + const float4 *gsrc = + reinterpret_cast(normed + (size_t)row * D_MODEL); + float4 *gdst = reinterpret_cast(row_stage); + for (int i = threadIdx.x; i < CH; i += THREADS) gdst[i] = gsrc[i]; + } + if ((int)threadIdx.x < BB_EXPERTS) + moe_w[(size_t)row * BB_EXPERTS + threadIdx.x] = 0.f; + __syncthreads(); + for (int e = warp; e < BB_EXPERTS; e += WARPS) { + const bf16 *wr = gate_w + (size_t)e * D_MODEL; + float acc = 0.f; + for (int i = lane; i < D_MODEL; i += 32) + acc += __bfloat162float(row_stage[i]) * + __bfloat162float(wr[i]); + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) score[e] = 1.f / (1.f + __expf(-acc)); + } + __syncthreads(); + if (threadIdx.x == 0) { + float remaining[BB_EXPERTS]; + int pick[BB_TOP_K]; + for (int e = 0; e < BB_EXPERTS; ++e) + remaining[e] = score[e] + __bfloat162float(gate_bias[e]); + float total = 0.f; + for (int k = 0; k < BB_TOP_K; ++k) { + int best = 0; + for (int e = 1; e < BB_EXPERTS; ++e) + if (remaining[e] > remaining[best]) best = e; + pick[k] = best; + total += score[best]; + remaining[best] = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + unsigned int bits = 0u; + for (int k = 0; k < BB_TOP_K; ++k) { + moe_w[row * BB_EXPERTS + pick[k]] = score[pick[k]] * norm; + bits |= 1u << pick[k]; + s_pick[k] = pick[k]; + } + if (tile_masks) atomicOr(&tile_masks[row / BATCH_PAD], bits); + } + __syncthreads(); + if (gath) { + if ((int)threadIdx.x < BB_TOP_K) { + const int k = threadIdx.x; + const int slot = atomicAdd(&ecnt[s_pick[k]], 1); + s_slot[k] = slot; + gidx[(size_t)s_pick[k] * estride + slot] = row; + } + __syncthreads(); + const float4 *src = reinterpret_cast(row_stage); + for (int k = 0; k < BB_TOP_K; ++k) { + float4 *dst = reinterpret_cast( + gath + ((size_t)s_pick[k] * estride + s_slot[k]) * D_MODEL); + for (int i = threadIdx.x; i < CH; i += THREADS) dst[i] = src[i]; + } + } + __syncthreads(); + } +} +__device__ __forceinline__ unsigned moe_row_mask(const float *routes, int row) { + unsigned m = 0u; + const float *rw = routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) if (rw[e] != 0.f) m |= 1u << e; + return m; +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx, + unsigned mask) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + for (unsigned m = mask; m; m &= m - 1u) + value = __fadd_rn(value, parts[base + (size_t)(__ffs(m) - 1) * D_MODEL]); + return __float2bfloat16( + __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL])); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + uint64_t *keys, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int n = threadIdx.x; n < SORT_N; n += THREADS) { + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + keys[n] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + __syncthreads(); + for (int k = 2; k <= SORT_N; k <<= 1) { + if (k >= 64) __syncthreads(); + for (int j = k >> 1; j > 0; j >>= 1) { + for (int i = threadIdx.x; i < SORT_N; i += THREADS) { + const int partner = i ^ j; + if (partner > i) { + const bool descending = ((i & k) == 0); + if ((keys[i] < keys[partner]) == descending) { + const uint64_t t = keys[i]; + keys[i] = keys[partner]; + keys[partner] = t; + } + } + } + if (j >= 32) __syncthreads(); + else __syncwarp(); + } + } + __syncthreads(); + if (threadIdx.x == 0) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float v; + int n; + if (!eos_used && eos_v > av) { + v = eos_v; + n = AUDIO_EOS; + eos_used = true; + } else { + v = av; + n = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = + v + gumbel[((size_t)row * N_POS + pos) * TOP_K + j]; + if (noised > best_noised) { best_noised = noised; best_code = n; } + } + out_codes[row] = best_code; + } + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void prefill_expert_masks(const float *routes, unsigned int *masks, int rows) { + const int lane=threadIdx.x&31, warp=threadIdx.x>>5; + const int nt=(rows+BATCH_PAD-1)/BATCH_PAD; + for (int t=blockIdx.x*WARPS+warp;t(gath + g * D_MODEL + c * 8); + if (local < ecnt[e]) + *dst = *reinterpret_cast( + normed + (size_t)gidx[g] * D_MODEL + c * 8); + else + *dst = make_float4(0.f, 0.f, 0.f, 0.f); + } +} +__device__ __forceinline__ int gemm_eff(int tiles, int reps) { + if (NCONS <= 1) return 1; + const int g = (int)gridDim.x; + const int rounds1 = (tiles * reps + g - 1) / g; + if constexpr (NCONS == 3) { + const int max_eff = reps <= BB_EXPERTS ? 2 : 3; + const int rounds2 = (((tiles + 1) / 2) * reps + g - 1) / g; + if (max_eff == 2) return rounds2 < rounds1 ? 2 : 1; + const int rounds3 = (((tiles + 2) / 3) * reps + g - 1) / g; + if (rounds3 < rounds2 && rounds3 < rounds1) return 3; + return rounds2 < rounds1 ? 2 : 1; + } else { + const int roundsN = (((tiles + NCONS - 1) / NCONS) * reps + g - 1) / g; + return roundsN < rounds1 ? NCONS : 1; + } +} +__device__ __forceinline__ int gemm_groups(int tiles, int reps) { + const int eff = gemm_eff(tiles, reps); + return (tiles + eff - 1) / eff; +} +__device__ __forceinline__ int gemm_consumer_id() { + const int g = warpgroup::groupid(); + if (g == 1) return -1; + const int c = g == 0 ? 0 : g - 1; + return c < NCONS ? c : -1; +} +__device__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][BATCH_PAD]; + __shared__ int moe_gidx_stage[NCONS][BATCH_PAD]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + a_s[s], a, {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + a_s[s], a, {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / BATCH_PAD) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, BATCH_PAD> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], a_s[s]); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * BATCH_PAD; + if (epi == EPI_BIAS) { + warpgroup::store(o_sf, accum); + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + dest[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(o_sf[int2{col, row}] + + __bfloat162float(bias[base + col])); + } + warpgroup::sync(1 + cons); + continue; + } + warpgroup::store(o_s[cons], accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_s[cons][int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__device__ void gemm_narrow( + const weight_gl &w, const act_gl_n &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = PAD_N, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][PAD_N]; + __shared__ int moe_gidx_stage[NCONS][PAD_N]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + PAD_N - 1) / PAD_N; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + PAD_N - 1) / PAD_N; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / PAD_N) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, PAD_N> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], + *reinterpret_cast(&a_s[s])); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * PAD_N; + out_tile_n &o_sn = *reinterpret_cast(&o_s[cons]); + warpgroup::store(o_sn, accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_sn[int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&w_s)[STAGES * NCONS] = + alloc.allocate(); + act_tile (&a_s)[STAGES] = alloc.allocate(); + out_tile (&o_s)[NCONS] = alloc.allocate(); +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 + static_assert(NCONS * sizeof(out_tile) >= sizeof(out_tile_f), + "aliased head staging must fit the output tiles"); + out_tile_f &o_sf = *reinterpret_cast(&o_s[0]); +#else + out_tile_f (&o_sf) = alloc.allocate(); +#endif + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_WARPS + 2]; +#if MK2_PREFILL_ONLY + bf16 *const route_stage = reinterpret_cast(&o_s[0]); + static_assert(sizeof(o_s) >= D_MODEL * sizeof(bf16), + "router row stage must fit the output tiles"); +#else + __shared__ __align__(16) uint64_t sort_keys[SORT_N]; + bf16 *const route_stage = reinterpret_cast(&sort_keys[0]); + static_assert(sizeof(sort_keys) >= D_MODEL * sizeof(bf16), + "router row stage must fit the sampler workspace"); +#endif + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS * NCONS, 0); + for (int c = 0; c < CONSUMER_WARPS * NCONS; ++c) arrive(drained[s]); + } + } + int gemm_iter = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map); + grid_barrier(g.barrier); + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map); + grid_barrier(g.barrier); + mamba_gated_norm(L, bb->zxbcdt, bb->gated); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, nullptr, 1, BATCH, + bb->moe_w, BB_MOE_INTER / OUT_ATOM); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + 1, BATCH, bb->moe_w, D_MODEL / OUT_ATOM); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_deterministic(bb->moe_out, idx, + (1u << BB_EXPERTS) - 1u); + grid_barrier(g.barrier); + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int original_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + bool identical_cohort = false; + if constexpr (BATCH > 1) { + if (bb->native_pf_cohort_check && bb->native_prefill_reset && + bb->pf_n_seq == BATCH && !bb->pf_fuse_decode && + !bb->pf_layer_out && !bb->moe_w_inject && !bb->moe_w_audit) { + if (blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[0] = 0; + grid_barrier(g.barrier); + const auto *bits = reinterpret_cast(bb->prompt); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < original_total * D_MODEL; idx += gridDim.x * THREADS) { + const int canonical = (idx / (BATCH * D_MODEL)) * BATCH * D_MODEL + idx % D_MODEL; + if (bits[idx] != bits[canonical]) atomicOr(bb->native_pf_cohort_check, 1u); + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + if (bb->pf_seq_slot[row] != row || bb->slot_map[row] != row) + atomicOr(bb->native_pf_cohort_check, 1u); + grid_barrier(g.barrier); + identical_cohort = bb->native_pf_cohort_check[0] == 0; + } + } + const int n_total = identical_cohort ? original_total / BATCH : original_total; + if (bb->native_pf_cohort_check && blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[1] = identical_cohort ? 1 : BATCH; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int reset_slots = identical_cohort ? 1 : BATCH; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * reset_slots * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % reset_slots; + const MambaLayer &L = bb->mamba[row_layer / reset_slots]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * reset_slots * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % reset_slots; + const MambaLayer &L = bb->mamba[row_layer / reset_slots]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = original_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + const int src = identical_cohort ? (idx / D_MODEL) * BATCH * D_MODEL + idx % D_MODEL : idx; + bb->pf_h[idx] = bb->prompt[src]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = identical_cohort ? 1 : (bb->pf_n_seq > 0 ? bb->pf_n_seq : 1); + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T, + (!bb->pf_layer_out && layer > 0 && + bb->moe_of_layer[layer - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + if (bb->native_pf_expert_masks) + for (int i = blockIdx.x * THREADS + threadIdx.x; i < n_tok; + i += gridDim.x * THREADS) + bb->native_pf_expert_masks[i] = 0u; + if (bb->pf_moe_ecnt) + for (int e = blockIdx.x * THREADS + threadIdx.x; e < BB_EXPERTS; + e += gridDim.x * THREADS) + bb->pf_moe_ecnt[e] = 0; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + mamba_dtda_prefill(L, bb->pf_zxbcdt, bb->pf_dtda, T); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_dtda, + bb->pf_gated, T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, + T, bb->native_pf_expert_masks, bb->pf_moe_gath, + bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, route_stage); + } + grid_barrier(g.barrier); + int max_ecnt = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + max_ecnt = max(max_ecnt, bb->pf_moe_ecnt[e]); + const bool narrow = BATCH_PAD > PAD_N && max_ecnt <= PAD_N; + const int etile = narrow ? PAD_N : BATCH_PAD; + int nt_sum = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_sum += (bb->pf_moe_ecnt[e] + etile - 1) / etile; + const int up_steps = + gemm_groups(BB_MOE_INTER / OUT_ATOM, nt_sum) * + nt_sum * (D_MODEL / KCHUNK); + const int shup_steps = + gemm_groups(BB_SHARED_INTER / OUT_ATOM, n_tok) * + n_tok * (D_MODEL / KCHUNK); + const int down_steps = + gemm_groups(D_MODEL / OUT_ATOM, nt_sum) * + nt_sum * (BB_MOE_INTER / KCHUNK); + const int shdown_steps = + gemm_groups(D_MODEL / OUT_ATOM, n_tok) * + n_tok * (BB_SHARED_INTER / KCHUNK); + if (narrow) + gemm_narrow(L.exp_up, bb->pf_moe_gath_vn, + BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, PAD_N, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + else + gemm(L.exp_up, bb->pf_moe_gath_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, BATCH_PAD, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + up_steps, up_steps + shup_steps); + grid_barrier(g.barrier); + if (narrow) + gemm_narrow(L.exp_down, bb->pf_moe_hc_vn, + BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, PAD_N, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + else + gemm(L.exp_down, bb->pf_moe_hc_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + down_steps, down_steps + shdown_steps); + grid_barrier(g.barrier); + if (bb->pf_layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * D_MODEL; idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = + bb_moe_fold_deterministic(bb->pf_moe_out, idx, + moe_row_mask(bb->pf_moe_w, idx / D_MODEL)); + grid_barrier(g.barrier); + } + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, + identical_cohort ? bb->pf_normed : bb->pf_final, ln_scratch, + n_total, + (!bb->pf_layer_out && n_backbone_layers > 0 && + bb->moe_of_layer[n_backbone_layers - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + grid_barrier(g.barrier); + if (identical_cohort) { + const int worker = blockIdx.x * THREADS + threadIdx.x; + const int workers = gridDim.x * THREADS; + for (int i = worker; i < original_total * D_MODEL; i += workers) + bb->pf_final[i] = bb->pf_normed[(i / (BATCH * D_MODEL)) * D_MODEL + i % D_MODEL]; + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = worker; i < 14 * (BATCH - 1) * CV; i += workers) { + const int cell = i % CV, pair = i / CV; + const int row = pair % (BATCH - 1) + 1; + const MambaLayer &L = bb->mamba[pair / (BATCH - 1)]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = L.conv_state[cell]; + } + for (int i = worker; i < 14 * (BATCH - 1) * SS; i += workers) { + const int cell = i % SS, pair = i / SS; + const int row = pair % (BATCH - 1) + 1; + const MambaLayer &L = bb->mamba[pair / (BATCH - 1)]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = L.ssm_state[cell]; + } + const int kv_columns = BB_KV_HEADS * BB_HEAD_DIM; + for (int i = worker; i < 3 * 2 * (BATCH - 1) * n_total * kv_columns; i += workers) { + int q = i / kv_columns; + const int col = i % kv_columns; + const int pos = q % n_total; q /= n_total; + const int row = q % (BATCH - 1) + 1; q /= BATCH - 1; + const int kind = q % 2, layer = q / 2; + const AttnLayer &L = bb->attn[layer]; + const int offset = kv_offset(L, pos, col / BB_HEAD_DIM) + col % BB_HEAD_DIM; + kv_slot(L, row, kind, pos)[offset] = kv_slot(L, 0, kind, pos)[offset]; + } + grid_barrier(g.barrier); + } + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } +#if MK2_PREFILL_ONLY + return; +#else + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, bb->phoneme_b); + grid_barrier(g.barrier); + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + if (threadIdx.x == 0) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = 0; n < BB_PHONEME; ++n) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + bb->phoneme_out[row] = arg; + } + __syncthreads(); + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + for (int i = lane; i < D_HEAD; i += 32) { + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + for (int t = 0; t <= pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + for (int i = lane; i < D_HEAD; i += 32) { + float acc = 0.f; + for (int t = 0; t <= pos; ++t) + acc += sc[t] * __bfloat162float( + g.vcache[cache_base + t * D_HEAD + i]); + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + i] = + __float2bfloat16(acc); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, g.sampled_pos, + sort_keys, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { +#ifdef ABLATE_SAMPLE + __syncthreads(); +#endif + } else { + grid_barrier(g.barrier); + } + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x < BATCH && blockIdx.x == 0) { + const int code = g.sampled_pos[threadIdx.x]; + g.sampled[threadIdx.x * N_POS + pos] = code; + if (g.native_codes) + g.native_codes[threadIdx.x * N_POS + pos] = (long long)code; + } + grid_barrier(g.barrier); + } + if (g.native_seq_len && blockIdx.x == 0 && threadIdx.x < BATCH) + g.native_seq_len[threadIdx.x] += 1; + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +#endif +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 2 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK), + 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = MK2_PREFILL_ONLY + ? out->sms * out->blocks_per_sm + : std::min(out->sms * out->blocks_per_sm, D_FFN / OUT_ATOM); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + 0, + nullptr, + nullptr, + nullptr, nullptr, nullptr, nullptr, nullptr, 0, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "dedup_native_prefill_impl.inc" diff --git a/deploy/nemotron_native/kernels/prefill/b64/mk2_abi.cuh b/deploy/nemotron_native/kernels/prefill/b64/mk2_abi.cuh new file mode 100644 index 0000000..95a253b --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b64/mk2_abi.cuh @@ -0,0 +1,82 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} diff --git a/deploy/nemotron_native/kernels/prefill/b64/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/prefill/b64/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b64/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/prefill/b64/prefill.cu b/deploy/nemotron_native/kernels/prefill/b64/prefill.cu new file mode 100644 index 0000000..3369db2 --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b64/prefill.cu @@ -0,0 +1,18 @@ +#define MK2_BATCH 64 +#define MK2_BATCH_TILE 64 +#define MK2_PREFILL_ONLY 1 +#define STAGES 3 +#define KCHUNK 128 +#define MK2_GEMM_CONSUMERS 3 +#define MK2_THREADS 512 +#define MK2_NO_MAIN 1 +#define MK2_DIRECT_INPUT 1 +#define MK2_B1_LOCAL_BARRIERS 1 +#define MK2_SKIP_HIDDEN_SNAPSHOT 1 +#define MK2_ACTIVE_EXPERTS 1 +#define MK2_SSM_SPLITS 1 +#define PF_SSM_PW_MULTI 12 +#define PF_ACTIVE_TOKEN_EXPERTS 1 +#include +#include +#include "dedup_predictor_mk2.cuh" diff --git a/deploy/nemotron_native/kernels/prefill/b8/dedup_native_prefill_impl.inc b/deploy/nemotron_native/kernels/prefill/b8/dedup_native_prefill_impl.inc new file mode 100644 index 0000000..f7e9802 --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b8/dedup_native_prefill_impl.inc @@ -0,0 +1,237 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 2 || BATCH == 4 || BATCH == 8 || BATCH == 16 || BATCH == 32 || BATCH == 64)) + throw std::runtime_error("native prefill requires a matching supported physical batch"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (padded > INT_MAX / ((BB_EXPERTS + 1) * D_MODEL)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + float *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(float))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(float))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_dtda, 2 * MB_HEADS); + { + const size_t estride = padded; + const size_t grows = size_t(BB_EXPERTS) * estride; + mk2_check(owned.alloc(&hb.pf_moe_gath, grows * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_gath, 0, grows * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_hc, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(cudaMemset(hb.pf_moe_hc, 0, grows * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&hb.pf_moe_gidx, grows * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_gidx, 0, grows * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ecnt, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ecnt, 0, BB_EXPERTS * sizeof(int))); + mk2_check(owned.alloc(&hb.pf_moe_ntiles, BB_EXPERTS * sizeof(int))); + mk2_check(cudaMemset(hb.pf_moe_ntiles, 0, BB_EXPERTS * sizeof(int))); + hb.pf_moe_estride = int(estride); + hb.pf_moe_gath_v = act_gl{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_v = act_gl{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + hb.pf_moe_gath_vn = act_gl_n{hb.pf_moe_gath, nullptr, nullptr, grows, + size_t(D_MODEL)}; + hb.pf_moe_hc_vn = act_gl_n{hb.pf_moe_hc, nullptr, nullptr, grows, + size_t(BB_MOE_INTER)}; + scratch_bytes += grows * (D_MODEL + BB_MOE_INTER) * sizeof(bf16); + } + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(owned.alloc(&hb.native_pf_cohort_check, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(hb.native_pf_cohort_check, 0, 2 * sizeof(unsigned int))); +#if PF_ACTIVE_TOKEN_EXPERTS + mk2_check(owned.alloc(&hb.native_pf_expert_masks, (padded/BATCH_PAD)*sizeof(unsigned int))); +#endif + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} +extern "C" int mk2_native_prefill_effective_sequences(MK2NativePrefill *p, int *out) { + try { + if (!p || !out) throw std::runtime_error("null effective-sequence query"); + auto *st = reinterpret_cast(p->base); + mk2_check(cudaMemcpy(out, st->host_bb.native_pf_cohort_check + 1, + sizeof(int), cudaMemcpyDeviceToHost)); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} diff --git a/deploy/nemotron_native/kernels/prefill/b8/dedup_predictor_mk2.cuh b/deploy/nemotron_native/kernels/prefill/b8/dedup_predictor_mk2.cuh new file mode 100644 index 0000000..a640917 --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b8/dedup_predictor_mk2.cuh @@ -0,0 +1,4588 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef PF_ACTIVE_TOKEN_EXPERTS +#define PF_ACTIVE_TOKEN_EXPERTS 0 +#endif +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +#ifndef MK2_PREFILL_ONLY +#define MK2_PREFILL_ONLY 0 +#endif +#ifndef MK2_BATCH_TILE +#define MK2_BATCH_TILE (((MK2_BATCH + 15) / 16) * 16) +#endif +constexpr int BATCH_PAD = MK2_BATCH_TILE; +static_assert(BATCH >= 1 && BATCH <= BATCH_PAD); +static_assert(BATCH_PAD == 16 || BATCH_PAD == 32 || BATCH_PAD == 64 || BATCH_PAD == 128, + "native prefill tile must be 16, 32, 64 or 128"); +static_assert(MK2_PREFILL_ONLY || BATCH_PAD == ((BATCH + 15) / 16) * 16, + "a wider token tile requires a prefill-only library"); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_PW_MULTI +#define PF_SSM_PW_MULTI 4 +#endif +#ifndef PF_SSM_PW_SINGLE +#define PF_SSM_PW_SINGLE 2 +#endif +static_assert(MB_HEAD_DIM % PF_SSM_PW_MULTI == 0 && + MB_HEAD_DIM % PF_SSM_PW_SINGLE == 0, + "head_dim must split evenly into scan channel groups"); +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef MK2_SMEM_SLACK +#define MK2_SMEM_SLACK 1024 +#endif +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +#ifndef MK2_THREADS +#define MK2_THREADS 256 +#endif +constexpr int THREADS = MK2_THREADS; +constexpr int NORM_THREADS = 256; +constexpr int NORM_WARPS = NORM_THREADS / 32; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +#ifndef MK2_GEMM_CONSUMERS +#define MK2_GEMM_CONSUMERS 1 +#endif +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 +#define MK2_HEAD_STAGE_BYTES 0 +#else +#define MK2_HEAD_STAGE_BYTES sizeof(out_tile_f) +#endif +constexpr int NCONS = MK2_GEMM_CONSUMERS; +static_assert(NCONS >= 1 && NCONS <= 3, "consumers occupy warpgroups 0, 2 and 3"); +static_assert(THREADS >= (NCONS + 1) * 128, "one producer plus NCONS consumers"); +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +constexpr int PAD_N = 32; +using act_tile_n = st_bf; +using out_tile_n = st_bf; +using act_gl_n = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_MOE_ACC_C, EPI_ACC_F32 }; +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + unsigned int *native_pf_expert_masks = nullptr; + float *pf_dtda = nullptr; + bf16 *pf_moe_gath = nullptr; + bf16 *pf_moe_hc = nullptr; + int *pf_moe_gidx = nullptr; + int *pf_moe_ecnt = nullptr; + int *pf_moe_ntiles = nullptr; + int pf_moe_estride = 0; + act_gl pf_moe_gath_v; + act_gl pf_moe_hc_v; + act_gl_n pf_moe_gath_vn; + act_gl_n pf_moe_hc_vn; + unsigned int *native_pf_cohort_check = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + volatile unsigned int *generation = state + 1; + const unsigned int seen = *generation; + if (atomicAdd(state, 1u) == gridDim.x - 1) { + atomicExch(state, 0u); + atomicAdd(state + 1, 1u); + } else { + while (*generation == seen) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const bf16 *in = src + (size_t)row * D_MODEL; + float v[PER_THREAD]; + float sum = 0.f, sq = 0.f; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (threadIdx.x < NORM_THREADS) { + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + v[j] = __bfloat162float(in[threadIdx.x + j * NORM_THREADS]); + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + if (lane == 0) { scratch[warp] = sum; scratch[NORM_WARPS + warp] = sq; } + } + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f, q = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) { s += scratch[w]; q += scratch[NORM_WARPS + w]; } + scratch[2 * NORM_WARPS] = s / D_MODEL; + scratch[2 * NORM_WARPS + 1] = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + } + __syncthreads(); + const float mean = scratch[2 * NORM_WARPS], inv = scratch[2 * NORM_WARPS + 1]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + if (threadIdx.x >= NORM_THREADS) break; + const int i = threadIdx.x + j * NORM_THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * weight[i]); + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, + int idx, unsigned mask); +__device__ __forceinline__ void norm_sub_sync(int sub) { + if (THREADS == NORM_THREADS) { __syncthreads(); return; } + asm volatile("bar.sync %0, %1;" :: "r"(6 + sub), "n"(NORM_THREADS) : "memory"); +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, const float *fold_src = nullptr, + const float *fold_routes = nullptr) { + constexpr int PER_THREAD = D_MODEL / NORM_THREADS; + constexpr int NSUB = THREADS / NORM_THREADS; + if (threadIdx.x >= NSUB * NORM_THREADS) return; + const int sub = threadIdx.x / NORM_THREADS; + const int tid = threadIdx.x - sub * NORM_THREADS; + float *const sc = scratch + sub * (NORM_WARPS + 1); + for (int row = blockIdx.x * NSUB + sub; row < n_rows; + row += gridDim.x * NSUB) { + const size_t base = (size_t)row * D_MODEL; + unsigned fold_mask = (1u << BB_EXPERTS) - 1u; + if (fold_src && fold_routes) { + fold_mask = 0u; + const float *rw = fold_routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) + if (rw[e] != 0.f) fold_mask |= 1u << e; + } + float v[PER_THREAD]; + float sq = 0.f; + const int warp = tid >> 5, lane = tid & 31; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + bf16 hv; + if (fold_src) { + hv = bb_moe_fold_deterministic(fold_src, (int)(base + i), fold_mask); + h[base + i] = hv; + } else { + hv = h[base + i]; + } + v[j] = __bfloat162float(hv) + __bfloat162float(r[base + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + if (lane == 0) sc[warp] = sq; + norm_sub_sync(sub); + if (tid == 0) { + float s = 0.f; + for (int w = 0; w < NORM_WARPS; ++w) s += sc[w]; + sc[NORM_WARPS] = rsqrtf(s / D_MODEL + RMS_EPS); + } + norm_sub_sync(sub); + const float inv = sc[NORM_WARPS]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = tid + j * NORM_THREADS; + out[base + i] = __float2bfloat16( + __bfloat162float(weight[i]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + } + norm_sub_sync(sub); + } +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + const int c = idx - row * MB_CONV_DIM; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + const float raw = __bfloat162float( + zxbcdt[(size_t)row * MB_IN_PROJ + MB_INNER + c]); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + xbc_act[(size_t)row * MB_CONV_DIM + c] = __float2bfloat16(acc); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + const int pairs = BATCH * MB_HEADS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_HEADS * split); + const int head = (pair / split) % MB_HEADS; + const int first_channel = (pair % split) * channels; + const int group = head / MB_HEADS_PER_GROUP; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + L.dt_bias[head]); + const float dA = __expf(L.A[head] * dt); + const float Dh = L.D[head]; + float *st = L.ssm_state + (size_t)slot_map[row] * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + for (int p = first_channel; p < first_channel + channels; ++p) { + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + const size_t at = (size_t)p * MB_STATE + n; + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + static_assert(MB_CONV_K == 4, "packed conv taps assume a 4-tap kernel"); + constexpr int CPT = 4; + static_assert(MB_CONV_DIM % CPT == 0, "channel vector width"); + constexpr int CGROUPS = MB_CONV_DIM / CPT; + const int total = n_tokens * CGROUPS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / CGROUPS; + const int c = (idx - t * CGROUPS) * CPT; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + float wv[CPT][MB_CONV_K]; + { + const float4 *wsrc = reinterpret_cast( + L.conv_w + (size_t)c * MB_CONV_K); + #pragma unroll + for (int j = 0; j < CPT / 2; ++j) { + const float4 w4 = wsrc[j]; + const unsigned uu0 = __float_as_uint(w4.x); + const unsigned uu1 = __float_as_uint(w4.y); + const unsigned uu2 = __float_as_uint(w4.z); + const unsigned uu3 = __float_as_uint(w4.w); + wv[2 * j][0] = __uint_as_float(uu0 << 16); + wv[2 * j][1] = __uint_as_float(uu0 & 0xffff0000u); + wv[2 * j][2] = __uint_as_float(uu1 << 16); + wv[2 * j][3] = __uint_as_float(uu1 & 0xffff0000u); + wv[2 * j + 1][0] = __uint_as_float(uu2 << 16); + wv[2 * j + 1][1] = __uint_as_float(uu2 & 0xffff0000u); + wv[2 * j + 1][2] = __uint_as_float(uu3 << 16); + wv[2 * j + 1][3] = __uint_as_float(uu3 & 0xffff0000u); + } + } + float acc[CPT]; + { + const uint2 b2 = *reinterpret_cast(L.conv_b + c); + acc[0] = __uint_as_float(b2.x << 16); + acc[1] = __uint_as_float(b2.x & 0xffff0000u); + acc[2] = __uint_as_float(b2.y << 16); + acc[3] = __uint_as_float(b2.y & 0xffff0000u); + } + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v[CPT]; + if (src < 0) { + #pragma unroll + for (int q = 0; q < CPT; ++q) + v[q] = __bfloat162float( + state[(size_t)q * (MB_CONV_K - 1) + pos + k]); + } else { + const uint2 p = *reinterpret_cast( + zxbcdt + (size_t)src * MB_IN_PROJ + MB_INNER + c); + v[0] = __uint_as_float(p.x << 16); + v[1] = __uint_as_float(p.x & 0xffff0000u); + v[2] = __uint_as_float(p.y << 16); + v[3] = __uint_as_float(p.y & 0xffff0000u); + } + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] += wv[q][k] * v[q]; + } + uint2 out; + #pragma unroll + for (int q = 0; q < CPT; ++q) acc[q] = acc[q] / (1.f + __expf(-acc[q])); + out.x = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[0])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[1])) << 16); + out.y = (unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[2])) | + ((unsigned)__bfloat16_as_ushort(__float2bfloat16(acc[3])) << 16); + *reinterpret_cast(xbc_act + (size_t)t * MB_CONV_DIM + c) = out; + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_dtda_prefill( + const MambaLayer &L, const bf16 *zxbcdt, float *dtda, int n_tokens) { + const int total = n_tokens * MB_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_HEADS; + const int head = idx - t * MB_HEADS; + const float dt = mamba_softplus( + __bfloat162float(zxbcdt[(size_t)t * MB_IN_PROJ + MB_INNER + + MB_CONV_DIM + head]) + L.dt_bias[head]); + dtda[(size_t)idx * 2] = dt; + dtda[(size_t)idx * 2 + 1] = __expf(L.A[head] * dt); + } +} +template +__device__ __forceinline__ void mamba_ssm_prefill_impl( + const MambaLayer &L, const bf16 *xbc_act, const float *dtda, bf16 *gated, + int n_pos, int n_seq, int slot, const int *seq_slot) { + constexpr int NS = MB_STATE / 32; + constexpr int PG = MB_HEAD_DIM / PW; + constexpr bool RS = (PW % 4) == 0; + constexpr int H1 = RS ? PW / 2 : 1; + constexpr int H2 = RS ? PW / 4 : 1; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const bool keep_hi = (lane & 16) != 0; + const bool keep_hi2 = (lane & 8) != 0; + const int warps_total = gridDim.x * WARPS; + const int units = n_seq * MB_HEADS * PG; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int pg = unit % PG; + const int rest = unit / PG; + const int head = rest % MB_HEADS; + const int seq = rest / MB_HEADS; + const int p0 = pg * PW; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Dh = L.D[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p0 * MB_STATE; + float sreg[PW][NS]; + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + sreg[q][j] = st[(size_t)q * MB_STATE + lane + j * 32]; + } + auto fetch = [&](int p, float (&oB)[NS], float (&oC)[NS], + float (&oX)[PW], float2 &od) { + const size_t tt = (size_t)p * n_seq + seq; + const bf16 *xa2 = xbc_act + tt * MB_CONV_DIM; + const bf16 *Bs = xa2 + MB_INNER + group * MB_STATE; + const bf16 *Cs = xa2 + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + od = *reinterpret_cast( + dtda + (tt * MB_HEADS + head) * 2); + #pragma unroll + for (int j = 0; j < NS; ++j) { + oB[j] = __bfloat162float(Bs[lane + j * 32]); + oC[j] = __bfloat162float(Cs[lane + j * 32]); + } + if constexpr (RS) { + const uint2 *xp = reinterpret_cast( + xa2 + head * MB_HEAD_DIM + p0); + #pragma unroll + for (int q = 0; q < PW; q += 4) { + const uint2 w = xp[q >> 2]; + oX[q] = __uint_as_float(w.x << 16); + oX[q + 1] = __uint_as_float(w.x & 0xffff0000u); + oX[q + 2] = __uint_as_float(w.y << 16); + oX[q + 3] = __uint_as_float(w.y & 0xffff0000u); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) + oX[q] = __bfloat162float(xa2[head * MB_HEAD_DIM + p0 + q]); + } + }; + float pB[NS], pC[NS], pX[PW]; + float2 pd; + fetch(0, pB, pC, pX, pd); + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const float dt = pd.x; + const float dA = pd.y; + float Bv[NS], Cv[NS]; + #pragma unroll + for (int j = 0; j < NS; ++j) { + Bv[j] = pB[j] * dt; + Cv[j] = pC[j]; + } + float acc[PW], xvs[PW]; + #pragma unroll + for (int q = 0; q < PW; ++q) xvs[q] = pX[q]; + if (pos + 1 < n_pos) fetch(pos + 1, pB, pC, pX, pd); + #pragma unroll + for (int q = 0; q < PW; ++q) { + const float xv = xvs[q]; + float a = 0.f; + #pragma unroll + for (int j = 0; j < NS; ++j) { + sreg[q][j] = sreg[q][j] * dA + Bv[j] * xv; + a += sreg[q][j] * Cv[j]; + } + acc[q] = a; + } + if constexpr (RS) { + float m1[H1], t1[H1], x1[H1]; + #pragma unroll + for (int q = 0; q < H1; ++q) { + m1[q] = keep_hi ? acc[q + H1] : acc[q]; + t1[q] = keep_hi ? acc[q] : acc[q + H1]; + x1[q] = keep_hi ? xvs[q + H1] : xvs[q]; + } + #pragma unroll + for (int q = 0; q < H1; ++q) + m1[q] += __shfl_xor_sync(0xffffffff, t1[q], 16); + float m2[H2], t2[H2], x2[H2]; + #pragma unroll + for (int q = 0; q < H2; ++q) { + m2[q] = keep_hi2 ? m1[q + H2] : m1[q]; + t2[q] = keep_hi2 ? m1[q] : m1[q + H2]; + x2[q] = keep_hi2 ? x1[q + H2] : x1[q]; + } + #pragma unroll + for (int q = 0; q < H2; ++q) + m2[q] += __shfl_xor_sync(0xffffffff, t2[q], 8); + #pragma unroll + for (int q = 0; q < H2; ++q) { + #pragma unroll + for (int off = 4; off > 0; off >>= 1) + m2[q] += __shfl_xor_sync(0xffffffff, m2[q], off); + } + if ((lane & 7) == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + + p0 + (keep_hi ? H1 : 0) + (keep_hi2 ? H2 : 0); + #pragma unroll + for (int q = 0; q < H2; ++q) + dst[q] = __float2bfloat16(m2[q] + x2[q] * Dh); + } + } else { + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc[q] += __shfl_xor_sync(0xffffffff, acc[q], off); + } + if (lane == 0) { + bf16 *const dst = gated + t * MB_INNER + head * MB_HEAD_DIM + p0; + #pragma unroll + for (int q = 0; q < PW; ++q) + dst[q] = __float2bfloat16(acc[q] + xvs[q] * Dh); + } + } + } + #pragma unroll + for (int q = 0; q < PW; ++q) { + #pragma unroll + for (int j = 0; j < NS; ++j) + st[(size_t)q * MB_STATE + lane + j * 32] = sreg[q][j]; + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, + const float *dtda, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + (void)zxbcdt; + if (n_seq > 1) + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens / n_seq, n_seq, slot, seq_slot); + else + mamba_ssm_prefill_impl( + L, xbc_act, dtda, gated, n_tokens, 1, slot, nullptr); +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16( + held[j] * inv * __bfloat162float(L.gnorm_w[idx])); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = BATCH * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / BB_ATT_HEADS; + const int head = pair % BB_ATT_HEADS; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const int kv = head / kv_per_q; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *q = base + head * BB_HEAD_DIM; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (head % kv_per_q == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + for (int tpos = 0; tpos < T; ++tpos) { +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new + : kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#endif + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m, s); + const float rescale = __expf(m - m_new); + const float w = __expf(s - m_new); + denom = denom * rescale + w; +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *vt = kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#else + const bf16 *vt = cur_token + ? v_new + : kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + w * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + constexpr int CHUNKS = BB_HEAD_DIM / 8; + const int total = n_tokens * BB_KV_HEADS * CHUNKS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int c = idx % CHUNKS; + const int rest = idx / CHUNKS; + const int t = rest / BB_KV_HEADS; + const int kv = rest - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM + c * 8; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM + c * 8; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv) + c * 8; + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv) + c * 8; + *reinterpret_cast(kc) = *reinterpret_cast(k_new); + *reinterpret_cast(vc) = *reinterpret_cast(v_new); + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float qv[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qv[j] = __bfloat162float(q[lane + j * 32]); + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + const int *const krow = + L.block_table + ((size_t)seq * 2 + 0) * L.max_blocks_per_seq; + const int *const vrow = + L.block_table + ((size_t)seq * 2 + 1) * L.max_blocks_per_seq; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + const int head_bias = kv * L.tokens_per_block * BB_HEAD_DIM; + int page = -1; + const bf16 *kpage = nullptr, *vpage = nullptr; + float kn[PER_LANE], vn[PER_LANE]; + { + page = 0; + kpage = L.pool_base + (size_t)(krow[0] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + (size_t)(vrow[0] & 0x7FFFFFFF) * elems_per_block; + const int at = head_bias + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + for (int s = 0; s <= pos; ++s) { + float kc[PER_LANE], vc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { kc[j] = kn[j]; vc[j] = vn[j]; } + if (s < pos) { + const int np = (s + 1) / L.tokens_per_block; + if (np != page) { + page = np; + kpage = L.pool_base + + (size_t)(krow[np] & 0x7FFFFFFF) * elems_per_block; + vpage = L.pool_base + + (size_t)(vrow[np] & 0x7FFFFFFF) * elems_per_block; + } + const int at = head_bias + + ((s + 1) % L.tokens_per_block) * BB_HEAD_DIM + lane; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + kn[j] = __bfloat162float(kpage[at + j * 32]); + vn[j] = __bfloat162float(vpage[at + j * 32]); + } + } + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) dot += qv[j] * kc[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[j] = acc[j] * rescale + wt * vc[j]; + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH, unsigned int *tile_masks = nullptr, + bf16 *gath = nullptr, int *gidx = nullptr, int *ecnt = nullptr, + int estride = 0, bf16 *row_stage = nullptr) { + __shared__ float score[BB_EXPERTS]; + __shared__ int s_pick[BB_TOP_K], s_slot[BB_TOP_K]; + constexpr int CH = D_MODEL / 8; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + { + const float4 *gsrc = + reinterpret_cast(normed + (size_t)row * D_MODEL); + float4 *gdst = reinterpret_cast(row_stage); + for (int i = threadIdx.x; i < CH; i += THREADS) gdst[i] = gsrc[i]; + } + if ((int)threadIdx.x < BB_EXPERTS) + moe_w[(size_t)row * BB_EXPERTS + threadIdx.x] = 0.f; + __syncthreads(); + for (int e = warp; e < BB_EXPERTS; e += WARPS) { + const bf16 *wr = gate_w + (size_t)e * D_MODEL; + float acc = 0.f; + for (int i = lane; i < D_MODEL; i += 32) + acc += __bfloat162float(row_stage[i]) * + __bfloat162float(wr[i]); + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) score[e] = 1.f / (1.f + __expf(-acc)); + } + __syncthreads(); + if (threadIdx.x == 0) { + float remaining[BB_EXPERTS]; + int pick[BB_TOP_K]; + for (int e = 0; e < BB_EXPERTS; ++e) + remaining[e] = score[e] + __bfloat162float(gate_bias[e]); + float total = 0.f; + for (int k = 0; k < BB_TOP_K; ++k) { + int best = 0; + for (int e = 1; e < BB_EXPERTS; ++e) + if (remaining[e] > remaining[best]) best = e; + pick[k] = best; + total += score[best]; + remaining[best] = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + unsigned int bits = 0u; + for (int k = 0; k < BB_TOP_K; ++k) { + moe_w[row * BB_EXPERTS + pick[k]] = score[pick[k]] * norm; + bits |= 1u << pick[k]; + s_pick[k] = pick[k]; + } + if (tile_masks) atomicOr(&tile_masks[row / BATCH_PAD], bits); + } + __syncthreads(); + if (gath) { + if ((int)threadIdx.x < BB_TOP_K) { + const int k = threadIdx.x; + const int slot = atomicAdd(&ecnt[s_pick[k]], 1); + s_slot[k] = slot; + gidx[(size_t)s_pick[k] * estride + slot] = row; + } + __syncthreads(); + const float4 *src = reinterpret_cast(row_stage); + for (int k = 0; k < BB_TOP_K; ++k) { + float4 *dst = reinterpret_cast( + gath + ((size_t)s_pick[k] * estride + s_slot[k]) * D_MODEL); + for (int i = threadIdx.x; i < CH; i += THREADS) dst[i] = src[i]; + } + } + __syncthreads(); + } +} +__device__ __forceinline__ unsigned moe_row_mask(const float *routes, int row) { + unsigned m = 0u; + const float *rw = routes + (size_t)row * BB_EXPERTS; + #pragma unroll + for (int e = 0; e < BB_EXPERTS; ++e) if (rw[e] != 0.f) m |= 1u << e; + return m; +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx, + unsigned mask) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + for (unsigned m = mask; m; m &= m - 1u) + value = __fadd_rn(value, parts[base + (size_t)(__ffs(m) - 1) * D_MODEL]); + return __float2bfloat16( + __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL])); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + uint64_t *keys, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int n = threadIdx.x; n < SORT_N; n += THREADS) { + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + keys[n] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + __syncthreads(); + for (int k = 2; k <= SORT_N; k <<= 1) { + if (k >= 64) __syncthreads(); + for (int j = k >> 1; j > 0; j >>= 1) { + for (int i = threadIdx.x; i < SORT_N; i += THREADS) { + const int partner = i ^ j; + if (partner > i) { + const bool descending = ((i & k) == 0); + if ((keys[i] < keys[partner]) == descending) { + const uint64_t t = keys[i]; + keys[i] = keys[partner]; + keys[partner] = t; + } + } + } + if (j >= 32) __syncthreads(); + else __syncwarp(); + } + } + __syncthreads(); + if (threadIdx.x == 0) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float v; + int n; + if (!eos_used && eos_v > av) { + v = eos_v; + n = AUDIO_EOS; + eos_used = true; + } else { + v = av; + n = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = + v + gumbel[((size_t)row * N_POS + pos) * TOP_K + j]; + if (noised > best_noised) { best_noised = noised; best_code = n; } + } + out_codes[row] = best_code; + } + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void prefill_expert_masks(const float *routes, unsigned int *masks, int rows) { + const int lane=threadIdx.x&31, warp=threadIdx.x>>5; + const int nt=(rows+BATCH_PAD-1)/BATCH_PAD; + for (int t=blockIdx.x*WARPS+warp;t(gath + g * D_MODEL + c * 8); + if (local < ecnt[e]) + *dst = *reinterpret_cast( + normed + (size_t)gidx[g] * D_MODEL + c * 8); + else + *dst = make_float4(0.f, 0.f, 0.f, 0.f); + } +} +__device__ __forceinline__ int gemm_eff(int tiles, int reps) { + if (NCONS <= 1) return 1; + const int g = (int)gridDim.x; + const int rounds1 = (tiles * reps + g - 1) / g; + if constexpr (NCONS == 3) { + const int max_eff = reps <= BB_EXPERTS ? 2 : 3; + const int rounds2 = (((tiles + 1) / 2) * reps + g - 1) / g; + if (max_eff == 2) return rounds2 < rounds1 ? 2 : 1; + const int rounds3 = (((tiles + 2) / 3) * reps + g - 1) / g; + if (rounds3 < rounds2 && rounds3 < rounds1) return 3; + return rounds2 < rounds1 ? 2 : 1; + } else { + const int roundsN = (((tiles + NCONS - 1) / NCONS) * reps + g - 1) / g; + return roundsN < rounds1 ? NCONS : 1; + } +} +__device__ __forceinline__ int gemm_groups(int tiles, int reps) { + const int eff = gemm_eff(tiles, reps); + return (tiles + eff - 1) / eff; +} +__device__ __forceinline__ int gemm_consumer_id() { + const int g = warpgroup::groupid(); + if (g == 1) return -1; + const int c = g == 0 ? 0 : g - 1; + return c < NCONS ? c : -1; +} +__device__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][BATCH_PAD]; + __shared__ int moe_gidx_stage[NCONS][BATCH_PAD]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + BATCH_PAD - 1) / BATCH_PAD; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + a_s[s], a, {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + a_s[s], a, {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / BATCH_PAD) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, BATCH_PAD> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], a_s[s]); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * BATCH_PAD; + if (epi == EPI_BIAS) { + warpgroup::store(o_sf, accum); + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + dest[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(o_sf[int2{col, row}] + + __bfloat162float(bias[base + col])); + } + warpgroup::sync(1 + cons); + continue; + } + warpgroup::store(o_s[cons], accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_s[cons][int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__device__ void gemm_narrow( + const weight_gl &w, const act_gl_n &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES * NCONS], act_tile (&a_s)[STAGES], + out_tile (&o_s)[NCONS], + semaphore *filled, semaphore *drained, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = PAD_N, + const float *active_routes = nullptr, int active_expert_tiles = 0, + const unsigned int *token_expert_masks = nullptr, int token_expert_tiles = 0, + const int *moe_ntiles = nullptr, const int *moe_gidx = nullptr, + const int *moe_ecnt = nullptr, int moe_estride = 0, + int phase_steps_before = 0, int phase_steps_total = 0) { + __shared__ float moe_wt_stage[NCONS][PAD_N]; + __shared__ int moe_gidx_stage[NCONS][PAD_N]; + __shared__ int s_cum[BB_EXPERTS + 1]; + __shared__ int s_nt[BB_EXPERTS]; + int eff = NCONS; + if (NCONS > 1 && moe_ntiles) { + int nt_tot = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_tot += (moe_ecnt[e] + PAD_N - 1) / PAD_N; + eff = gemm_eff(out_tiles_per_expert, nt_tot); + } + const int otpe_g = out_tiles_per_expert / eff; + if (moe_ntiles) { + if (threadIdx.x == 0) { + int c = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + const int nt = (moe_ecnt[e] + PAD_N - 1) / PAD_N; + s_nt[e] = nt; + s_cum[e] = c; + c += otpe_g * nt; + } + s_cum[BB_EXPERTS] = c; + } + __syncthreads(); + } + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + if (!moe_ntiles) eff = gemm_eff(n_tiles, n_tok); + const int n_groups = (n_tiles + eff - 1) / eff; + const int iters = k_total / KCHUNK; + const int total_work = moe_ntiles ? s_cum[BB_EXPERTS] : n_groups * n_tok; + int flat_lo, flat_hi; + if (phase_steps_total > 0) { + const long long lo = + (long long)phase_steps_total * blockIdx.x / gridDim.x - phase_steps_before; + const long long hi = + (long long)phase_steps_total * (blockIdx.x + 1) / gridDim.x - phase_steps_before; + const long long qlo = lo <= 0 ? 0 : (lo + iters - 1) / iters; + const long long qhi = hi <= 0 ? 0 : (hi + iters - 1) / iters; + flat_lo = (int)(qlo > total_work ? total_work : qlo); + flat_hi = (int)(qhi > total_work ? total_work : qhi); + } else { + flat_lo = (int)((long long)total_work * blockIdx.x / gridDim.x); + flat_hi = (int)((long long)total_work * (blockIdx.x + 1) / gridDim.x); + } + const int iter_base = iter_io; + int work_token_tiles=flat_hi-flat_lo; + if(token_expert_masks) { + work_token_tiles=0; + for(int f=flat_lo;f= 8) + tma::load_async<2, cache_policy::EVICT_LAST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s * NCONS + c], w, {tw, i}, filled[s]); + } + if (reuse >= 8) + tma::load_async<2, cache_policy::EVICT_FIRST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + else + tma::load_async<2, cache_policy::EVICT_LAST>( + *reinterpret_cast(&a_s[s]), a, + {tok, a_col}, filled[s]); + } + } + } + } + } else if (gemm_consumer_id() >= 0 && gemm_consumer_id() < eff) { + const int cons = gemm_consumer_id(); + const uint32_t arrivals = + (uint32_t)(NCONS / eff + (cons < NCONS % eff)); + int iter = iter_base; + for (int f = flat_lo; f < flat_hi; ++f) { + int work, tok, tile_hi; + if (moe_ntiles) { int e = 0; while (s_cum[e+1] <= f) ++e; + const int rem = f - s_cum[e], nt = s_nt[e], j = rem / nt; + work = e * out_tiles_per_expert + j * eff; + tile_hi = (e + 1) * out_tiles_per_expert; + tok = e * (moe_estride / PAD_N) + rem - j * nt; } + else { const int gidx = f / n_tok; tok = f - gidx * n_tok; + work = gidx * eff; tile_hi = n_tiles; } + const bool mine = work + cons < tile_hi; + const int tile_id = mk2_expert_tile( + mine ? work + cons : tile_hi - 1, active_mask, compact_tiles); + if(token_expert_masks && !(token_expert_masks[tok] & + (1u<<(tile_id/token_expert_tiles)))) continue; + rt_fl<16, PAD_N> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s * NCONS + cons], + *reinterpret_cast(&a_s[s])); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s], arrivals); + } + if (!mine) continue; + const int base = moe_ntiles + ? (tile_id % out_tiles_per_expert) * OUT_ATOM + : tile_id * OUT_ATOM; + const int grow0 = tok * PAD_N; + out_tile_n &o_sn = *reinterpret_cast(&o_s[cons]); + warpgroup::store(o_sn, accum); + if (epi == EPI_MOE_ACC) { + const int expert_w = tile_id / out_tiles_per_expert; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) + moe_wt_stage[cons][r] = + moe_w[(size_t)(grow0 + r) * BB_EXPERTS + expert_w]; + } else if (epi == EPI_MOE_ACC_C) { + const int expert_w = tile_id / out_tiles_per_expert; + const int cnt = moe_ecnt[expert_w]; + for (int r = warpgroup::laneid(); r < rows_valid; r += 128) { + const int local = grow0 + r - expert_w * moe_estride; + const int gr = local < cnt ? moe_gidx[grow0 + r] : -1; + moe_gidx_stage[cons][r] = gr; + moe_wt_stage[cons][r] = gr >= 0 + ? moe_w[(size_t)gr * BB_EXPERTS + expert_w] : 0.f; + } + } + warpgroup::sync(1 + cons); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_sn[int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_wt_stage[cons][row]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + __fmul_rn(wt, v); + break; + } + case EPI_MOE_ACC_C: { + const int grow_o = moe_gidx_stage[cons][row]; + if (grow_o >= 0) + moe_out[((size_t)grow_o * (BB_EXPERTS + 1) + + tile_id / out_tiles_per_expert) * D_MODEL + + base + col] = + __fmul_rn(moe_wt_stage[cons][row], v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1 + cons); + } + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&w_s)[STAGES * NCONS] = + alloc.allocate(); + act_tile (&a_s)[STAGES] = alloc.allocate(); + out_tile (&o_s)[NCONS] = alloc.allocate(); +#if MK2_PREFILL_ONLY && MK2_GEMM_CONSUMERS >= 2 + static_assert(NCONS * sizeof(out_tile) >= sizeof(out_tile_f), + "aliased head staging must fit the output tiles"); + out_tile_f &o_sf = *reinterpret_cast(&o_s[0]); +#else + out_tile_f (&o_sf) = alloc.allocate(); +#endif + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * NORM_WARPS + 2]; +#if MK2_PREFILL_ONLY + bf16 *const route_stage = reinterpret_cast(&o_s[0]); + static_assert(sizeof(o_s) >= D_MODEL * sizeof(bf16), + "router row stage must fit the output tiles"); +#else + __shared__ __align__(16) uint64_t sort_keys[SORT_N]; + bf16 *const route_stage = reinterpret_cast(&sort_keys[0]); + static_assert(sizeof(sort_keys) >= D_MODEL * sizeof(bf16), + "router row stage must fit the sampler workspace"); +#endif + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS * NCONS, 0); + for (int c = 0; c < CONSUMER_WARPS * NCONS; ++c) arrive(drained[s]); + } + } + int gemm_iter = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map); + grid_barrier(g.barrier); + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map); + grid_barrier(g.barrier); + mamba_gated_norm(L, bb->zxbcdt, bb->gated); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w, + BATCH, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, nullptr, 1, BATCH, + bb->moe_w, BB_MOE_INTER / OUT_ATOM); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + 1, BATCH, bb->moe_w, D_MODEL / OUT_ATOM); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_deterministic(bb->moe_out, idx, + (1u << BB_EXPERTS) - 1u); + grid_barrier(g.barrier); + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int original_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + bool identical_cohort = false; + if constexpr (BATCH > 1) { + if (bb->native_pf_cohort_check && bb->native_prefill_reset && + bb->pf_n_seq == BATCH && !bb->pf_fuse_decode && + !bb->pf_layer_out && !bb->moe_w_inject && !bb->moe_w_audit) { + if (blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[0] = 0; + grid_barrier(g.barrier); + const auto *bits = reinterpret_cast(bb->prompt); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < original_total * D_MODEL; idx += gridDim.x * THREADS) { + const int canonical = (idx / (BATCH * D_MODEL)) * BATCH * D_MODEL + idx % D_MODEL; + if (bits[idx] != bits[canonical]) atomicOr(bb->native_pf_cohort_check, 1u); + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + if (bb->pf_seq_slot[row] != row || bb->slot_map[row] != row) + atomicOr(bb->native_pf_cohort_check, 1u); + grid_barrier(g.barrier); + identical_cohort = bb->native_pf_cohort_check[0] == 0; + } + } + const int n_total = identical_cohort ? original_total / BATCH : original_total; + if (bb->native_pf_cohort_check && blockIdx.x == 0 && threadIdx.x == 0) + bb->native_pf_cohort_check[1] = identical_cohort ? 1 : BATCH; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int reset_slots = identical_cohort ? 1 : BATCH; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * reset_slots * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % reset_slots; + const MambaLayer &L = bb->mamba[row_layer / reset_slots]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * reset_slots * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % reset_slots; + const MambaLayer &L = bb->mamba[row_layer / reset_slots]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = original_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + const int src = identical_cohort ? (idx / D_MODEL) * BATCH * D_MODEL + idx % D_MODEL : idx; + bb->pf_h[idx] = bb->prompt[src]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = identical_cohort ? 1 : (bb->pf_n_seq > 0 ? bb->pf_n_seq : 1); + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T, + (!bb->pf_layer_out && layer > 0 && + bb->moe_of_layer[layer - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + if (bb->native_pf_expert_masks) + for (int i = blockIdx.x * THREADS + threadIdx.x; i < n_tok; + i += gridDim.x * THREADS) + bb->native_pf_expert_masks[i] = 0u; + if (bb->pf_moe_ecnt) + for (int e = blockIdx.x * THREADS + threadIdx.x; e < BB_EXPERTS; + e += gridDim.x * THREADS) + bb->pf_moe_ecnt[e] = 0; + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + mamba_dtda_prefill(L, bb->pf_zxbcdt, bb->pf_dtda, T); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_dtda, + bb->pf_gated, T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T, nullptr, nullptr, nullptr, nullptr, 0, + route_stage); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, + T, bb->native_pf_expert_masks, bb->pf_moe_gath, + bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, route_stage); + } + grid_barrier(g.barrier); + int max_ecnt = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + max_ecnt = max(max_ecnt, bb->pf_moe_ecnt[e]); + const bool narrow = BATCH_PAD > PAD_N && max_ecnt <= PAD_N; + const int etile = narrow ? PAD_N : BATCH_PAD; + int nt_sum = 0; + for (int e = 0; e < BB_EXPERTS; ++e) + nt_sum += (bb->pf_moe_ecnt[e] + etile - 1) / etile; + const int up_steps = + gemm_groups(BB_MOE_INTER / OUT_ATOM, nt_sum) * + nt_sum * (D_MODEL / KCHUNK); + const int shup_steps = + gemm_groups(BB_SHARED_INTER / OUT_ATOM, n_tok) * + n_tok * (D_MODEL / KCHUNK); + const int down_steps = + gemm_groups(D_MODEL / OUT_ATOM, nt_sum) * + nt_sum * (BB_MOE_INTER / KCHUNK); + const int shdown_steps = + gemm_groups(D_MODEL / OUT_ATOM, n_tok) * + n_tok * (BB_SHARED_INTER / KCHUNK); + if (narrow) + gemm_narrow(L.exp_up, bb->pf_moe_gath_vn, + BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, PAD_N, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + else + gemm(L.exp_up, bb->pf_moe_gath_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_hc, BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, BB_MOE_INTER / OUT_ATOM, 0, nullptr, nullptr, + n_tok, BATCH_PAD, nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, up_steps + shup_steps); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + up_steps, up_steps + shup_steps); + grid_barrier(g.barrier); + if (narrow) + gemm_narrow(L.exp_down, bb->pf_moe_hc_vn, + BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, PAD_N, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + else + gemm(L.exp_down, bb->pf_moe_hc_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC_C, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, 0, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, + bb->pf_moe_ntiles, bb->pf_moe_gidx, bb->pf_moe_ecnt, + bb->pf_moe_estride, 0, down_steps + shdown_steps); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD, + nullptr, 0, nullptr, 0, nullptr, nullptr, nullptr, 0, + down_steps, down_steps + shdown_steps); + grid_barrier(g.barrier); + if (bb->pf_layer_out) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * D_MODEL; idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = + bb_moe_fold_deterministic(bb->pf_moe_out, idx, + moe_row_mask(bb->pf_moe_w, idx / D_MODEL)); + grid_barrier(g.barrier); + } + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, + identical_cohort ? bb->pf_normed : bb->pf_final, ln_scratch, + n_total, + (!bb->pf_layer_out && n_backbone_layers > 0 && + bb->moe_of_layer[n_backbone_layers - 1] >= 0) + ? bb->pf_moe_out : nullptr, + bb->pf_moe_w); + grid_barrier(g.barrier); + if (identical_cohort) { + const int worker = blockIdx.x * THREADS + threadIdx.x; + const int workers = gridDim.x * THREADS; + for (int i = worker; i < original_total * D_MODEL; i += workers) + bb->pf_final[i] = bb->pf_normed[(i / (BATCH * D_MODEL)) * D_MODEL + i % D_MODEL]; + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = worker; i < 14 * (BATCH - 1) * CV; i += workers) { + const int cell = i % CV, pair = i / CV; + const int row = pair % (BATCH - 1) + 1; + const MambaLayer &L = bb->mamba[pair / (BATCH - 1)]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = L.conv_state[cell]; + } + for (int i = worker; i < 14 * (BATCH - 1) * SS; i += workers) { + const int cell = i % SS, pair = i / SS; + const int row = pair % (BATCH - 1) + 1; + const MambaLayer &L = bb->mamba[pair / (BATCH - 1)]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = L.ssm_state[cell]; + } + const int kv_columns = BB_KV_HEADS * BB_HEAD_DIM; + for (int i = worker; i < 3 * 2 * (BATCH - 1) * n_total * kv_columns; i += workers) { + int q = i / kv_columns; + const int col = i % kv_columns; + const int pos = q % n_total; q /= n_total; + const int row = q % (BATCH - 1) + 1; q /= BATCH - 1; + const int kind = q % 2, layer = q / 2; + const AttnLayer &L = bb->attn[layer]; + const int offset = kv_offset(L, pos, col / BB_HEAD_DIM) + col % BB_HEAD_DIM; + kv_slot(L, row, kind, pos)[offset] = kv_slot(L, 0, kind, pos)[offset]; + } + grid_barrier(g.barrier); + } + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } +#if MK2_PREFILL_ONLY + return; +#else + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, bb->phoneme_b); + grid_barrier(g.barrier); + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + if (threadIdx.x == 0) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = 0; n < BB_PHONEME; ++n) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + bb->phoneme_out[row] = arg; + } + __syncthreads(); + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + for (int i = lane; i < D_HEAD; i += 32) { + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + for (int t = 0; t <= pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + for (int i = lane; i < D_HEAD; i += 32) { + float acc = 0.f; + for (int t = 0; t <= pos; ++t) + acc += sc[t] * __bfloat162float( + g.vcache[cache_base + t * D_HEAD + i]); + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + i] = + __float2bfloat16(acc); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, g.sampled_pos, + sort_keys, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { +#ifdef ABLATE_SAMPLE + __syncthreads(); +#endif + } else { + grid_barrier(g.barrier); + } + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x < BATCH && blockIdx.x == 0) { + const int code = g.sampled_pos[threadIdx.x]; + g.sampled[threadIdx.x * N_POS + pos] = code; + if (g.native_codes) + g.native_codes[threadIdx.x * N_POS + pos] = (long long)code; + } + grid_barrier(g.barrier); + } + if (g.native_seq_len && blockIdx.x == 0 && threadIdx.x < BATCH) + g.native_seq_len[threadIdx.x] += 1; + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +#endif +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 2 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * (NCONS * sizeof(weight_tile) + sizeof(act_tile)) + + NCONS * sizeof(out_tile) + MK2_HEAD_STAGE_BYTES + MK2_SMEM_SLACK), + 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = MK2_PREFILL_ONLY + ? out->sms * out->blocks_per_sm + : std::min(out->sms * out->blocks_per_sm, D_FFN / OUT_ATOM); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + 0, + nullptr, + nullptr, + nullptr, nullptr, nullptr, nullptr, nullptr, 0, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl_n{moe_h, nullptr, nullptr, (size_t)BATCH, (size_t)BB_MOE_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + if constexpr (MK2_PREFILL_ONLY) { + g_last_error = "prefill-only tile library: use a separate Stage0/predictor library"; + return 1; + } + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "dedup_native_prefill_impl.inc" diff --git a/deploy/nemotron_native/kernels/prefill/b8/mk2_abi.cuh b/deploy/nemotron_native/kernels/prefill/b8/mk2_abi.cuh new file mode 100644 index 0000000..95a253b --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b8/mk2_abi.cuh @@ -0,0 +1,82 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} diff --git a/deploy/nemotron_native/kernels/prefill/b8/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/prefill/b8/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b8/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/prefill/b8/prefill.cu b/deploy/nemotron_native/kernels/prefill/b8/prefill.cu new file mode 100644 index 0000000..ea72250 --- /dev/null +++ b/deploy/nemotron_native/kernels/prefill/b8/prefill.cu @@ -0,0 +1,18 @@ +#define MK2_BATCH 8 +#define MK2_BATCH_TILE 64 +#define MK2_PREFILL_ONLY 1 +#define STAGES 3 +#define KCHUNK 128 +#define MK2_GEMM_CONSUMERS 3 +#define MK2_THREADS 512 +#define MK2_NO_MAIN 1 +#define MK2_DIRECT_INPUT 1 +#define MK2_B1_LOCAL_BARRIERS 1 +#define MK2_SKIP_HIDDEN_SNAPSHOT 1 +#define MK2_ACTIVE_EXPERTS 1 +#define MK2_SSM_SPLITS 1 +#define PF_SSM_PW_MULTI 12 +#define PF_ACTIVE_TOKEN_EXPERTS 1 +#include +#include +#include "dedup_predictor_mk2.cuh" diff --git a/deploy/nemotron_native/kernels/recipes.json b/deploy/nemotron_native/kernels/recipes.json new file mode 100644 index 0000000..453d9bb --- /dev/null +++ b/deploy/nemotron_native/kernels/recipes.json @@ -0,0 +1,2936 @@ +{ + "schema": "nemotron-native-kernel-recipes.v2", + "architecture": "sm_90a", + "recipes": { + "b1/stage0": { + "width": 1, + "role": "stage0", + "source": "stage0/b1/predictor_mk2.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-Xcompiler=-fPIC", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v", + "-lcuda", + "-DMK2_BATCH=1", + "-DSTAGES=3", + "-DKCHUNK=256", + "-DMK2_NO_MAIN=1", + "-DMK2_DIRECT_INPUT=1", + "-DMK2_B1_LOCAL_BARRIERS=1", + "-DMK2_SKIP_HIDDEN_SNAPSHOT=1", + "-DMK2_ACTIVE_EXPERTS=1", + "-DMK2_SSM_SPLITS=6", + "-DMK2_PARALLEL_ROUTER=1", + "-DMK2_PARALLEL_SAMPLE=1", + "-DMK2_PARALLEL_PHONEME=1" + ], + "include_dirs": [ + "stage0/b1", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype" + ], + "dependencies": [ + "inflight/progress.cuh", + "inflight/runtime.cuh", + "stage0/b1/bb_router_selection.cuh", + "stage0/b1/mk2_abi.cuh", + "stage0/b1/native_prefill_abi.cuh", + "stage0/b1/native_prefill_impl.inc", + "stage0/b1/predictor_mk2.cu" + ], + "output": "programs/b1/artifacts/kernels/b1/stage0.so" + }, + "b1/prefill": { + "width": 1, + "role": "prefill", + "source": "prefill/b1/prefill.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-Xcompiler=-fPIC", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v", + "-lcuda" + ], + "include_dirs": [ + "prefill/b1", + "third_party/thunderkittens-v4-prefill/include" + ], + "dependencies": [ + "prefill/b1/dedup_native_prefill_impl.inc", + "prefill/b1/dedup_predictor_mk2.cuh", + "prefill/b1/mk2_abi.cuh", + "prefill/b1/native_prefill_abi.cuh", + "prefill/b1/prefill.cu" + ], + "output": "programs/b1/artifacts/kernels/b1/prefill.so", + "qualified_library_sha256": "0d32efb9bf04ddc1ccadeae02394e58728f122192a9a87065ee81ba64580bfe1" + }, + "b1/codec": { + "width": 1, + "role": "codec", + "source": "codec/base/native_stage1.cu", + "flags": [ + "-shared", + "-rdc=false", + "-gencode=arch=compute_90a,code=sm_90a", + "-std=c++17", + "-O3", + "-lineinfo", + "-Xcompiler=-fPIC", + "-Xptxas=-v", + "-lcudadevrt", + "--maxrregcount", + "128", + "-DNATIVE_CODEC_PREFETCH=1", + "-DNATIVE_CODEC_PACK_SUBCTAS=1", + "-DD7W=96" + ], + "include_dirs": [ + "codec/base", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype" + ], + "dependencies": [ + "codec/base/kernel_kf_codec.h", + "codec/base/legacy_codec.cuh", + "codec/base/native_stage1.cu", + "codec/base/native_stage1.h", + "codec/base/repack3.inc", + "codec/base/tail_kernel.inc", + "codec/base/wg2_kernel.inc", + "codec/base/wg3_kernel.inc", + "codec/base/wg4_step_widths.inc", + "codec/base/wg5_step_widths.inc", + "codec/base/wg_kernel.inc", + "codec/base/wgmma_macro.inc" + ], + "output": "programs/b1/artifacts/kernels/b1/codec.so", + "qualified_library_sha256": "2f1e6aebb31a0d5841b62ed4d84d8b423c2c512cfdd884f4a1b6a3be8070f0a5" + }, + "b2/stage0": { + "width": 2, + "role": "stage0", + "source": "stage0/small/predictor_mk2.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-Xcompiler=-fPIC", + "-Xlinker=-Bsymbolic", + "-lcuda", + "-DMK2_BATCH=2", + "-DSTAGES=3", + "-DKCHUNK=256", + "-DMK2_NO_MAIN=1", + "-DMK2_DIRECT_INPUT=1", + "-DMK2_B1_LOCAL_BARRIERS=1", + "-DMK2_SKIP_HIDDEN_SNAPSHOT=1", + "-DMK2_ACTIVE_EXPERTS=1", + "-DMK2_SSM_SPLITS=1", + "-DMK2_SPARSE_MOE=1", + "-Xptxas=-v" + ], + "include_dirs": [ + "stage0/small", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype" + ], + "dependencies": [ + "inflight/progress.cuh", + "inflight/runtime.cuh", + "stage0/small/mk2_abi.cuh", + "stage0/small/native_prefill_abi.cuh", + "stage0/small/native_prefill_impl.inc", + "stage0/small/predictor_mk2.cu", + "stage0/small/sparse_moe_impl.cuh" + ], + "output": "programs/b2/artifacts/kernels/stage0/native.so", + "qualified_library_sha256": "5fb72da3199aa9dd5c6b4a8b4433f5951cb8f9eb501b4e5dfa8dd04a60da7b21" + }, + "b2/prefill": { + "width": 2, + "role": "prefill", + "source": "prefill/b2/prefill.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-Xcompiler=-fPIC", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v", + "-lcuda" + ], + "include_dirs": [ + "prefill/b2", + "third_party/thunderkittens-v4-prefill/include" + ], + "dependencies": [ + "prefill/b2/dedup_native_prefill_impl.inc", + "prefill/b2/dedup_predictor_mk2.cuh", + "prefill/b2/mk2_abi.cuh", + "prefill/b2/native_prefill_abi.cuh", + "prefill/b2/prefill.cu" + ], + "output": "programs/b2/artifacts/kernels/prefill/native.so", + "qualified_library_sha256": "67a89acb9c228d0cca67ca658d216883e35682b48f8f99cb7a5f33be010090e6" + }, + "b2/codec": { + "width": 2, + "role": "codec", + "source": "codec/small/native_stage1.cu", + "flags": [ + "-shared", + "-rdc=false", + "-gencode=arch=compute_90a,code=sm_90a", + "-std=c++17", + "-O3", + "-lineinfo", + "-Xcompiler=-fPIC", + "-Xptxas=-v", + "-lcudadevrt", + "--maxrregcount", + "128", + "-DNATIVE_CODEC_PREFETCH=1", + "-DNATIVE_CODEC_PACK_SUBCTAS=1", + "-DD7W=96", + "-DSMALL_NATIVE_BATCH=2" + ], + "include_dirs": [ + "codec/small", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype" + ], + "dependencies": [ + "codec/small/kernel_kf_codec.h", + "codec/small/legacy_codec.cuh", + "codec/small/native_stage1.cu", + "codec/small/native_stage1.h", + "codec/small/repack3.inc", + "codec/small/tail_kernel.inc", + "codec/small/wg2_kernel.inc", + "codec/small/wg3_kernel.inc", + "codec/small/wg4_step_widths.inc", + "codec/small/wg5_step_widths.inc", + "codec/small/wg_kernel.inc", + "codec/small/wgmma_macro.inc" + ], + "output": "programs/b2/artifacts/kernels/codec/native.so", + "qualified_library_sha256": "ee1c038a088d2cfbb84311da2c804563afa31bca3187a2b09537bb61c328c165" + }, + "b4/stage0": { + "width": 4, + "role": "stage0", + "source": "stage0/small/predictor_mk2.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-Xcompiler=-fPIC", + "-Xlinker=-Bsymbolic", + "-lcuda", + "-DMK2_BATCH=4", + "-DSTAGES=3", + "-DKCHUNK=256", + "-DMK2_NO_MAIN=1", + "-DMK2_DIRECT_INPUT=1", + "-DMK2_B1_LOCAL_BARRIERS=1", + "-DMK2_SKIP_HIDDEN_SNAPSHOT=1", + "-DMK2_ACTIVE_EXPERTS=1", + "-DMK2_SSM_SPLITS=1", + "-DMK2_SPARSE_MOE=1", + "-Xptxas=-v" + ], + "include_dirs": [ + "stage0/small", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype" + ], + "dependencies": [ + "inflight/progress.cuh", + "inflight/runtime.cuh", + "stage0/small/mk2_abi.cuh", + "stage0/small/native_prefill_abi.cuh", + "stage0/small/native_prefill_impl.inc", + "stage0/small/predictor_mk2.cu", + "stage0/small/sparse_moe_impl.cuh" + ], + "output": "programs/b4/artifacts/kernels/stage0/native.so", + "qualified_library_sha256": "e462525a21a47b7156384a3b86a69671734095106a97992af1ff4e73c1b374b6" + }, + "b4/prefill": { + "width": 4, + "role": "prefill", + "source": "prefill/b4/prefill.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-Xcompiler=-fPIC", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v", + "-lcuda" + ], + "include_dirs": [ + "prefill/b4", + "third_party/thunderkittens-v4-prefill/include" + ], + "dependencies": [ + "prefill/b4/dedup_native_prefill_impl.inc", + "prefill/b4/dedup_predictor_mk2.cuh", + "prefill/b4/mk2_abi.cuh", + "prefill/b4/native_prefill_abi.cuh", + "prefill/b4/prefill.cu" + ], + "output": "programs/b4/artifacts/kernels/prefill/native.so", + "qualified_library_sha256": "fa4d889479309f0cd82a9f1d9d59c71229eac4a346b13edeeb90bdfe7298e670" + }, + "b4/codec": { + "width": 4, + "role": "codec", + "source": "codec/small/native_stage1.cu", + "flags": [ + "-shared", + "-rdc=false", + "-gencode=arch=compute_90a,code=sm_90a", + "-std=c++17", + "-O3", + "-lineinfo", + "-Xcompiler=-fPIC", + "-Xptxas=-v", + "-lcudadevrt", + "--maxrregcount", + "128", + "-DNATIVE_CODEC_PREFETCH=1", + "-DNATIVE_CODEC_PACK_SUBCTAS=1", + "-DD7W=96", + "-DSMALL_NATIVE_BATCH=4" + ], + "include_dirs": [ + "codec/small", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype" + ], + "dependencies": [ + "codec/small/kernel_kf_codec.h", + "codec/small/legacy_codec.cuh", + "codec/small/native_stage1.cu", + "codec/small/native_stage1.h", + "codec/small/repack3.inc", + "codec/small/tail_kernel.inc", + "codec/small/wg2_kernel.inc", + "codec/small/wg3_kernel.inc", + "codec/small/wg4_step_widths.inc", + "codec/small/wg5_step_widths.inc", + "codec/small/wg_kernel.inc", + "codec/small/wgmma_macro.inc" + ], + "output": "programs/b4/artifacts/kernels/codec/native.so", + "qualified_library_sha256": "685d1795ba50aedfd762a78fdb1c8e45132eaaf8f5049013342b1402e73b7f90" + }, + "b8/stage0": { + "width": 8, + "role": "stage0", + "source": "stage0/small/predictor_mk2.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-Xcompiler=-fPIC", + "-Xlinker=-Bsymbolic", + "-lcuda", + "-DMK2_BATCH=8", + "-DSTAGES=3", + "-DKCHUNK=256", + "-DMK2_NO_MAIN=1", + "-DMK2_DIRECT_INPUT=1", + "-DMK2_B1_LOCAL_BARRIERS=1", + "-DMK2_SKIP_HIDDEN_SNAPSHOT=1", + "-DMK2_ACTIVE_EXPERTS=1", + "-DMK2_SSM_SPLITS=1", + "-DMK2_SPARSE_MOE=1", + "-Xptxas=-v" + ], + "include_dirs": [ + "stage0/small", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype" + ], + "dependencies": [ + "inflight/progress.cuh", + "inflight/runtime.cuh", + "stage0/small/mk2_abi.cuh", + "stage0/small/native_prefill_abi.cuh", + "stage0/small/native_prefill_impl.inc", + "stage0/small/predictor_mk2.cu", + "stage0/small/sparse_moe_impl.cuh" + ], + "output": "programs/b8/artifacts/kernels/stage0/native.so", + "qualified_library_sha256": "c1f606ff84636a4b418f09a52891d25fe9ec4b558b5b5e05e34207d4b1cfdc06" + }, + "b8/prefill": { + "width": 8, + "role": "prefill", + "source": "prefill/b8/prefill.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-Xcompiler=-fPIC", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v", + "-lcuda" + ], + "include_dirs": [ + "prefill/b8", + "third_party/thunderkittens-v4-prefill/include" + ], + "dependencies": [ + "prefill/b8/dedup_native_prefill_impl.inc", + "prefill/b8/dedup_predictor_mk2.cuh", + "prefill/b8/mk2_abi.cuh", + "prefill/b8/native_prefill_abi.cuh", + "prefill/b8/prefill.cu" + ], + "output": "programs/b8/artifacts/kernels/prefill/native.so", + "qualified_library_sha256": "51e2fdcf82400127761d4208fca81e915604ffa6b92e9a7bd8a895bdc8317b09" + }, + "b8/codec": { + "width": 8, + "role": "codec", + "source": "codec/small/native_stage1.cu", + "flags": [ + "-shared", + "-rdc=false", + "-gencode=arch=compute_90a,code=sm_90a", + "-std=c++17", + "-O3", + "-lineinfo", + "-Xcompiler=-fPIC", + "-Xptxas=-v", + "-lcudadevrt", + "--maxrregcount", + "128", + "-DNATIVE_CODEC_PREFETCH=1", + "-DNATIVE_CODEC_PACK_SUBCTAS=1", + "-DD7W=96", + "-DSMALL_NATIVE_BATCH=8" + ], + "include_dirs": [ + "codec/small", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype" + ], + "dependencies": [ + "codec/small/kernel_kf_codec.h", + "codec/small/legacy_codec.cuh", + "codec/small/native_stage1.cu", + "codec/small/native_stage1.h", + "codec/small/repack3.inc", + "codec/small/tail_kernel.inc", + "codec/small/wg2_kernel.inc", + "codec/small/wg3_kernel.inc", + "codec/small/wg4_step_widths.inc", + "codec/small/wg5_step_widths.inc", + "codec/small/wg_kernel.inc", + "codec/small/wgmma_macro.inc" + ], + "output": "programs/b8/artifacts/kernels/codec/native.so", + "qualified_library_sha256": "993caed1fffc38c9fb257af0f4f6c99203a23a7b7f28744d28ec313ac6a47cb6" + }, + "b16/stage0": { + "width": 16, + "role": "stage0", + "source": "stage0/b16/predictor_mk2.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-Xcompiler=-fPIC", + "-Xlinker=-Bsymbolic", + "-lcuda", + "-DMK2_BATCH=16", + "-DSTAGES=3", + "-DKCHUNK=256", + "-DMK2_NO_MAIN=1", + "-DMK2_DIRECT_INPUT=1", + "-DMK2_B1_LOCAL_BARRIERS=1", + "-DMK2_SKIP_HIDDEN_SNAPSHOT=1", + "-DMK2_ACTIVE_EXPERTS=1", + "-DMK2_SSM_SPLITS=1", + "-DMK2_SPARSE_MOE=1", + "-Xptxas=-v" + ], + "include_dirs": [ + "stage0/b16", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype" + ], + "dependencies": [ + "inflight/progress.cuh", + "inflight/runtime.cuh", + "stage0/b16/mk2_abi.cuh", + "stage0/b16/native_prefill_abi.cuh", + "stage0/b16/native_prefill_impl.inc", + "stage0/b16/predictor_mk2.cu", + "stage0/b16/sparse_moe_impl.cuh" + ], + "output": "programs/b16/artifacts/kernels/b16/stage0/native.so", + "qualified_library_sha256": "72b6bc359d41916e7f7e3041135a79c085d0965ea801f309792cad156fae0939" + }, + "b16/prefill": { + "width": 16, + "role": "prefill", + "source": "prefill/b16/prefill.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-Xcompiler=-fPIC", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v", + "-lcuda" + ], + "include_dirs": [ + "prefill/b16", + "third_party/thunderkittens-v4-prefill/include" + ], + "dependencies": [ + "prefill/b16/dedup_native_prefill_impl.inc", + "prefill/b16/dedup_predictor_mk2.cuh", + "prefill/b16/mk2_abi.cuh", + "prefill/b16/native_prefill_abi.cuh", + "prefill/b16/prefill.cu" + ], + "output": "programs/b16/artifacts/kernels/b16/prefill/native.so", + "qualified_library_sha256": "3bb6471b3bdb20ebaecec393f1f571223ac7e2dc588de842149b2e11d105064a" + }, + "b16/codec": { + "width": 16, + "role": "codec", + "source": "codec/base/native_stage1.cu", + "flags": [ + "-shared", + "-rdc=false", + "-gencode=arch=compute_90a,code=sm_90a", + "-std=c++17", + "-O3", + "-lineinfo", + "-Xcompiler=-fPIC", + "-Xptxas=-v", + "-lcudadevrt", + "--maxrregcount", + "128", + "-DNATIVE_CODEC_PREFETCH=1", + "-DNATIVE_CODEC_PACK_SUBCTAS=1", + "-DD7W=96" + ], + "include_dirs": [ + "codec/base", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype" + ], + "dependencies": [ + "codec/base/kernel_kf_codec.h", + "codec/base/legacy_codec.cuh", + "codec/base/native_stage1.cu", + "codec/base/native_stage1.h", + "codec/base/repack3.inc", + "codec/base/tail_kernel.inc", + "codec/base/wg2_kernel.inc", + "codec/base/wg3_kernel.inc", + "codec/base/wg4_step_widths.inc", + "codec/base/wg5_step_widths.inc", + "codec/base/wg_kernel.inc", + "codec/base/wgmma_macro.inc" + ], + "output": "programs/b16/artifacts/kernels/b16/codec/libnative_stage1.so", + "qualified_library_sha256": "2f1e6aebb31a0d5841b62ed4d84d8b423c2c512cfdd884f4a1b6a3be8070f0a5", + "alias_of": "b1/codec" + }, + "b32/stage0": { + "width": 32, + "role": "stage0", + "source": "b32/stage0.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v" + ], + "include_dirs": [ + "b32", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b32/native" + ], + "dependencies": [ + "b32/native/stage0/mk2_abi.cuh", + "b32/native/stage0/native_prefill_abi.cuh", + "b32/native/stage0/native_prefill_impl.inc", + "b32/native/stage0/predictor_mk2.cu", + "b32/native/stage0/sparse_moe_impl.cuh", + "b32/stage0.cu", + "inflight/progress.cuh", + "inflight/runtime.cuh" + ], + "output": "programs/b32/artifacts/kernels/b32/stage0/native.so", + "qualified_library_sha256": "29c0766225355307795351c7c273ecbf04d71b5a3b7960441d8415bbc30ea49f" + }, + "b32/prefill": { + "width": 32, + "role": "prefill", + "source": "b32/prefill.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v" + ], + "include_dirs": [ + "b32", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b32/native" + ], + "dependencies": [ + "b32/native/prefill/dedup_native_prefill_impl.inc", + "b32/native/prefill/dedup_predictor_mk2.cuh", + "b32/native/prefill/mk2_abi.cuh", + "b32/native/prefill/native_prefill_abi.cuh", + "b32/native/prefill/prefill.cu", + "b32/prefill.cu" + ], + "output": "programs/b32/artifacts/kernels/b32/prefill/native.so", + "qualified_library_sha256": "0be8c225371b0ff51f9c27f86a083243608fe624af6603621460b2daea36fa18" + }, + "b32/codec": { + "width": 32, + "role": "codec", + "source": "b32/codec.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v" + ], + "include_dirs": [ + "b32", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b32/native" + ], + "dependencies": [ + "b32/codec.cu", + "b32/native/codec/kernel_kf_codec.h", + "b32/native/codec/legacy_codec.cuh", + "b32/native/codec/native_stage1.cu", + "b32/native/codec/native_stage1.h", + "b32/native/codec/repack3.inc", + "b32/native/codec/tail_kernel.inc", + "b32/native/codec/wg2_kernel.inc", + "b32/native/codec/wg3_kernel.inc", + "b32/native/codec/wg4_step_widths.inc", + "b32/native/codec/wg5_step_widths.inc", + "b32/native/codec/wg_kernel.inc", + "b32/native/codec/wgmma_macro.inc" + ], + "output": "programs/b32/artifacts/kernels/b32/codec/native.so", + "qualified_library_sha256": "8400b9ecd6da72abf7a771a9a33f7dc34faef13cbeda083c339092143e161749" + }, + "b64/stage0": { + "width": 64, + "role": "stage0", + "source": "stage0/b64/predictor_mk2.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-Xcompiler=-fPIC", + "-Xlinker=-Bsymbolic", + "-lcuda", + "-DMK2_BATCH=64", + "-DSTAGES=2", + "-DKCHUNK=256", + "-DMK2_NO_MAIN=1", + "-DMK2_DIRECT_INPUT=1", + "-DMK2_B1_LOCAL_BARRIERS=1", + "-DMK2_SKIP_HIDDEN_SNAPSHOT=1", + "-DMK2_ACTIVE_EXPERTS=1", + "-DMK2_SSM_SPLITS=1", + "-Xptxas=-v" + ], + "include_dirs": [ + "stage0/b64", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype" + ], + "dependencies": [ + "inflight/progress.cuh", + "inflight/runtime.cuh", + "stage0/b64/mk2_abi.cuh", + "stage0/b64/native_prefill_abi.cuh", + "stage0/b64/native_prefill_impl.inc", + "stage0/b64/predictor_mk2.cu" + ], + "output": "programs/b64/artifacts/kernels/b64/stage0/native.so", + "qualified_library_sha256": "ea803b85b30ab99e3871cd6dbf682f43d0d10f42d01732bbc40088dfbfcbd1f5" + }, + "b64/prefill": { + "width": 64, + "role": "prefill", + "source": "prefill/b64/prefill.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-Xcompiler=-fPIC", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v", + "-lcuda" + ], + "include_dirs": [ + "prefill/b64", + "third_party/thunderkittens-v4-prefill/include" + ], + "dependencies": [ + "prefill/b64/dedup_native_prefill_impl.inc", + "prefill/b64/dedup_predictor_mk2.cuh", + "prefill/b64/mk2_abi.cuh", + "prefill/b64/native_prefill_abi.cuh", + "prefill/b64/prefill.cu" + ], + "output": "programs/b64/artifacts/kernels/b64/prefill/native.so", + "qualified_library_sha256": "2423ff118b41156758fcfd762af4f12cb61dbd5fde5971b41d6a7aa7bb4d72f3" + }, + "b64/codec": { + "width": 64, + "role": "codec", + "source": "codec/base/native_stage1.cu", + "flags": [ + "-shared", + "-rdc=false", + "-gencode=arch=compute_90a,code=sm_90a", + "-std=c++17", + "-O3", + "-lineinfo", + "-Xcompiler=-fPIC", + "-Xptxas=-v", + "-lcudadevrt", + "--maxrregcount", + "128", + "-DNATIVE_CODEC_PREFETCH=1", + "-DNATIVE_CODEC_PACK_SUBCTAS=1", + "-DD7W=96" + ], + "include_dirs": [ + "codec/base", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype" + ], + "dependencies": [ + "codec/base/kernel_kf_codec.h", + "codec/base/legacy_codec.cuh", + "codec/base/native_stage1.cu", + "codec/base/native_stage1.h", + "codec/base/repack3.inc", + "codec/base/tail_kernel.inc", + "codec/base/wg2_kernel.inc", + "codec/base/wg3_kernel.inc", + "codec/base/wg4_step_widths.inc", + "codec/base/wg5_step_widths.inc", + "codec/base/wg_kernel.inc", + "codec/base/wgmma_macro.inc" + ], + "output": "programs/b64/artifacts/kernels/b64/codec/libnative_stage1.so", + "qualified_library_sha256": "2f1e6aebb31a0d5841b62ed4d84d8b423c2c512cfdd884f4a1b6a3be8070f0a5", + "alias_of": "b1/codec" + }, + "b128/stage0": { + "width": 128, + "role": "stage0", + "source": "b128/stage0.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v" + ], + "include_dirs": [ + "b128", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b128/regions/stage0/native", + "b128/regions/prefill/native" + ], + "dependencies": [ + "b128/regions/stage0/native/stage0/mk2_abi.cuh", + "b128/regions/stage0/native/stage0/native_prefill_abi.cuh", + "b128/regions/stage0/native/stage0/native_prefill_impl.inc", + "b128/regions/stage0/native/stage0/predictor_mk2.cu", + "b128/regions/stage0/stage0.cu", + "b128/stage0.cu", + "inflight/progress.cuh", + "inflight/runtime.cuh" + ], + "output": "programs/b128/artifacts/kernels/b128/stage0.so", + "qualified_library_sha256": "4a40b7845c62e36a6a9362a4f8d6d53ca136f33d901a2d8ece217b11791edbb6" + }, + "b128/prefill": { + "width": 128, + "role": "prefill", + "source": "b128/prefill.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v" + ], + "include_dirs": [ + "b128", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b128/regions/stage0/native", + "b128/regions/prefill/native" + ], + "dependencies": [ + "b128/prefill.cu", + "b128/regions/prefill/native/prefill/dedup_native_prefill_impl.inc", + "b128/regions/prefill/native/prefill/dedup_predictor_mk2.cuh", + "b128/regions/prefill/native/prefill/mk2_abi.cuh", + "b128/regions/prefill/native/prefill/native_prefill_abi.cuh", + "b128/regions/prefill/native/prefill/prefill.cu", + "b128/regions/prefill/prefill.cu" + ], + "output": "programs/b128/artifacts/kernels/b128/prefill.so", + "qualified_library_sha256": "e828a6f778cb758df2492194d8fd5ec86545a0fbb9e9cd471f114416bfc2d70d" + }, + "b128/codec": { + "width": 128, + "role": "codec", + "source": "codec/first_packet/codec.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v" + ], + "include_dirs": [ + "codec/first_packet", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "codec/first_packet/regions/stage1/native" + ], + "dependencies": [ + "codec/first_packet/codec.cu", + "codec/first_packet/regions/stage1/codec.cu", + "codec/first_packet/regions/stage1/native/codec/kernel_kf_codec.h", + "codec/first_packet/regions/stage1/native/codec/legacy_codec.cuh", + "codec/first_packet/regions/stage1/native/codec/native_stage1.cu", + "codec/first_packet/regions/stage1/native/codec/native_stage1.h", + "codec/first_packet/regions/stage1/native/codec/repack3.inc", + "codec/first_packet/regions/stage1/native/codec/tail_kernel.inc", + "codec/first_packet/regions/stage1/native/codec/wg2_kernel.inc", + "codec/first_packet/regions/stage1/native/codec/wg3_kernel.inc", + "codec/first_packet/regions/stage1/native/codec/wg4_step_widths.inc", + "codec/first_packet/regions/stage1/native/codec/wg5_step_widths.inc", + "codec/first_packet/regions/stage1/native/codec/wg_kernel.inc", + "codec/first_packet/regions/stage1/native/codec/wgmma_macro.inc" + ], + "output": "programs/b128/artifacts/kernels/b128/codec.so", + "qualified_library_sha256": "6f1e03419e8dec694403e614a4d732975f7e724b910df65236e57f8bab2249fb" + }, + "b256/stage0": { + "width": 256, + "role": "stage0", + "source": "b256/stage0.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v" + ], + "include_dirs": [ + "b256", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b256/regions/stage0/native", + "b256/regions/prefill/native" + ], + "dependencies": [ + "b256/regions/stage0/native/stage0/mk2_abi.cuh", + "b256/regions/stage0/native/stage0/native_prefill_abi.cuh", + "b256/regions/stage0/native/stage0/native_prefill_impl.inc", + "b256/regions/stage0/native/stage0/predictor_mk2.cu", + "b256/regions/stage0/stage0.cu", + "b256/stage0.cu", + "inflight/progress.cuh", + "inflight/runtime.cuh" + ], + "output": "programs/b256/artifacts/kernels/b256/stage0.so", + "qualified_library_sha256": "4dd9dd7d091decad0a10cf09951a5c70497a6c9ce3dc9f53a6b649d9f195f1d3" + }, + "b256/prefill": { + "width": 256, + "role": "prefill", + "source": "b256/prefill.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v" + ], + "include_dirs": [ + "b256", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b256/regions/stage0/native", + "b256/regions/prefill/native" + ], + "dependencies": [ + "b256/prefill.cu", + "b256/regions/prefill/native/prefill/dedup_native_prefill_impl.inc", + "b256/regions/prefill/native/prefill/dedup_predictor_mk2.cuh", + "b256/regions/prefill/native/prefill/mk2_abi.cuh", + "b256/regions/prefill/native/prefill/native_prefill_abi.cuh", + "b256/regions/prefill/native/prefill/prefill.cu", + "b256/regions/prefill/prefill.cu" + ], + "output": "programs/b256/artifacts/kernels/b256/prefill.so", + "qualified_library_sha256": "8705c48bc0e8164de25f04f70cc4cafebff7bab81ddb0505877a8a1e795751b6" + }, + "b256/codec": { + "width": 256, + "role": "codec", + "source": "codec/first_packet/codec.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v" + ], + "include_dirs": [ + "codec/first_packet", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "codec/first_packet/regions/stage1/native" + ], + "dependencies": [ + "codec/first_packet/codec.cu", + "codec/first_packet/regions/stage1/codec.cu", + "codec/first_packet/regions/stage1/native/codec/kernel_kf_codec.h", + "codec/first_packet/regions/stage1/native/codec/legacy_codec.cuh", + "codec/first_packet/regions/stage1/native/codec/native_stage1.cu", + "codec/first_packet/regions/stage1/native/codec/native_stage1.h", + "codec/first_packet/regions/stage1/native/codec/repack3.inc", + "codec/first_packet/regions/stage1/native/codec/tail_kernel.inc", + "codec/first_packet/regions/stage1/native/codec/wg2_kernel.inc", + "codec/first_packet/regions/stage1/native/codec/wg3_kernel.inc", + "codec/first_packet/regions/stage1/native/codec/wg4_step_widths.inc", + "codec/first_packet/regions/stage1/native/codec/wg5_step_widths.inc", + "codec/first_packet/regions/stage1/native/codec/wg_kernel.inc", + "codec/first_packet/regions/stage1/native/codec/wgmma_macro.inc" + ], + "output": "programs/b256/artifacts/kernels/b256/codec.so", + "qualified_library_sha256": "6f1e03419e8dec694403e614a4d732975f7e724b910df65236e57f8bab2249fb", + "alias_of": "b128/codec" + }, + "b512/stage0": { + "width": 512, + "role": "stage0", + "source": "b512/stage0.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v" + ], + "include_dirs": [ + "b512", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b512/native" + ], + "dependencies": [ + "b512/native/stage0/mk2_abi.cuh", + "b512/native/stage0/native_prefill_abi.cuh", + "b512/native/stage0/native_prefill_impl.inc", + "b512/native/stage0/predictor_mk2.cu", + "b512/stage0.cu", + "inflight/progress.cuh", + "inflight/runtime.cuh" + ], + "output": "programs/b512/artifacts/kernels/b512/stage0.so", + "qualified_library_sha256": "8b47823086cca473da0bd6abd5cb34632cd04254bfd6dfe019fa95acaac5943b" + }, + "b512/prefill": { + "width": 512, + "role": "prefill", + "source": "b512/prefill.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v" + ], + "include_dirs": [ + "b512", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b512/native" + ], + "dependencies": [ + "b512/native/prefill/dedup_native_prefill_impl.inc", + "b512/native/prefill/dedup_predictor_mk2.cuh", + "b512/native/prefill/mk2_abi.cuh", + "b512/native/prefill/native_prefill_abi.cuh", + "b512/native/prefill/prefill.cu", + "b512/prefill.cu" + ], + "output": "programs/b512/artifacts/kernels/b512/prefill.so", + "qualified_library_sha256": "16ce40283cf11fd7f62660043892bef4b6bd2421411ccea054da38e588b3d429" + }, + "b512/codec": { + "width": 512, + "role": "codec", + "source": "b512/codec.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v" + ], + "include_dirs": [ + "b512", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b512/native" + ], + "dependencies": [ + "b512/codec.cu", + "b512/native/codec/kernel_kf_codec.h", + "b512/native/codec/legacy_codec.cuh", + "b512/native/codec/native_stage1.cu", + "b512/native/codec/native_stage1.h", + "b512/native/codec/repack3.inc", + "b512/native/codec/tail_kernel.inc", + "b512/native/codec/wg2_kernel.inc", + "b512/native/codec/wg3_kernel.inc", + "b512/native/codec/wg4_step_widths.inc", + "b512/native/codec/wg5_step_widths.inc", + "b512/native/codec/wg_kernel.inc", + "b512/native/codec/wgmma_macro.inc" + ], + "output": "programs/b512/artifacts/kernels/b512/codec.so", + "qualified_library_sha256": "59632662d0468367415c84581bcab0145ef4118b2502a817c1fd3802e0034ba4" + }, + "b1024/stage0": { + "width": 1024, + "role": "stage0", + "source": "b1024/stage0.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v" + ], + "include_dirs": [ + "b1024", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b1024/regions/stage0/native", + "b1024/regions/prefill/native" + ], + "dependencies": [ + "b1024/regions/stage0/native/stage0/mk2_abi.cuh", + "b1024/regions/stage0/native/stage0/native_prefill_abi.cuh", + "b1024/regions/stage0/native/stage0/native_prefill_impl.inc", + "b1024/regions/stage0/native/stage0/predictor_mk2.cu", + "b1024/regions/stage0/stage0.cu", + "b1024/stage0.cu", + "inflight/progress.cuh", + "inflight/runtime.cuh" + ], + "output": "programs/b1024/artifacts/kernels/b1024/stage0.so", + "qualified_library_sha256": "efb0a37d55c1a290c9aa794bd989f7aaa56b4bf31c231d18b54b9a0376346dbe" + }, + "b1024/prefill": { + "width": 1024, + "role": "prefill", + "source": "b1024/prefill.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v" + ], + "include_dirs": [ + "b1024", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b1024/regions/stage0/native", + "b1024/regions/prefill/native" + ], + "dependencies": [ + "b1024/prefill.cu", + "b1024/regions/prefill/native/prefill/dedup_native_prefill_impl.inc", + "b1024/regions/prefill/native/prefill/dedup_predictor_mk2.cuh", + "b1024/regions/prefill/native/prefill/mk2_abi.cuh", + "b1024/regions/prefill/native/prefill/native_prefill_abi.cuh", + "b1024/regions/prefill/native/prefill/prefill.cu", + "b1024/regions/prefill/prefill.cu" + ], + "output": "programs/b1024/artifacts/kernels/b1024/prefill.so", + "qualified_library_sha256": "28a987f2447254e18aa1b26f9b6460e2d8cf6c5b08a8bd2ed2fa1b998ece6ce9" + }, + "b1024/codec": { + "width": 1024, + "role": "codec", + "source": "codec/first_packet/codec.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v" + ], + "include_dirs": [ + "codec/first_packet", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "codec/first_packet/regions/stage1/native" + ], + "dependencies": [ + "codec/first_packet/codec.cu", + "codec/first_packet/regions/stage1/codec.cu", + "codec/first_packet/regions/stage1/native/codec/kernel_kf_codec.h", + "codec/first_packet/regions/stage1/native/codec/legacy_codec.cuh", + "codec/first_packet/regions/stage1/native/codec/native_stage1.cu", + "codec/first_packet/regions/stage1/native/codec/native_stage1.h", + "codec/first_packet/regions/stage1/native/codec/repack3.inc", + "codec/first_packet/regions/stage1/native/codec/tail_kernel.inc", + "codec/first_packet/regions/stage1/native/codec/wg2_kernel.inc", + "codec/first_packet/regions/stage1/native/codec/wg3_kernel.inc", + "codec/first_packet/regions/stage1/native/codec/wg4_step_widths.inc", + "codec/first_packet/regions/stage1/native/codec/wg5_step_widths.inc", + "codec/first_packet/regions/stage1/native/codec/wg_kernel.inc", + "codec/first_packet/regions/stage1/native/codec/wgmma_macro.inc" + ], + "output": "programs/b1024/artifacts/kernels/b1024/codec.so", + "qualified_library_sha256": "6f1e03419e8dec694403e614a4d732975f7e724b910df65236e57f8bab2249fb", + "alias_of": "b128/codec" + }, + "b2048/stage0": { + "width": 2048, + "role": "stage0", + "source": "b2048/stage0.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v" + ], + "include_dirs": [ + "b2048", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b2048/regions/stage0/native", + "b2048/regions/prefill/native" + ], + "dependencies": [ + "b2048/regions/stage0/native/stage0/mk2_abi.cuh", + "b2048/regions/stage0/native/stage0/native_prefill_abi.cuh", + "b2048/regions/stage0/native/stage0/native_prefill_impl.inc", + "b2048/regions/stage0/native/stage0/predictor_mk2.cu", + "b2048/regions/stage0/stage0.cu", + "b2048/stage0.cu", + "inflight/progress.cuh", + "inflight/runtime.cuh" + ], + "output": "programs/b2048/artifacts/kernels/b2048/stage0.so", + "qualified_library_sha256": "7a7041c9d8cc0f45ca8170eac59986e6acbb7e7299a8e223016362a152277213" + }, + "b2048/prefill": { + "width": 2048, + "role": "prefill", + "source": "b2048/prefill.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v" + ], + "include_dirs": [ + "b2048", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b2048/regions/stage0/native", + "b2048/regions/prefill/native" + ], + "dependencies": [ + "b2048/prefill.cu", + "b2048/regions/prefill/native/prefill/dedup_native_prefill_impl.inc", + "b2048/regions/prefill/native/prefill/dedup_predictor_mk2.cuh", + "b2048/regions/prefill/native/prefill/mk2_abi.cuh", + "b2048/regions/prefill/native/prefill/native_prefill_abi.cuh", + "b2048/regions/prefill/native/prefill/prefill.cu", + "b2048/regions/prefill/prefill.cu" + ], + "output": "programs/b2048/artifacts/kernels/b2048/prefill.so", + "qualified_library_sha256": "ec8fa24c56be358ea9a30df0033d98363f3a15b8afdd5aee3b9a83b2375e14fc" + }, + "b2048/codec": { + "width": 2048, + "role": "codec", + "source": "codec/first_packet/codec.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v" + ], + "include_dirs": [ + "codec/first_packet", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "codec/first_packet/regions/stage1/native" + ], + "dependencies": [ + "codec/first_packet/codec.cu", + "codec/first_packet/regions/stage1/codec.cu", + "codec/first_packet/regions/stage1/native/codec/kernel_kf_codec.h", + "codec/first_packet/regions/stage1/native/codec/legacy_codec.cuh", + "codec/first_packet/regions/stage1/native/codec/native_stage1.cu", + "codec/first_packet/regions/stage1/native/codec/native_stage1.h", + "codec/first_packet/regions/stage1/native/codec/repack3.inc", + "codec/first_packet/regions/stage1/native/codec/tail_kernel.inc", + "codec/first_packet/regions/stage1/native/codec/wg2_kernel.inc", + "codec/first_packet/regions/stage1/native/codec/wg3_kernel.inc", + "codec/first_packet/regions/stage1/native/codec/wg4_step_widths.inc", + "codec/first_packet/regions/stage1/native/codec/wg5_step_widths.inc", + "codec/first_packet/regions/stage1/native/codec/wg_kernel.inc", + "codec/first_packet/regions/stage1/native/codec/wgmma_macro.inc" + ], + "output": "programs/b2048/artifacts/kernels/b2048/codec.so", + "qualified_library_sha256": "6f1e03419e8dec694403e614a4d732975f7e724b910df65236e57f8bab2249fb", + "alias_of": "b128/codec" + }, + "b4096/stage0": { + "width": 4096, + "role": "stage0", + "source": "b4096/stage0.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v" + ], + "include_dirs": [ + "b4096", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b4096/regions/stage0/native", + "b4096/regions/prefill/native", + "b4096/regions/stage1/native" + ], + "dependencies": [ + "b4096/regions/stage0/native/stage0/mk2_abi.cuh", + "b4096/regions/stage0/native/stage0/native_prefill_abi.cuh", + "b4096/regions/stage0/native/stage0/native_prefill_impl.inc", + "b4096/regions/stage0/native/stage0/predictor_mk2.cu", + "b4096/regions/stage0/stage0.cu", + "b4096/stage0.cu", + "inflight/progress.cuh", + "inflight/runtime.cuh" + ], + "output": "programs/b4096/artifacts/kernels/b4096/stage0.so", + "qualified_library_sha256": "2e3bd3269a2cff8fa46f8944548154bc01a5aa7f34387490a954b0da738aa503" + }, + "b4096/prefill": { + "width": 4096, + "role": "prefill", + "source": "b4096/prefill.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v" + ], + "include_dirs": [ + "b4096", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b4096/regions/stage0/native", + "b4096/regions/prefill/native", + "b4096/regions/stage1/native" + ], + "dependencies": [ + "b4096/prefill.cu", + "b4096/regions/prefill/native/prefill/dedup_native_prefill_impl.inc", + "b4096/regions/prefill/native/prefill/dedup_predictor_mk2.cuh", + "b4096/regions/prefill/native/prefill/mk2_abi.cuh", + "b4096/regions/prefill/native/prefill/native_prefill_abi.cuh", + "b4096/regions/prefill/native/prefill/prefill.cu", + "b4096/regions/prefill/prefill.cu" + ], + "output": "programs/b4096/artifacts/kernels/b4096/prefill.so", + "qualified_library_sha256": "20e8c1f6514923ca4ec8cca1965c78c2ce20f5c0936547e5477763215f9aeb22" + }, + "b4096/codec": { + "width": 4096, + "role": "codec", + "source": "b4096/codec.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v" + ], + "include_dirs": [ + "b4096", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b4096/regions/stage0/native", + "b4096/regions/prefill/native", + "b4096/regions/stage1/native" + ], + "dependencies": [ + "b4096/codec.cu", + "b4096/regions/stage1/codec.cu", + "b4096/regions/stage1/native/codec/kernel_kf_codec.h", + "b4096/regions/stage1/native/codec/legacy_codec.cuh", + "b4096/regions/stage1/native/codec/native_stage1.cu", + "b4096/regions/stage1/native/codec/native_stage1.h", + "b4096/regions/stage1/native/codec/repack3.inc", + "b4096/regions/stage1/native/codec/tail_kernel.inc", + "b4096/regions/stage1/native/codec/wg2_kernel.inc", + "b4096/regions/stage1/native/codec/wg3_kernel.inc", + "b4096/regions/stage1/native/codec/wg4_step_widths.inc", + "b4096/regions/stage1/native/codec/wg5_step_widths.inc", + "b4096/regions/stage1/native/codec/wg_kernel.inc", + "b4096/regions/stage1/native/codec/wgmma_macro.inc" + ], + "output": "programs/b4096/artifacts/kernels/b4096/codec.so", + "qualified_library_sha256": "63d7d6efa78f4e3068d959587ba01279474ac167a74a79f1a4c94114ce2cc8b8" + }, + "b32/contraction_stage0": { + "width": 32, + "role": "contraction_stage0", + "source": "stage0/b16/predictor_mk2.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-Xcompiler=-fPIC", + "-Xlinker=-Bsymbolic", + "-lcuda", + "-DMK2_BATCH=16", + "-DSTAGES=3", + "-DKCHUNK=256", + "-DMK2_NO_MAIN=1", + "-DMK2_DIRECT_INPUT=1", + "-DMK2_B1_LOCAL_BARRIERS=1", + "-DMK2_SKIP_HIDDEN_SNAPSHOT=1", + "-DMK2_ACTIVE_EXPERTS=1", + "-DMK2_SSM_SPLITS=1", + "-DMK2_SPARSE_MOE=1", + "-Xptxas=-v" + ], + "include_dirs": [ + "stage0/b16", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype" + ], + "dependencies": [ + "stage0/b16/mk2_abi.cuh", + "stage0/b16/native_prefill_abi.cuh", + "stage0/b16/native_prefill_impl.inc", + "stage0/b16/predictor_mk2.cu", + "stage0/b16/sparse_moe_impl.cuh" + ], + "output": "programs/b32/artifacts/contraction/b32/stage0.so", + "qualified_library_sha256": "72b6bc359d41916e7f7e3041135a79c085d0965ea801f309792cad156fae0939", + "alias_of": "b16/stage0" + }, + "b64/contraction_stage0": { + "width": 64, + "role": "contraction_stage0", + "source": "stage0/b16/predictor_mk2.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-Xcompiler=-fPIC", + "-Xlinker=-Bsymbolic", + "-lcuda", + "-DMK2_BATCH=16", + "-DSTAGES=3", + "-DKCHUNK=256", + "-DMK2_NO_MAIN=1", + "-DMK2_DIRECT_INPUT=1", + "-DMK2_B1_LOCAL_BARRIERS=1", + "-DMK2_SKIP_HIDDEN_SNAPSHOT=1", + "-DMK2_ACTIVE_EXPERTS=1", + "-DMK2_SSM_SPLITS=1", + "-DMK2_SPARSE_MOE=1", + "-Xptxas=-v" + ], + "include_dirs": [ + "stage0/b16", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype" + ], + "dependencies": [ + "stage0/b16/mk2_abi.cuh", + "stage0/b16/native_prefill_abi.cuh", + "stage0/b16/native_prefill_impl.inc", + "stage0/b16/predictor_mk2.cu", + "stage0/b16/sparse_moe_impl.cuh" + ], + "output": "programs/b64/artifacts/contraction/b64/stage0.so", + "qualified_library_sha256": "72b6bc359d41916e7f7e3041135a79c085d0965ea801f309792cad156fae0939", + "alias_of": "b16/stage0" + }, + "b32/contraction_transfer": { + "width": 32, + "role": "contraction_transfer", + "source": "contraction/lane_transfer.cu", + "flags": [ + "-shared", + "-O3", + "-std=c++17", + "-Xcompiler=-fPIC", + "-gencode", + "arch=compute_90a,code=sm_90a", + "-Xptxas=-v" + ], + "include_dirs": [ + "contraction" + ], + "dependencies": [ + "contraction/lane_transfer.cu", + "contraction/mk2_abi.cuh" + ], + "output": "programs/b32/artifacts/contraction/b32/transfer.so", + "qualified_library_sha256": "58f49b8b770ff5d1ad10280a43a41eca4b5beb671d5c428c29d639c08370aa35" + }, + "b64/contraction_transfer": { + "width": 64, + "role": "contraction_transfer", + "source": "contraction/lane_transfer.cu", + "flags": [ + "-shared", + "-O3", + "-std=c++17", + "-Xcompiler=-fPIC", + "-gencode", + "arch=compute_90a,code=sm_90a", + "-Xptxas=-v" + ], + "include_dirs": [ + "contraction" + ], + "dependencies": [ + "contraction/lane_transfer.cu", + "contraction/mk2_abi.cuh" + ], + "output": "programs/b64/artifacts/contraction/b64/transfer.so", + "qualified_library_sha256": "58f49b8b770ff5d1ad10280a43a41eca4b5beb671d5c428c29d639c08370aa35", + "alias_of": "b32/contraction_transfer" + }, + "b32/contraction_bf16": { + "width": 32, + "role": "contraction_bf16", + "source": "contraction/bf16/predictor_mk2.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-Xcompiler=-fPIC", + "-Xlinker=-Bsymbolic", + "-lcuda", + "-DMK2_BATCH=16", + "-DSTAGES=3", + "-DKCHUNK=256", + "-DMK2_NO_MAIN=1", + "-DMK2_DIRECT_INPUT=1", + "-DMK2_B1_LOCAL_BARRIERS=1", + "-DMK2_SKIP_HIDDEN_SNAPSHOT=1", + "-DMK2_ACTIVE_EXPERTS=1", + "-DMK2_SSM_SPLITS=1", + "-DMK2_SPARSE_MOE=1", + "-Xptxas=-v" + ], + "include_dirs": [ + "contraction/bf16", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype" + ], + "dependencies": [ + "contraction/bf16/mk2_abi.cuh", + "contraction/bf16/native_prefill_abi.cuh", + "contraction/bf16/native_prefill_impl.inc", + "contraction/bf16/predictor_mk2.cu", + "contraction/bf16/sparse_moe_impl.cuh" + ], + "output": "programs/b32/artifacts/contraction/b32/bf16-native.so", + "qualified_library_sha256": "ad02193d44baa580d5cd5affbc6004af1aa81dd4db8319276a059ad15a87e88d" + }, + "b32/stage0_inflight": { + "width": 32, + "role": "stage0_inflight", + "source": "b32/stage0.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v", + "-DMK2_INFLIGHT=1" + ], + "include_dirs": [ + "b32", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b32/native" + ], + "dependencies": [ + "b32/native/stage0/mk2_abi.cuh", + "b32/native/stage0/native_prefill_abi.cuh", + "b32/native/stage0/native_prefill_impl.inc", + "b32/native/stage0/predictor_mk2.cu", + "b32/native/stage0/sparse_moe_impl.cuh", + "b32/stage0.cu", + "inflight/progress.cuh", + "inflight/runtime.cuh" + ], + "output": "programs/b32/artifacts/kernels/b32/stage0/inflight.so" + }, + "b1/stage0_inflight": { + "width": 1, + "role": "stage0_inflight", + "source": "stage0/b1/predictor_mk2.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-Xcompiler=-fPIC", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v", + "-lcuda", + "-DMK2_BATCH=1", + "-DSTAGES=3", + "-DKCHUNK=256", + "-DMK2_NO_MAIN=1", + "-DMK2_DIRECT_INPUT=1", + "-DMK2_B1_LOCAL_BARRIERS=1", + "-DMK2_SKIP_HIDDEN_SNAPSHOT=1", + "-DMK2_ACTIVE_EXPERTS=1", + "-DMK2_SSM_SPLITS=6", + "-DMK2_PARALLEL_ROUTER=1", + "-DMK2_PARALLEL_SAMPLE=1", + "-DMK2_PARALLEL_PHONEME=1", + "-DMK2_INFLIGHT=1" + ], + "include_dirs": [ + "stage0/b1", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype" + ], + "dependencies": [ + "inflight/progress.cuh", + "inflight/runtime.cuh", + "stage0/b1/bb_router_selection.cuh", + "stage0/b1/mk2_abi.cuh", + "stage0/b1/native_prefill_abi.cuh", + "stage0/b1/native_prefill_impl.inc", + "stage0/b1/predictor_mk2.cu" + ], + "output": "programs/b1/artifacts/kernels/b1/inflight.so" + }, + "b2/stage0_inflight": { + "width": 2, + "role": "stage0_inflight", + "source": "stage0/small/predictor_mk2.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-Xcompiler=-fPIC", + "-Xlinker=-Bsymbolic", + "-lcuda", + "-DMK2_BATCH=2", + "-DSTAGES=3", + "-DKCHUNK=256", + "-DMK2_NO_MAIN=1", + "-DMK2_DIRECT_INPUT=1", + "-DMK2_B1_LOCAL_BARRIERS=1", + "-DMK2_SKIP_HIDDEN_SNAPSHOT=1", + "-DMK2_ACTIVE_EXPERTS=1", + "-DMK2_SSM_SPLITS=1", + "-DMK2_SPARSE_MOE=1", + "-Xptxas=-v", + "-DMK2_INFLIGHT=1" + ], + "include_dirs": [ + "stage0/small", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype" + ], + "dependencies": [ + "inflight/progress.cuh", + "inflight/runtime.cuh", + "stage0/small/mk2_abi.cuh", + "stage0/small/native_prefill_abi.cuh", + "stage0/small/native_prefill_impl.inc", + "stage0/small/predictor_mk2.cu", + "stage0/small/sparse_moe_impl.cuh" + ], + "output": "programs/b2/artifacts/kernels/stage0/inflight.so" + }, + "b4/stage0_inflight": { + "width": 4, + "role": "stage0_inflight", + "source": "stage0/small/predictor_mk2.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-Xcompiler=-fPIC", + "-Xlinker=-Bsymbolic", + "-lcuda", + "-DMK2_BATCH=4", + "-DSTAGES=3", + "-DKCHUNK=256", + "-DMK2_NO_MAIN=1", + "-DMK2_DIRECT_INPUT=1", + "-DMK2_B1_LOCAL_BARRIERS=1", + "-DMK2_SKIP_HIDDEN_SNAPSHOT=1", + "-DMK2_ACTIVE_EXPERTS=1", + "-DMK2_SSM_SPLITS=1", + "-DMK2_SPARSE_MOE=1", + "-Xptxas=-v", + "-DMK2_INFLIGHT=1" + ], + "include_dirs": [ + "stage0/small", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype" + ], + "dependencies": [ + "inflight/progress.cuh", + "inflight/runtime.cuh", + "stage0/small/mk2_abi.cuh", + "stage0/small/native_prefill_abi.cuh", + "stage0/small/native_prefill_impl.inc", + "stage0/small/predictor_mk2.cu", + "stage0/small/sparse_moe_impl.cuh" + ], + "output": "programs/b4/artifacts/kernels/stage0/inflight.so" + }, + "b8/stage0_inflight": { + "width": 8, + "role": "stage0_inflight", + "source": "stage0/small/predictor_mk2.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-Xcompiler=-fPIC", + "-Xlinker=-Bsymbolic", + "-lcuda", + "-DMK2_BATCH=8", + "-DSTAGES=3", + "-DKCHUNK=256", + "-DMK2_NO_MAIN=1", + "-DMK2_DIRECT_INPUT=1", + "-DMK2_B1_LOCAL_BARRIERS=1", + "-DMK2_SKIP_HIDDEN_SNAPSHOT=1", + "-DMK2_ACTIVE_EXPERTS=1", + "-DMK2_SSM_SPLITS=1", + "-DMK2_SPARSE_MOE=1", + "-Xptxas=-v", + "-DMK2_INFLIGHT=1" + ], + "include_dirs": [ + "stage0/small", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype" + ], + "dependencies": [ + "inflight/progress.cuh", + "inflight/runtime.cuh", + "stage0/small/mk2_abi.cuh", + "stage0/small/native_prefill_abi.cuh", + "stage0/small/native_prefill_impl.inc", + "stage0/small/predictor_mk2.cu", + "stage0/small/sparse_moe_impl.cuh" + ], + "output": "programs/b8/artifacts/kernels/stage0/inflight.so" + }, + "b16/stage0_inflight": { + "width": 16, + "role": "stage0_inflight", + "source": "stage0/b16/predictor_mk2.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-Xcompiler=-fPIC", + "-Xlinker=-Bsymbolic", + "-lcuda", + "-DMK2_BATCH=16", + "-DSTAGES=3", + "-DKCHUNK=256", + "-DMK2_NO_MAIN=1", + "-DMK2_DIRECT_INPUT=1", + "-DMK2_B1_LOCAL_BARRIERS=1", + "-DMK2_SKIP_HIDDEN_SNAPSHOT=1", + "-DMK2_ACTIVE_EXPERTS=1", + "-DMK2_SSM_SPLITS=1", + "-DMK2_SPARSE_MOE=1", + "-Xptxas=-v", + "-DMK2_INFLIGHT=1" + ], + "include_dirs": [ + "stage0/b16", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype" + ], + "dependencies": [ + "inflight/progress.cuh", + "inflight/runtime.cuh", + "stage0/b16/mk2_abi.cuh", + "stage0/b16/native_prefill_abi.cuh", + "stage0/b16/native_prefill_impl.inc", + "stage0/b16/predictor_mk2.cu", + "stage0/b16/sparse_moe_impl.cuh" + ], + "output": "programs/b16/artifacts/kernels/b16/stage0/inflight.so" + }, + "b64/stage0_inflight": { + "width": 64, + "role": "stage0_inflight", + "source": "stage0/b64/predictor_mk2.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-Xcompiler=-fPIC", + "-Xlinker=-Bsymbolic", + "-lcuda", + "-DMK2_BATCH=64", + "-DSTAGES=2", + "-DKCHUNK=256", + "-DMK2_NO_MAIN=1", + "-DMK2_DIRECT_INPUT=1", + "-DMK2_B1_LOCAL_BARRIERS=1", + "-DMK2_SKIP_HIDDEN_SNAPSHOT=1", + "-DMK2_ACTIVE_EXPERTS=1", + "-DMK2_SSM_SPLITS=1", + "-Xptxas=-v", + "-DMK2_INFLIGHT=1" + ], + "include_dirs": [ + "stage0/b64", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype" + ], + "dependencies": [ + "inflight/progress.cuh", + "inflight/runtime.cuh", + "stage0/b64/mk2_abi.cuh", + "stage0/b64/native_prefill_abi.cuh", + "stage0/b64/native_prefill_impl.inc", + "stage0/b64/predictor_mk2.cu" + ], + "output": "programs/b64/artifacts/kernels/b64/stage0/inflight.so" + }, + "b128/stage0_inflight": { + "width": 128, + "role": "stage0_inflight", + "source": "b128/stage0.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v", + "-DMK2_INFLIGHT=1" + ], + "include_dirs": [ + "b128", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b128/regions/stage0/native", + "b128/regions/prefill/native" + ], + "dependencies": [ + "b128/regions/stage0/native/stage0/mk2_abi.cuh", + "b128/regions/stage0/native/stage0/native_prefill_abi.cuh", + "b128/regions/stage0/native/stage0/native_prefill_impl.inc", + "b128/regions/stage0/native/stage0/predictor_mk2.cu", + "b128/regions/stage0/stage0.cu", + "b128/stage0.cu", + "inflight/progress.cuh", + "inflight/runtime.cuh" + ], + "output": "programs/b128/artifacts/kernels/b128/inflight.so" + }, + "b256/stage0_inflight": { + "width": 256, + "role": "stage0_inflight", + "source": "b256/stage0.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v", + "-DMK2_INFLIGHT=1" + ], + "include_dirs": [ + "b256", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b256/regions/stage0/native", + "b256/regions/prefill/native" + ], + "dependencies": [ + "b256/regions/stage0/native/stage0/mk2_abi.cuh", + "b256/regions/stage0/native/stage0/native_prefill_abi.cuh", + "b256/regions/stage0/native/stage0/native_prefill_impl.inc", + "b256/regions/stage0/native/stage0/predictor_mk2.cu", + "b256/regions/stage0/stage0.cu", + "b256/stage0.cu", + "inflight/progress.cuh", + "inflight/runtime.cuh" + ], + "output": "programs/b256/artifacts/kernels/b256/inflight.so" + }, + "b512/stage0_inflight": { + "width": 512, + "role": "stage0_inflight", + "source": "b512/stage0.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v", + "-DMK2_INFLIGHT=1" + ], + "include_dirs": [ + "b512", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b512/native" + ], + "dependencies": [ + "b512/native/stage0/mk2_abi.cuh", + "b512/native/stage0/native_prefill_abi.cuh", + "b512/native/stage0/native_prefill_impl.inc", + "b512/native/stage0/predictor_mk2.cu", + "b512/stage0.cu", + "inflight/progress.cuh", + "inflight/runtime.cuh" + ], + "output": "programs/b512/artifacts/kernels/b512/inflight.so" + }, + "b1024/stage0_inflight": { + "width": 1024, + "role": "stage0_inflight", + "source": "b1024/stage0.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v", + "-DMK2_INFLIGHT=1" + ], + "include_dirs": [ + "b1024", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b1024/regions/stage0/native", + "b1024/regions/prefill/native" + ], + "dependencies": [ + "b1024/regions/stage0/native/stage0/mk2_abi.cuh", + "b1024/regions/stage0/native/stage0/native_prefill_abi.cuh", + "b1024/regions/stage0/native/stage0/native_prefill_impl.inc", + "b1024/regions/stage0/native/stage0/predictor_mk2.cu", + "b1024/regions/stage0/stage0.cu", + "b1024/stage0.cu", + "inflight/progress.cuh", + "inflight/runtime.cuh" + ], + "output": "programs/b1024/artifacts/kernels/b1024/inflight.so" + }, + "b2048/stage0_inflight": { + "width": 2048, + "role": "stage0_inflight", + "source": "b2048/stage0.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v", + "-DMK2_INFLIGHT=1" + ], + "include_dirs": [ + "b2048", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b2048/regions/stage0/native", + "b2048/regions/prefill/native" + ], + "dependencies": [ + "b2048/regions/stage0/native/stage0/mk2_abi.cuh", + "b2048/regions/stage0/native/stage0/native_prefill_abi.cuh", + "b2048/regions/stage0/native/stage0/native_prefill_impl.inc", + "b2048/regions/stage0/native/stage0/predictor_mk2.cu", + "b2048/regions/stage0/stage0.cu", + "b2048/stage0.cu", + "inflight/progress.cuh", + "inflight/runtime.cuh" + ], + "output": "programs/b2048/artifacts/kernels/b2048/inflight.so" + }, + "b4096/stage0_inflight": { + "width": 4096, + "role": "stage0_inflight", + "source": "b4096/stage0.cu", + "flags": [ + "-shared", + "-O3", + "--use_fast_math", + "-std=c++20", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-DKITTENS_SM90", + "-diag-suppress=2361", + "-gencode=arch=compute_90a,code=sm_90a", + "-rdc=false", + "-lineinfo", + "-Xcompiler=-O3,-std=c++20", + "-Xcompiler=-fPIC", + "-lcuda", + "-lcudadevrt", + "-ldl", + "-Xlinker=-Bsymbolic", + "-Xptxas=-v", + "-DMK2_INFLIGHT=1" + ], + "include_dirs": [ + "b4096", + "third_party/thunderkittens/include", + "third_party/thunderkittens/prototype", + "b4096/regions/stage0/native", + "b4096/regions/prefill/native", + "b4096/regions/stage1/native" + ], + "dependencies": [ + "b4096/regions/stage0/native/stage0/mk2_abi.cuh", + "b4096/regions/stage0/native/stage0/native_prefill_abi.cuh", + "b4096/regions/stage0/native/stage0/native_prefill_impl.inc", + "b4096/regions/stage0/native/stage0/predictor_mk2.cu", + "b4096/regions/stage0/stage0.cu", + "b4096/stage0.cu", + "inflight/progress.cuh", + "inflight/runtime.cuh" + ], + "output": "programs/b4096/artifacts/kernels/b4096/inflight.so" + } + }, + "files": { + "b1024/prefill.cu": "e4012b41741ddeee7e62e48a170408d30e4b3a462f7c8add6245823f2745fa6f", + "b1024/regions/prefill/native/prefill/dedup_native_prefill_impl.inc": "3fd42d110796793798eb037c18abc9027a44c856630f4bd84490d2543546eeba", + "b1024/regions/prefill/native/prefill/dedup_predictor_mk2.cuh": "e876d9afb263b6a13cfd7b27b940a93b4e0451d8505bc53686e8181595733eaf", + "b1024/regions/prefill/native/prefill/mk2_abi.cuh": "03b263ec47f68bd0377a1ee8379fbe2e696ea22e806a4f13518ab9ff4855a0b9", + "b1024/regions/prefill/native/prefill/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "b1024/regions/prefill/native/prefill/prefill.cu": "da3c81ddecf8ab953e4a847cb6cb0d629d0c0f3e472f11b7624155f131160ac5", + "b1024/regions/prefill/prefill.cu": "10adeb7015313ab047cc5637b9f2fc85477a6e838703bf8d9921b29fee66257c", + "b1024/regions/stage0/native/stage0/mk2_abi.cuh": "84b6eec09584db81d2f18cf7eecd7497bb19e559c8cf767e67a9eb9199504584", + "b1024/regions/stage0/native/stage0/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "b1024/regions/stage0/native/stage0/native_prefill_impl.inc": "3ec9966d1cbefb16794547dca92632828480c4d60623bee535caaf47b25ff36e", + "b1024/regions/stage0/native/stage0/predictor_mk2.cu": "f5234f80d2268dbafad035a712c30de1cb966ea8a56b5a10e3b8926fdeb12643", + "b1024/regions/stage0/stage0.cu": "527d758cc6effc1dd34cac57da9f5db42b6da2e77fe2978b370618e98ce59bb0", + "b1024/stage0.cu": "d1c8222e9ef5c87e6e2b69f7d143004a277b77aa9541085532b31c1ce6ae752e", + "b128/prefill.cu": "e4012b41741ddeee7e62e48a170408d30e4b3a462f7c8add6245823f2745fa6f", + "b128/regions/prefill/native/prefill/dedup_native_prefill_impl.inc": "3fd42d110796793798eb037c18abc9027a44c856630f4bd84490d2543546eeba", + "b128/regions/prefill/native/prefill/dedup_predictor_mk2.cuh": "e876d9afb263b6a13cfd7b27b940a93b4e0451d8505bc53686e8181595733eaf", + "b128/regions/prefill/native/prefill/mk2_abi.cuh": "03b263ec47f68bd0377a1ee8379fbe2e696ea22e806a4f13518ab9ff4855a0b9", + "b128/regions/prefill/native/prefill/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "b128/regions/prefill/native/prefill/prefill.cu": "4598cf0aaf84e7cf35c687ea60ec8b1d58a927dad1025fb6d5ae05bba236dcb0", + "b128/regions/prefill/prefill.cu": "2f5971bc419952e640e1710f04c18588e85bcdaaa2db8cfc9ed57ebb58968fe5", + "b128/regions/stage0/native/stage0/mk2_abi.cuh": "84b6eec09584db81d2f18cf7eecd7497bb19e559c8cf767e67a9eb9199504584", + "b128/regions/stage0/native/stage0/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "b128/regions/stage0/native/stage0/native_prefill_impl.inc": "3ec9966d1cbefb16794547dca92632828480c4d60623bee535caaf47b25ff36e", + "b128/regions/stage0/native/stage0/predictor_mk2.cu": "96b7f8f67b683c561dcdd9b6c08d0dd4a6a3dce7861d2862547426740e52d353", + "b128/regions/stage0/stage0.cu": "667ee80b78ce45a77d8fde8151e898018b94eaaaa167fb30f4a165fb33e370c5", + "b128/stage0.cu": "d1c8222e9ef5c87e6e2b69f7d143004a277b77aa9541085532b31c1ce6ae752e", + "b2048/prefill.cu": "e4012b41741ddeee7e62e48a170408d30e4b3a462f7c8add6245823f2745fa6f", + "b2048/regions/prefill/native/prefill/dedup_native_prefill_impl.inc": "3fd42d110796793798eb037c18abc9027a44c856630f4bd84490d2543546eeba", + "b2048/regions/prefill/native/prefill/dedup_predictor_mk2.cuh": "e876d9afb263b6a13cfd7b27b940a93b4e0451d8505bc53686e8181595733eaf", + "b2048/regions/prefill/native/prefill/mk2_abi.cuh": "03b263ec47f68bd0377a1ee8379fbe2e696ea22e806a4f13518ab9ff4855a0b9", + "b2048/regions/prefill/native/prefill/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "b2048/regions/prefill/native/prefill/prefill.cu": "3be52144944eff16562fb597b0f89146dfc3f4304a0acb15032bf233b6a915d5", + "b2048/regions/prefill/prefill.cu": "fddd9d7cea0d58c8c0ddb7cc550487c4b7750837fa82e87220143fe952a4daa8", + "b2048/regions/stage0/native/stage0/mk2_abi.cuh": "84b6eec09584db81d2f18cf7eecd7497bb19e559c8cf767e67a9eb9199504584", + "b2048/regions/stage0/native/stage0/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "b2048/regions/stage0/native/stage0/native_prefill_impl.inc": "3ec9966d1cbefb16794547dca92632828480c4d60623bee535caaf47b25ff36e", + "b2048/regions/stage0/native/stage0/predictor_mk2.cu": "80336e56373eebbe0dcb351930039263302c4894b343c4d3b52c056d2d561b5c", + "b2048/regions/stage0/stage0.cu": "a53ed90565797634a5286d2bda5f71683f5c4797c9755e6cafc65effee996f1e", + "b2048/stage0.cu": "d1c8222e9ef5c87e6e2b69f7d143004a277b77aa9541085532b31c1ce6ae752e", + "b256/prefill.cu": "e4012b41741ddeee7e62e48a170408d30e4b3a462f7c8add6245823f2745fa6f", + "b256/regions/prefill/native/prefill/dedup_native_prefill_impl.inc": "3fd42d110796793798eb037c18abc9027a44c856630f4bd84490d2543546eeba", + "b256/regions/prefill/native/prefill/dedup_predictor_mk2.cuh": "e876d9afb263b6a13cfd7b27b940a93b4e0451d8505bc53686e8181595733eaf", + "b256/regions/prefill/native/prefill/mk2_abi.cuh": "03b263ec47f68bd0377a1ee8379fbe2e696ea22e806a4f13518ab9ff4855a0b9", + "b256/regions/prefill/native/prefill/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "b256/regions/prefill/native/prefill/prefill.cu": "5e9f9057b39d42b00cf6121867687a0e02e734c7d179982589b7a23c1bf370fa", + "b256/regions/prefill/prefill.cu": "98c820b73f0e43be0b4daeae49c618c28660cdca856e19f4909c79722c3dda22", + "b256/regions/stage0/native/stage0/mk2_abi.cuh": "84b6eec09584db81d2f18cf7eecd7497bb19e559c8cf767e67a9eb9199504584", + "b256/regions/stage0/native/stage0/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "b256/regions/stage0/native/stage0/native_prefill_impl.inc": "3ec9966d1cbefb16794547dca92632828480c4d60623bee535caaf47b25ff36e", + "b256/regions/stage0/native/stage0/predictor_mk2.cu": "96b7f8f67b683c561dcdd9b6c08d0dd4a6a3dce7861d2862547426740e52d353", + "b256/regions/stage0/stage0.cu": "c9ec1d440ff3f9a61048399066753a5dfa844114bebb99119c087a4e9ec9bb85", + "b256/stage0.cu": "d1c8222e9ef5c87e6e2b69f7d143004a277b77aa9541085532b31c1ce6ae752e", + "b32/codec.cu": "30ac9245e757171491565bd9e6ef6d439cad3dbbf3b9328c84bb41adebf9c165", + "b32/native/codec/kernel_kf_codec.h": "d7106461a0e6510ae3530377a0d8aea72f2178b9f4ec35ae44198bcb66ba1270", + "b32/native/codec/legacy_codec.cuh": "a60b890c3b65d6e065945c5054b33de01567c3a3b12eb83f5511825513e49eed", + "b32/native/codec/native_stage1.cu": "e638009f1ec9ee9460e2b5622adcd129a2016bf71884259875872836b16790e4", + "b32/native/codec/native_stage1.h": "acabf975a8525c4dca3f237585930acea94c4976b6b47a000988db5e46ed9e97", + "b32/native/codec/repack3.inc": "32250e761b2cd976ccb5deff397f52f9ff5264f2e7d38595d08aeb90c53f3d3f", + "b32/native/codec/tail_kernel.inc": "9530bfd8d7e44b8e82e090a756940afcf1a21aa56f9fd83fd647c447cc448a81", + "b32/native/codec/wg2_kernel.inc": "c235ef28d4eeab2213cb62d7b6634aec6049f003763aca70f408aed4c269f01d", + "b32/native/codec/wg3_kernel.inc": "35c5367c9d84e79cf49c7c836723a63ed5c7f76c086b60dde79b793ed0d4a5b3", + "b32/native/codec/wg4_step_widths.inc": "4845ae88f89b6ad7c0069b48c0489cfcda9522eba06876307a294fbd7cee92d9", + "b32/native/codec/wg5_step_widths.inc": "77da92f8ada4299016a352ebd416533fd425998296f57c907d21faf9a4e71365", + "b32/native/codec/wg_kernel.inc": "0efcc685572d02618855a0348b766b17bc74038714e6ba52eeeeaf97c1b05cc7", + "b32/native/codec/wgmma_macro.inc": "b39bedc91469a871ea6801a0bf89093b2f2ade3c643285a95ead82912996d228", + "b32/native/prefill/dedup_native_prefill_impl.inc": "1bf52676e894069d93cdd7d119909dafc795432542bcf9a96204df4105370be3", + "b32/native/prefill/dedup_predictor_mk2.cuh": "7210b152a4bb75a06ddfdcd30dad06430f3246025db1b3ff7f56e550956128dc", + "b32/native/prefill/mk2_abi.cuh": "03b263ec47f68bd0377a1ee8379fbe2e696ea22e806a4f13518ab9ff4855a0b9", + "b32/native/prefill/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "b32/native/prefill/prefill.cu": "aeac0d899e21085a6b8028e492346ecd69f3b462e628abc7c15530da194a06c8", + "b32/native/stage0/mk2_abi.cuh": "d59126ce70f6d91a42baa7ca13f7c5ff40aef9b4ad73b6afe1f3773525ec045d", + "b32/native/stage0/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "b32/native/stage0/native_prefill_impl.inc": "3ec9966d1cbefb16794547dca92632828480c4d60623bee535caaf47b25ff36e", + "b32/native/stage0/predictor_mk2.cu": "db2798ac8320cf49143af5a3b30b66a10bca8f166fabc39288b6b59cce16ac24", + "b32/native/stage0/sparse_moe_impl.cuh": "82751d425742f615df2c9fe3bc7c4aaa932d7c2b10931d30d19ce516facd7a22", + "b32/prefill.cu": "9908e718bc2c3f63d4c69fb43c2e59ac8d196ffa2d2b2fe1567638d9ebac1325", + "b32/stage0.cu": "67c7245729d2cb39b53850f0615f695a797b58cf7a991b06912527b1822bb91e", + "b4096/codec.cu": "8d32b92eb209a6ee2569a9c78958e3ab8cdca7f96dcdfb7a3098828bd3e0fa4b", + "b4096/prefill.cu": "e4012b41741ddeee7e62e48a170408d30e4b3a462f7c8add6245823f2745fa6f", + "b4096/regions/prefill/native/prefill/dedup_native_prefill_impl.inc": "9536807a0546de42f06a254ab45daa06eda8ee39eff67ea5556abdf3d2b03e78", + "b4096/regions/prefill/native/prefill/dedup_predictor_mk2.cuh": "e876d9afb263b6a13cfd7b27b940a93b4e0451d8505bc53686e8181595733eaf", + "b4096/regions/prefill/native/prefill/mk2_abi.cuh": "03b263ec47f68bd0377a1ee8379fbe2e696ea22e806a4f13518ab9ff4855a0b9", + "b4096/regions/prefill/native/prefill/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "b4096/regions/prefill/native/prefill/prefill.cu": "21c76e3bf6e447aaf86513b53d1815656079f874b11c3a6133b7815f31c7faff", + "b4096/regions/prefill/prefill.cu": "470084d73160937d17c2665b6dc26b4cc89aaba984d6043c5c878ccc8912ed08", + "b4096/regions/stage0/native/stage0/mk2_abi.cuh": "84b6eec09584db81d2f18cf7eecd7497bb19e559c8cf767e67a9eb9199504584", + "b4096/regions/stage0/native/stage0/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "b4096/regions/stage0/native/stage0/native_prefill_impl.inc": "3ec9966d1cbefb16794547dca92632828480c4d60623bee535caaf47b25ff36e", + "b4096/regions/stage0/native/stage0/predictor_mk2.cu": "467c5ab8fab594e3aed2082a1a3ae9deca9fab9de38387a09cc3664f9655f4a1", + "b4096/regions/stage0/stage0.cu": "2522442a92f3836849690031abf3ed5d1e8c331bd134465a0953d83ac3bb0762", + "b4096/regions/stage1/codec.cu": "30ac9245e757171491565bd9e6ef6d439cad3dbbf3b9328c84bb41adebf9c165", + "b4096/regions/stage1/native/codec/kernel_kf_codec.h": "d7106461a0e6510ae3530377a0d8aea72f2178b9f4ec35ae44198bcb66ba1270", + "b4096/regions/stage1/native/codec/legacy_codec.cuh": "5f9df78d7f9cf3b1159a69bdd1253d9ac48f232f004f205d9bce4f030b8aa1e8", + "b4096/regions/stage1/native/codec/native_stage1.cu": "c9ecbb9e0334ddfe120adeed05e292353073bba8b5998c84eae77f5e6deb1a8e", + "b4096/regions/stage1/native/codec/native_stage1.h": "acabf975a8525c4dca3f237585930acea94c4976b6b47a000988db5e46ed9e97", + "b4096/regions/stage1/native/codec/repack3.inc": "32250e761b2cd976ccb5deff397f52f9ff5264f2e7d38595d08aeb90c53f3d3f", + "b4096/regions/stage1/native/codec/tail_kernel.inc": "9530bfd8d7e44b8e82e090a756940afcf1a21aa56f9fd83fd647c447cc448a81", + "b4096/regions/stage1/native/codec/wg2_kernel.inc": "b7eb02869fa0a68a786656d02f21703037c96ac0e5e7993f853495dd6e657da5", + "b4096/regions/stage1/native/codec/wg3_kernel.inc": "35c5367c9d84e79cf49c7c836723a63ed5c7f76c086b60dde79b793ed0d4a5b3", + "b4096/regions/stage1/native/codec/wg4_step_widths.inc": "d71f5b887655da190336af52f9f070d34cad9bd94a49205c85e631fc67cabd76", + "b4096/regions/stage1/native/codec/wg5_step_widths.inc": "fe0c9a2935961f693ab180b9ca142c4e293ce196de61512ab1cee93ccff2447f", + "b4096/regions/stage1/native/codec/wg_kernel.inc": "0efcc685572d02618855a0348b766b17bc74038714e6ba52eeeeaf97c1b05cc7", + "b4096/regions/stage1/native/codec/wgmma_macro.inc": "b39bedc91469a871ea6801a0bf89093b2f2ade3c643285a95ead82912996d228", + "b4096/stage0.cu": "d1c8222e9ef5c87e6e2b69f7d143004a277b77aa9541085532b31c1ce6ae752e", + "b512/codec.cu": "30ac9245e757171491565bd9e6ef6d439cad3dbbf3b9328c84bb41adebf9c165", + "b512/native/codec/kernel_kf_codec.h": "d7106461a0e6510ae3530377a0d8aea72f2178b9f4ec35ae44198bcb66ba1270", + "b512/native/codec/legacy_codec.cuh": "af5fb2d4e67d2ff5f60d211d6ad06116672a20aa9b805321a165838b8154d37d", + "b512/native/codec/native_stage1.cu": "423c6f8e0bad38c374dd23732748fdfe8f2ff6daabe7de39516cdd8cdedc2815", + "b512/native/codec/native_stage1.h": "acabf975a8525c4dca3f237585930acea94c4976b6b47a000988db5e46ed9e97", + "b512/native/codec/repack3.inc": "32250e761b2cd976ccb5deff397f52f9ff5264f2e7d38595d08aeb90c53f3d3f", + "b512/native/codec/tail_kernel.inc": "62e673a505da1e09ababaf58385d53e5396cc81a0c0df364e8381cf236962d58", + "b512/native/codec/wg2_kernel.inc": "b7eb02869fa0a68a786656d02f21703037c96ac0e5e7993f853495dd6e657da5", + "b512/native/codec/wg3_kernel.inc": "35c5367c9d84e79cf49c7c836723a63ed5c7f76c086b60dde79b793ed0d4a5b3", + "b512/native/codec/wg4_step_widths.inc": "0877f7ac9ad705a20b6223266ed36fe9d6766afd2b13c88462ea319ecaada6ee", + "b512/native/codec/wg5_step_widths.inc": "6a7427f82f2cc7968395f1a04da654349ba5f07a8b1092e8141b3ef906f89dba", + "b512/native/codec/wg_kernel.inc": "0efcc685572d02618855a0348b766b17bc74038714e6ba52eeeeaf97c1b05cc7", + "b512/native/codec/wgmma_macro.inc": "b39bedc91469a871ea6801a0bf89093b2f2ade3c643285a95ead82912996d228", + "b512/native/prefill/dedup_native_prefill_impl.inc": "3387cb8ce201d59264375839e2c6a5be343205d3266a851864ccb4168b9b7ef3", + "b512/native/prefill/dedup_predictor_mk2.cuh": "28341b82853a7598aa4ff38654c7e5fe1561df4edf9507fab539bbe6fdfb6fae", + "b512/native/prefill/mk2_abi.cuh": "03b263ec47f68bd0377a1ee8379fbe2e696ea22e806a4f13518ab9ff4855a0b9", + "b512/native/prefill/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "b512/native/prefill/prefill.cu": "ebae1fbc72a64c37c8037cbb52ca0be6e295747472afbdd9e2c001ea49fa9c7a", + "b512/native/stage0/mk2_abi.cuh": "84b6eec09584db81d2f18cf7eecd7497bb19e559c8cf767e67a9eb9199504584", + "b512/native/stage0/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "b512/native/stage0/native_prefill_impl.inc": "3ec9966d1cbefb16794547dca92632828480c4d60623bee535caaf47b25ff36e", + "b512/native/stage0/predictor_mk2.cu": "f4536289eb58147bf467316819ac7a88d26584a7fdceeadc2815df1e27f7a823", + "b512/prefill.cu": "004bec3ce9fe55b0452ff88d885ed7e6f9ed688f147a499bc74b0847832520fb", + "b512/stage0.cu": "a35402c5c9b6e0c65a4ebf9beb74ba1308c3f26e9a5dd2885dd4b5d229027efb", + "codec/base/kernel_kf_codec.h": "d7106461a0e6510ae3530377a0d8aea72f2178b9f4ec35ae44198bcb66ba1270", + "codec/base/legacy_codec.cuh": "faeec72569687e324c8699e2df7e3aae1cbce051c81476c73b2805b3491efb20", + "codec/base/native_stage1.cu": "5dccbe3ff041fde33f500218cb6d495adfec3dd6dbba57e14e9c363cdd86e807", + "codec/base/native_stage1.h": "acabf975a8525c4dca3f237585930acea94c4976b6b47a000988db5e46ed9e97", + "codec/base/repack3.inc": "32250e761b2cd976ccb5deff397f52f9ff5264f2e7d38595d08aeb90c53f3d3f", + "codec/base/tail_kernel.inc": "9530bfd8d7e44b8e82e090a756940afcf1a21aa56f9fd83fd647c447cc448a81", + "codec/base/wg2_kernel.inc": "b7eb02869fa0a68a786656d02f21703037c96ac0e5e7993f853495dd6e657da5", + "codec/base/wg3_kernel.inc": "35c5367c9d84e79cf49c7c836723a63ed5c7f76c086b60dde79b793ed0d4a5b3", + "codec/base/wg4_step_widths.inc": "4845ae88f89b6ad7c0069b48c0489cfcda9522eba06876307a294fbd7cee92d9", + "codec/base/wg5_step_widths.inc": "5ea6b07a1410057a8c9a431a8ac61b1f447c855b3e7d9705b23ff19252b67574", + "codec/base/wg_kernel.inc": "0efcc685572d02618855a0348b766b17bc74038714e6ba52eeeeaf97c1b05cc7", + "codec/base/wgmma_macro.inc": "b39bedc91469a871ea6801a0bf89093b2f2ade3c643285a95ead82912996d228", + "codec/first_packet/codec.cu": "8d32b92eb209a6ee2569a9c78958e3ab8cdca7f96dcdfb7a3098828bd3e0fa4b", + "codec/first_packet/regions/stage1/codec.cu": "30ac9245e757171491565bd9e6ef6d439cad3dbbf3b9328c84bb41adebf9c165", + "codec/first_packet/regions/stage1/native/codec/kernel_kf_codec.h": "d7106461a0e6510ae3530377a0d8aea72f2178b9f4ec35ae44198bcb66ba1270", + "codec/first_packet/regions/stage1/native/codec/legacy_codec.cuh": "932a87f714b990948b7d293ef35581596734aa03360be25fea921c90bf7bb243", + "codec/first_packet/regions/stage1/native/codec/native_stage1.cu": "b3bc0075273fdea6a8199c697110bd84bf9f67d753f35c48c3d9329a5fee9ed8", + "codec/first_packet/regions/stage1/native/codec/native_stage1.h": "acabf975a8525c4dca3f237585930acea94c4976b6b47a000988db5e46ed9e97", + "codec/first_packet/regions/stage1/native/codec/repack3.inc": "32250e761b2cd976ccb5deff397f52f9ff5264f2e7d38595d08aeb90c53f3d3f", + "codec/first_packet/regions/stage1/native/codec/tail_kernel.inc": "62e673a505da1e09ababaf58385d53e5396cc81a0c0df364e8381cf236962d58", + "codec/first_packet/regions/stage1/native/codec/wg2_kernel.inc": "b7eb02869fa0a68a786656d02f21703037c96ac0e5e7993f853495dd6e657da5", + "codec/first_packet/regions/stage1/native/codec/wg3_kernel.inc": "35c5367c9d84e79cf49c7c836723a63ed5c7f76c086b60dde79b793ed0d4a5b3", + "codec/first_packet/regions/stage1/native/codec/wg4_step_widths.inc": "a8a4d1bcdef4d723ffa8f3dcedd4f301b75b8eb11fcea051cf990abe1bc84f6f", + "codec/first_packet/regions/stage1/native/codec/wg5_step_widths.inc": "dede5279dd92ffda8d58c63adf5f3d143e8522ff483538e55cbee55095cca938", + "codec/first_packet/regions/stage1/native/codec/wg_kernel.inc": "0efcc685572d02618855a0348b766b17bc74038714e6ba52eeeeaf97c1b05cc7", + "codec/first_packet/regions/stage1/native/codec/wgmma_macro.inc": "b39bedc91469a871ea6801a0bf89093b2f2ade3c643285a95ead82912996d228", + "codec/small/kernel_kf_codec.h": "d7106461a0e6510ae3530377a0d8aea72f2178b9f4ec35ae44198bcb66ba1270", + "codec/small/legacy_codec.cuh": "faeec72569687e324c8699e2df7e3aae1cbce051c81476c73b2805b3491efb20", + "codec/small/native_stage1.cu": "aa58aaa759728540813e61c60e0d32b548b6af20d1909fd9cea7ba9a9dec3265", + "codec/small/native_stage1.h": "acabf975a8525c4dca3f237585930acea94c4976b6b47a000988db5e46ed9e97", + "codec/small/repack3.inc": "32250e761b2cd976ccb5deff397f52f9ff5264f2e7d38595d08aeb90c53f3d3f", + "codec/small/tail_kernel.inc": "9530bfd8d7e44b8e82e090a756940afcf1a21aa56f9fd83fd647c447cc448a81", + "codec/small/wg2_kernel.inc": "b7eb02869fa0a68a786656d02f21703037c96ac0e5e7993f853495dd6e657da5", + "codec/small/wg3_kernel.inc": "35c5367c9d84e79cf49c7c836723a63ed5c7f76c086b60dde79b793ed0d4a5b3", + "codec/small/wg4_step_widths.inc": "4845ae88f89b6ad7c0069b48c0489cfcda9522eba06876307a294fbd7cee92d9", + "codec/small/wg5_step_widths.inc": "5ea6b07a1410057a8c9a431a8ac61b1f447c855b3e7d9705b23ff19252b67574", + "codec/small/wg_kernel.inc": "0efcc685572d02618855a0348b766b17bc74038714e6ba52eeeeaf97c1b05cc7", + "codec/small/wgmma_macro.inc": "b39bedc91469a871ea6801a0bf89093b2f2ade3c643285a95ead82912996d228", + "contraction/bf16/mk2_abi.cuh": "2d326d7ef226810bd8f22d1a5d8269fbd925388b95e63d88d39bf7f819e2b84f", + "contraction/bf16/native_prefill_abi.cuh": "01d10a274b74b381151f3884ab4ff9b7b1df4ba5ba6a739e9bf26cdb93fcbfb3", + "contraction/bf16/native_prefill_impl.inc": "cd50fc333849f8cd7c41f44abb836d36e6539b74b91bf87fc848dadcc3555fe5", + "contraction/bf16/predictor_mk2.cu": "7440c2ce45c9f27f7f54efeaea4eac13cb261bc21ab4266ab033028b1206eada", + "contraction/bf16/sparse_moe_impl.cuh": "f09100a2b339a8419d726cfd30fa158b74d2d1dd6c0b78c84f006f94a1016bf6", + "contraction/lane_transfer.cu": "f3dd2f3e84687c55e7d69e499e961d4db8d84551813c60a0e1a1a5cc1b788ffc", + "contraction/mk2_abi.cuh": "2d326d7ef226810bd8f22d1a5d8269fbd925388b95e63d88d39bf7f819e2b84f", + "inflight/progress.cuh": "8aea468073324060049bfd7c26ae949af9be941fc82649e9a52c382bea5dd2ea", + "inflight/runtime.cuh": "159e0355654292830e857e93b59bdf6f97f3641c40c4f3543ef712bec81fd810", + "prefill/b1/dedup_native_prefill_impl.inc": "e4b2dfc229499feca0fa4092e7b0b3378781ee2e27d3515e1dcba8a06888c3d6", + "prefill/b1/dedup_predictor_mk2.cuh": "f3eab0595791b524c52bc607eca22d90ae9662f391ccd9a9145694b7dc4430a6", + "prefill/b1/mk2_abi.cuh": "03b263ec47f68bd0377a1ee8379fbe2e696ea22e806a4f13518ab9ff4855a0b9", + "prefill/b1/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "prefill/b1/prefill.cu": "a78f95b0f678baa55e61bdf7940ae9a562eb0e44beb613da1e982fd5daac8117", + "prefill/b16/dedup_native_prefill_impl.inc": "e4b2dfc229499feca0fa4092e7b0b3378781ee2e27d3515e1dcba8a06888c3d6", + "prefill/b16/dedup_predictor_mk2.cuh": "f3eab0595791b524c52bc607eca22d90ae9662f391ccd9a9145694b7dc4430a6", + "prefill/b16/mk2_abi.cuh": "03b263ec47f68bd0377a1ee8379fbe2e696ea22e806a4f13518ab9ff4855a0b9", + "prefill/b16/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "prefill/b16/prefill.cu": "9d81336c17c866e5087130eb9502ac6ecf2872139c584e137bc5e6f269a44b3e", + "prefill/b2/dedup_native_prefill_impl.inc": "2817210269730783c849453a86087d5ac86d328961c0fbf4728f08a4271ec3ab", + "prefill/b2/dedup_predictor_mk2.cuh": "f3eab0595791b524c52bc607eca22d90ae9662f391ccd9a9145694b7dc4430a6", + "prefill/b2/mk2_abi.cuh": "03b263ec47f68bd0377a1ee8379fbe2e696ea22e806a4f13518ab9ff4855a0b9", + "prefill/b2/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "prefill/b2/prefill.cu": "715ec30d321d9d02f1689bcb2ef733b1c27435050cec10d86f7c0cc0a6a1acaf", + "prefill/b4/dedup_native_prefill_impl.inc": "2817210269730783c849453a86087d5ac86d328961c0fbf4728f08a4271ec3ab", + "prefill/b4/dedup_predictor_mk2.cuh": "f3eab0595791b524c52bc607eca22d90ae9662f391ccd9a9145694b7dc4430a6", + "prefill/b4/mk2_abi.cuh": "03b263ec47f68bd0377a1ee8379fbe2e696ea22e806a4f13518ab9ff4855a0b9", + "prefill/b4/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "prefill/b4/prefill.cu": "4ece2fb3df94c0d8dece734dc50485a2d5d22a335fdf52a993fa52a4cac02d05", + "prefill/b64/dedup_native_prefill_impl.inc": "e4b2dfc229499feca0fa4092e7b0b3378781ee2e27d3515e1dcba8a06888c3d6", + "prefill/b64/dedup_predictor_mk2.cuh": "f3eab0595791b524c52bc607eca22d90ae9662f391ccd9a9145694b7dc4430a6", + "prefill/b64/mk2_abi.cuh": "03b263ec47f68bd0377a1ee8379fbe2e696ea22e806a4f13518ab9ff4855a0b9", + "prefill/b64/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "prefill/b64/prefill.cu": "163ca6fc96f168e75dc1ba1dd9e0219466afe37d549c8efe4e47e7b099c70fe3", + "prefill/b8/dedup_native_prefill_impl.inc": "2817210269730783c849453a86087d5ac86d328961c0fbf4728f08a4271ec3ab", + "prefill/b8/dedup_predictor_mk2.cuh": "f3eab0595791b524c52bc607eca22d90ae9662f391ccd9a9145694b7dc4430a6", + "prefill/b8/mk2_abi.cuh": "03b263ec47f68bd0377a1ee8379fbe2e696ea22e806a4f13518ab9ff4855a0b9", + "prefill/b8/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "prefill/b8/prefill.cu": "6dafe5f34ee31df74b115995c074d845f8c6aab23d8b02afaf173dbc715da65b", + "stage0/b1/bb_router_selection.cuh": "52bd3ec25f8a791e3c8580ff99f7d0c3f66331db0f1289df6ddbfb21dc0d15e5", + "stage0/b1/mk2_abi.cuh": "84b6eec09584db81d2f18cf7eecd7497bb19e559c8cf767e67a9eb9199504584", + "stage0/b1/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "stage0/b1/native_prefill_impl.inc": "3ec9966d1cbefb16794547dca92632828480c4d60623bee535caaf47b25ff36e", + "stage0/b1/predictor_mk2.cu": "a049e6b0110159f096721e63babb7d3f7edfb5ce9a4662fc641403c645d44dbb", + "stage0/b16/mk2_abi.cuh": "84b6eec09584db81d2f18cf7eecd7497bb19e559c8cf767e67a9eb9199504584", + "stage0/b16/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "stage0/b16/native_prefill_impl.inc": "3ec9966d1cbefb16794547dca92632828480c4d60623bee535caaf47b25ff36e", + "stage0/b16/predictor_mk2.cu": "14aaeefd94a0c039a5f82169410da9b0a4e77599e5896fe8609a27e52c8a85d2", + "stage0/b16/sparse_moe_impl.cuh": "18e739ef359c8ea8d525c0bad25a5775ed260b9029a4b4b14e3209b14307cc4c", + "stage0/b64/mk2_abi.cuh": "84b6eec09584db81d2f18cf7eecd7497bb19e559c8cf767e67a9eb9199504584", + "stage0/b64/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "stage0/b64/native_prefill_impl.inc": "3ec9966d1cbefb16794547dca92632828480c4d60623bee535caaf47b25ff36e", + "stage0/b64/predictor_mk2.cu": "b5b210b8b3cf8978995ee47e701396e437d81d8486d9761fb4b84f314500576c", + "stage0/small/mk2_abi.cuh": "84b6eec09584db81d2f18cf7eecd7497bb19e559c8cf767e67a9eb9199504584", + "stage0/small/native_prefill_abi.cuh": "3c2d96461c9bc71e26e829ba0bf309b2cb7ed2b449665a2e2264146c96b6704e", + "stage0/small/native_prefill_impl.inc": "f012fc3d6b6c1a2ca1ecbc7ded02555160abf3e08b851268f7de9cc08663adea", + "stage0/small/predictor_mk2.cu": "5fcb97529ff1bfb361959f55beacb801869683b229c59e80d62cdca103016493", + "stage0/small/sparse_moe_impl.cuh": "18e739ef359c8ea8d525c0bad25a5775ed260b9029a4b4b14e3209b14307cc4c", + "third_party/thunderkittens-v4-prefill/LICENSE": "e705a1075640344d86b3b3d4f27dc06bf37cb5bf3b24676a9ca49a19c327ddae", + "third_party/thunderkittens-v4-prefill/include/common/base_ops.cuh": "c33d4ba051c0a7bc27878809263d75f06167ed2e781ba781372d0d90c95373c5", + "third_party/thunderkittens-v4-prefill/include/common/base_types.cuh": "6d09ee3d85bb952ae88fa9aeb79c9403717f4b473737f4a2e28055fe6fbac92f", + "third_party/thunderkittens-v4-prefill/include/common/common.cuh": "486ded03bbe43ca4ed3fcce3988585c954494f8366ecbbfec088d215aaf5ca3e", + "third_party/thunderkittens-v4-prefill/include/common/multimem.cuh": "a8aae2e97b5934f3849c9cb9ad5ef7994ca283eaff0c202d37c283636c0f3181", + "third_party/thunderkittens-v4-prefill/include/common/util.cuh": "c715f83d149160261bfeed56c5d1ce4b4d1e225f156d4f3a532cc694655260ff", + "third_party/thunderkittens-v4-prefill/include/kittens.cuh": "437d60a0d1c215019cc8f5e6b31c02043f1fa7c18f4bac4a6f95a94022f77715", + "third_party/thunderkittens-v4-prefill/include/ops/group/group.cuh": "8c33856f083a99735807d82d48ee64192da4dfa6f6ea674267fec210620838ad", + "third_party/thunderkittens-v4-prefill/include/ops/group/memory/memory.cuh": "d325ffeb850d7c2c942ac90e49e9cea37d4e565897682bfdd1270470c35733f6", + "third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/complex_global_to_register.cuh": "cbf4b20658b85a721903f2a82faeb9aea826b30ef9546bc2087f71fb9843567a", + "third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/complex_global_to_shared.cuh": "33a417e594becfc9e6284b304293cf5aba6a9d9e2fac8706442f466da0eca26d", + "third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/complex_shared_to_register.cuh": "ffa32c513bf74ffad6f19f5aff8ca596e514e6a95e050babd1008207c3121c69", + "third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/global_to_register.cuh": "04fcac6fdb32a785b304752e00f7a2681cf34695de85c13783c60e61457efd9a", + "third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/global_to_shared.cuh": "266d97ee9a74b5c7c2706ecb9477e4b13b2ab52f8bc4324397a6c41a182282e0", + "third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/parallel_global_to_global.cuh": "36e0923ba349ab9fefb277685d260f221c5949118c0dcd5d4daffcf46b11eeaa", + "third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/shared_to_register.cuh": "ae047978ed23d2da0f690ae58b6e14c90c1aaa9b3d18fcda7b87f691bd4740e3", + "third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/tensor_to_register.cuh": "f15ababf491e4b04f63d8e14f8cc60df669984853ce71cd1da57246f92d4f6c9", + "third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/tile.cuh": "e0c270ec738a925281733ed9de450b518693d86652c018bba48baec767fe808f", + "third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/tma.cuh": "2892ebb566b922c70f4c16cc7a5555dd329b55b55a5a4987d9e226c142b6b5fa", + "third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/tma_cluster.cuh": "54bbdd60b28f7ad1cbec58867b3e3a0291f1aab3cc386a1938eaf969c0c19a1f", + "third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/global_to_register.cuh": "f784f3e7f43ff4ae299fa660a91ba9b0690538334ca7fa2a7386c2fe0de50fe7", + "third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/global_to_shared.cuh": "b9e0483deeb85e0887f92443d8f886148c2cdf491b44792f13056e94656f60e1", + "third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/shared_to_register.cuh": "55f3546b89e00f8bc84f7805cd150e30cead7d2933e7992cb9ec720974d204e4", + "third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/tma.cuh": "4b93ba91d05f7defae9ea8c66da29ad687d0780ae148d96271f14d60cadc3fbf", + "third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/tma_cluster.cuh": "00ac7bdfc5b5c0ca38b037f4e262b19fa73f4164921a56c303a33e49349e1743", + "third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/vec.cuh": "b7ee8a98c36dc495f2201f91ac9a3be5af88bf5f07a14961b84a304f655d7c3b", + "third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x112.impl": "a04d5742a0ea37c6a6ca9379f4912dc8ef7424372342add231867d98e96df232", + "third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x128.impl": "2508e03800c05c778f345f8c5fe96415b6fcae8075fa7110ffe5372d3497ade1", + "third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x144.impl": "7bcf837979c41e7aca984b94bff07fd13b22eb7a59fe2e48855cc07f2d4a9dbd", + "third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x16.impl": "25c681bdbb3d3fd4d5bc6fbc9c6070583ff08d593538248e0a0d8f58b18e1042", + "third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x160.impl": "56378c1989a098599697193a7bc163617cc4488de13a96b76fa7ef9daf4d8472", + "third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x176.impl": "72e9e4865b5fcf92d67fb37e4e3114c1493f4cdaa8738772021048e1ffc8845b", + "third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x192.impl": "91f2c3d0ed837be3706daa603ec445d5a6ad3c20112ccd1e8420b05e9a181391", + "third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x208.impl": "4676ec206611e28b447cae7f0da5e9d032b2e834ab37f0663dce4248f3933128", + "third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x224.impl": "2bccdb907ededc0ce9ed55dbc4b6c95801aa2f32c9c7edff0742c492f8df5a45", + "third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x240.impl": "2ef3ef8c0cf374f97962731777a4352f241ec7d964f6f54f3b45e719d16efd7b", + "third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x256.impl": "0c7515f4b02cf49ab10e4ece84e97cf3b822e5fed93287cde77db5f708c2dee1", + "third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x32.impl": "bbedf1bc2e24bc9f152c5f9ce029bb178a40c9f7c510077bccd3e5f925361f8e", + "third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x48.impl": "bf6fab4e0611376788664f47c5e295bcd58e09e3faf42b47a6531bae61cdd9ba", + "third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x64.impl": "fb5a176e96e579f834750e9d86f80a399b1aa0dd34a4991ff639da0eab1f6e9e", + "third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x80.impl": "15b0c22259f07fc17fe70bcb2341e5b673ecef55a37dc7aceca1299ed7022888", + "third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x96.impl": "948e526120979fa31c88f2232d5a433e7c8a0af7168ea1be01738f8d46f8d996", + "third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/base.cuh": "0d97d610d412ee7580d199c3180c94a6f467bbe1273f09b9a4dafd641f9353f9", + "third_party/thunderkittens-v4-prefill/include/ops/group/mma/mma.cuh": "015245bd39b75b54cb999328ab887c1d0d04541b6570cc9af24dd716d0a8d8f4", + "third_party/thunderkittens-v4-prefill/include/ops/group/mma/tcgen05.cuh": "a2cccdfbc9a9843ce10dc04d2a0811dbef9f65d4065433191b5384c27399898b", + "third_party/thunderkittens-v4-prefill/include/ops/group/mma/warp.cuh": "426346462672db5bee203a82f1fc98100bb580085db2d5b61868a5131ba854a4", + "third_party/thunderkittens-v4-prefill/include/ops/group/mma/warpgroup.cuh": "574a3227339eae9c4b88b2e2ad420ce47b7cf9eab08d49422984b93210cdf3b2", + "third_party/thunderkittens-v4-prefill/include/ops/group/register/register.cuh": "d325ffeb850d7c2c942ac90e49e9cea37d4e565897682bfdd1270470c35733f6", + "third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/complex_conversions.cuh": "6072843c18334a812d786ba021564986b07568369858b1c6c1c8926479edaa40", + "third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/complex_maps.cuh": "44ecc28580def5b189df2265b72fb66169ad1fb959af241f2215dc741ad0301b", + "third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/conversions.cuh": "cfc15594b75ffc80383ddf49fe565a99855629b07e969502f5da0c059597f9be", + "third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/maps.cuh": "42bda47c15e10c35176ba3868924b00255f033973dd8f07848b7627cea413c57", + "third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/reductions.cuh": "e8914eb995b169aa72104d2d6505acd27546812ded7089723404f5c99476c87a", + "third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/tile.cuh": "7eae6e436a2dd01cf1bc3dddce26349b4a7cfca509a3f3b841b4b6dcf4b14400", + "third_party/thunderkittens-v4-prefill/include/ops/group/register/vec/conversions.cuh": "e34b95719637bf0449781afa2a63b7d1066080af551515e05d70fcc72e50f130", + "third_party/thunderkittens-v4-prefill/include/ops/group/register/vec/maps.cuh": "e8caa7feffdf8324bb0a4b6df3a44af44f291bb26aace54c27c8900deb45f486", + "third_party/thunderkittens-v4-prefill/include/ops/group/register/vec/reductions.cuh": "395ee946b9288b2add58272851d3031bccd8aec5c8d9de452b42df61ff3f8c2d", + "third_party/thunderkittens-v4-prefill/include/ops/group/register/vec/vec.cuh": "c9325ef900273918f92d68f6853c92a51c2c2897ca1d69a8a8cea9892292563a", + "third_party/thunderkittens-v4-prefill/include/ops/group/shared/shared.cuh": "d325ffeb850d7c2c942ac90e49e9cea37d4e565897682bfdd1270470c35733f6", + "third_party/thunderkittens-v4-prefill/include/ops/group/shared/tile/conversions.cuh": "1464ef684ec8ed7ecb2459551c5968593b56177e60cad7ed1c6f137b6cd95916", + "third_party/thunderkittens-v4-prefill/include/ops/group/shared/tile/maps.cuh": "a35a5659cf24e8f11559e26d5fe4d969925f25e9a86b6c8f7087f0523b09adb9", + "third_party/thunderkittens-v4-prefill/include/ops/group/shared/tile/reductions.cuh": "fcb9ea4a52056a17619c71a41041ffff2175c3b13308afb92d8fff3ef8cb162e", + "third_party/thunderkittens-v4-prefill/include/ops/group/shared/tile/tile.cuh": "2c45011828ef5bf5622e390a56addd502e0f503baf7c8cd07311b23b2ec7a634", + "third_party/thunderkittens-v4-prefill/include/ops/group/shared/vec/conversions.cuh": "81bad9ad68323425d809dd6d1acf15614980a7eb8536deb211e19abaf0bd8d0d", + "third_party/thunderkittens-v4-prefill/include/ops/group/shared/vec/maps.cuh": "d8304875784cc1d3d5811e8128d783557750bee62e1301b9bbe38aa15c705065", + "third_party/thunderkittens-v4-prefill/include/ops/group/shared/vec/reductions.cuh": "d0639cb18a83a69d823003ed2f1b3152ad9aa5916c42b60b62fb6544aee53b95", + "third_party/thunderkittens-v4-prefill/include/ops/group/shared/vec/vec.cuh": "595dc8698d47cac6d1c4fabcf29ff80e395b308d6f59dabfe1c45a35dd60e4ac", + "third_party/thunderkittens-v4-prefill/include/ops/group/util/sync.cuh": "76e287ae4e52b0d50925ce61317ead2df69774211ee9b13af8ab5d9ae50b3d70", + "third_party/thunderkittens-v4-prefill/include/ops/group/util/tma.cuh": "331f5b697b81960f8a0efafefdd26234b9ef4252eadb9f435b4f72c9d4a0a2f1", + "third_party/thunderkittens-v4-prefill/include/ops/group/util/tma_cluster.cuh": "0827e3e9a79d890a4d95e31bcf72f8bd6daaf3a40380b3b30cb88c189c5bdf4c", + "third_party/thunderkittens-v4-prefill/include/ops/group/util/util.cuh": "9fd98c5b83daab9fb7fd84369e193961b0b6642ec4fe1a69dfa03ba26ee6e562", + "third_party/thunderkittens-v4-prefill/include/ops/ops.cuh": "df42b59e04cf73de846c4a440d2de677a34c862e52713643deb764a0d9c13608", + "third_party/thunderkittens-v4-prefill/include/ops/thread/memory/memory.cuh": "acdd06fd053038f88b6b142144d2b3596de0129151a4aea51cc061c63c0a1991", + "third_party/thunderkittens-v4-prefill/include/ops/thread/memory/tile/shared_to_tensor.cuh": "59feea31e66399111dd97becdbb3581cc3f6a4c8770c7bd1536db941272bcff1", + "third_party/thunderkittens-v4-prefill/include/ops/thread/memory/tile/tile.cuh": "1b8f761a88dad695784f69e795d47613c4347d064c63d8ec904b3a7db6b6505c", + "third_party/thunderkittens-v4-prefill/include/ops/thread/memory/tile/tma.cuh": "46a7ce8b187aad3fb745e8682013ab2dc86e79b790fc79db536811b8d6a6813f", + "third_party/thunderkittens-v4-prefill/include/ops/thread/memory/vec/tma.cuh": "e27158462b970aa82979edab403e827a57d211fe453b8988eaeeb8f88edefcab", + "third_party/thunderkittens-v4-prefill/include/ops/thread/memory/vec/vec.cuh": "d17cbcd1e753c4eeaabff904e8e8d47aaff41aacd8b3522b2ece077e54685913", + "third_party/thunderkittens-v4-prefill/include/ops/thread/mma/mma.cuh": "f6b00ee91d69e91776f672731a9945876a300dec4747b353401d5a6ff9741463", + "third_party/thunderkittens-v4-prefill/include/ops/thread/mma/tcgen05.cuh": "5e62c26bc99fd1ee7e1e501c15c6727fc4841c32f9cd1abeb134e15863ec5dab", + "third_party/thunderkittens-v4-prefill/include/ops/thread/thread.cuh": "e0cc0207cb876321dc8207a9b2fb91e4293c9cab0615a6235800270c44795c4e", + "third_party/thunderkittens-v4-prefill/include/ops/thread/util/sync.cuh": "b12f7d742679a6ac742ffc3a406768b59f7502ace643dfa9abd8a7c9d7041b08", + "third_party/thunderkittens-v4-prefill/include/ops/thread/util/tma.cuh": "a607438cc8ccd3407c0c90876eb0687095dc268375d910f7b823e59f11a670bf", + "third_party/thunderkittens-v4-prefill/include/ops/thread/util/util.cuh": "b02f88d1f2d5439ac5ea74450194d1909601e1c765b1b25c1990144a73685c1c", + "third_party/thunderkittens-v4-prefill/include/pyutils/broker.cuh": "94c4ed1fe61d037415045ab1eb68667d4b4ece77761383a4e99e142ee49dd398", + "third_party/thunderkittens-v4-prefill/include/pyutils/club.cuh": "0683c1d6f5cbcedc649ef82cb83d0872ab02cf31150e26b5c956aac2b566f7ce", + "third_party/thunderkittens-v4-prefill/include/pyutils/parallel_tensor.cuh": "0f9d245e9317d7449550aeb097d5302ebcfdfc650d41502311f48174385275d9", + "third_party/thunderkittens-v4-prefill/include/pyutils/profiler.cuh": "5878ea5db1a2ff531240bb8c8774aedc8bdbf4c437c95943b1b8baf35a0a8bdc", + "third_party/thunderkittens-v4-prefill/include/pyutils/pyutils.cuh": "785ff56292e5f6d5b82a92ec20049a3548d27c0542d3fc7b8a8a3a7350d8a51a", + "third_party/thunderkittens-v4-prefill/include/pyutils/torchutils.cuh": "6346609e35d7f3cafbaddea0166b5f733d16153346452a263632740ac8fcae29", + "third_party/thunderkittens-v4-prefill/include/types/global/cgl.cuh": "c7bad4a47bcaee615d744accbe2b35ca4d2981b85cc27c57be9ac45209a83466", + "third_party/thunderkittens-v4-prefill/include/types/global/gl.cuh": "0f68aad174bc75fdf758e54cecbd855a26046df2bcfdba59d2409e7977983681", + "third_party/thunderkittens-v4-prefill/include/types/global/global.cuh": "dfbbece449be8aa0a58ba99525d86335b097e5d20ce9d9750e3d3b35e0252d09", + "third_party/thunderkittens-v4-prefill/include/types/global/tma.cuh": "a0f2a5839c036918ced4caabd32e20fc1d3ce277ecaa8ef1458ed2bc043c2323", + "third_party/thunderkittens-v4-prefill/include/types/global/util.cuh": "941e4efe9376d73da53f6a6192e9f15a2b51ae574c47e599755e81f4a26eca00", + "third_party/thunderkittens-v4-prefill/include/types/register/crt.cuh": "e89d53f1fb3bbe16e21db6b6b9ca72437777078dafd63a29ad9dc11ee2a78dbc", + "third_party/thunderkittens-v4-prefill/include/types/register/crv.cuh": "2c6be7504402ac2afd092f2ce97693fe88e01e934729cdf3f07a016f87d58dcb", + "third_party/thunderkittens-v4-prefill/include/types/register/register.cuh": "ce6f326f9400a3020b679395c7adc3d7a676e3e91cfecac51a7f23b93cf35680", + "third_party/thunderkittens-v4-prefill/include/types/register/rt.cuh": "cbcc742c3f2164dc529b149682f45e44113060e60d5eb4c948e200a9da646470", + "third_party/thunderkittens-v4-prefill/include/types/register/rt_base.cuh": "a76083987ac029aaade27af96bb3e01cf54a773eb9a53e4eae2dabca24facdee", + "third_party/thunderkittens-v4-prefill/include/types/register/rt_layout.cuh": "0411f109ceb319d1fea7345976e71fcb7b9a1b35a0ae891d01b128eea3c47e44", + "third_party/thunderkittens-v4-prefill/include/types/register/rv.cuh": "84171a5201a53e0e8921ebc092fe6dcf5d3017c9c49cb8c48d5e3ee31ae21190", + "third_party/thunderkittens-v4-prefill/include/types/register/rv_layout.cuh": "1b472f2dfb093f14f85cf0e52bf5793b11507ff93fb1b73f4d82f26aa3f877a1", + "third_party/thunderkittens-v4-prefill/include/types/shared/cst.cuh": "eb4ce7286acbee11edd4d1e4b63f21bbc052b020db937ed83f1ca0763ddd2d7f", + "third_party/thunderkittens-v4-prefill/include/types/shared/csv.cuh": "1dae21a98b4201ddfaafe05d38d05a19aa324e2aaa5e6e821622c66609c8b5ee", + "third_party/thunderkittens-v4-prefill/include/types/shared/descriptor.cuh": "3865654a0c58cbf243313e02ee7287afdad2f4debe6e5cb40e707cedf883685f", + "third_party/thunderkittens-v4-prefill/include/types/shared/shared.cuh": "ce957ac1e9ee5c5d9b0f5bd8d083ab045ed2e939f455a3531138fb4eeb96c55d", + "third_party/thunderkittens-v4-prefill/include/types/shared/st.cuh": "3e5bde8898e440630a56059387bcdf6668965e8f22659d237ae14ddb8f1692f8", + "third_party/thunderkittens-v4-prefill/include/types/shared/sv.cuh": "00108b12b666c1e26269f5ceb7c2fdaf64dc4296ec4bdbaecd57ed7a1b4f56a7", + "third_party/thunderkittens-v4-prefill/include/types/system/ipc.cuh": "a0d7a08f4819511c58c7305f8fa759cf6cb4d960334353aca1e3d91342becd1b", + "third_party/thunderkittens-v4-prefill/include/types/system/pgl.cuh": "b16663ef25f1c08fa4a6cd6c212389dd6fe63c95bf5788935ac88b9269157be2", + "third_party/thunderkittens-v4-prefill/include/types/system/system.cuh": "434a007f4ac8b4baade8406b56f357db210ef237b0b96489530fc565afcce070", + "third_party/thunderkittens-v4-prefill/include/types/system/vmm.cuh": "4f73f2163bd4cbfea6c97ed7e549de760b56b53d1e0993f0a3c75b5347504f51", + "third_party/thunderkittens-v4-prefill/include/types/tensor/tensor.cuh": "f5685b3e636558eed83800e678e609a0e9c630a5e66ed93cf4775732b91dda9a", + "third_party/thunderkittens-v4-prefill/include/types/tensor/tt.cuh": "b052dfa1b4f13c4d10eac8fa14c6983b79b7a4f7fa514d5babbfab4a7240f133", + "third_party/thunderkittens-v4-prefill/include/types/types.cuh": "fe49f38a06f554c33b0e65ddb3af1f9e6026b652276898f6200b72b1b0626cf2", + "third_party/thunderkittens/LICENSE": "e705a1075640344d86b3b3d4f27dc06bf37cb5bf3b24676a9ca49a19c327ddae", + "third_party/thunderkittens/include/common/base_ops.cuh": "c33d4ba051c0a7bc27878809263d75f06167ed2e781ba781372d0d90c95373c5", + "third_party/thunderkittens/include/common/base_types.cuh": "6d09ee3d85bb952ae88fa9aeb79c9403717f4b473737f4a2e28055fe6fbac92f", + "third_party/thunderkittens/include/common/common.cuh": "486ded03bbe43ca4ed3fcce3988585c954494f8366ecbbfec088d215aaf5ca3e", + "third_party/thunderkittens/include/common/multimem.cuh": "a8aae2e97b5934f3849c9cb9ad5ef7994ca283eaff0c202d37c283636c0f3181", + "third_party/thunderkittens/include/common/util.cuh": "c715f83d149160261bfeed56c5d1ce4b4d1e225f156d4f3a532cc694655260ff", + "third_party/thunderkittens/include/kittens.cuh": "437d60a0d1c215019cc8f5e6b31c02043f1fa7c18f4bac4a6f95a94022f77715", + "third_party/thunderkittens/include/ops/group/group.cuh": "8c33856f083a99735807d82d48ee64192da4dfa6f6ea674267fec210620838ad", + "third_party/thunderkittens/include/ops/group/memory/memory.cuh": "d325ffeb850d7c2c942ac90e49e9cea37d4e565897682bfdd1270470c35733f6", + "third_party/thunderkittens/include/ops/group/memory/tile/complex_global_to_register.cuh": "cbf4b20658b85a721903f2a82faeb9aea826b30ef9546bc2087f71fb9843567a", + "third_party/thunderkittens/include/ops/group/memory/tile/complex_global_to_shared.cuh": "33a417e594becfc9e6284b304293cf5aba6a9d9e2fac8706442f466da0eca26d", + "third_party/thunderkittens/include/ops/group/memory/tile/complex_shared_to_register.cuh": "ffa32c513bf74ffad6f19f5aff8ca596e514e6a95e050babd1008207c3121c69", + "third_party/thunderkittens/include/ops/group/memory/tile/global_to_register.cuh": "04fcac6fdb32a785b304752e00f7a2681cf34695de85c13783c60e61457efd9a", + "third_party/thunderkittens/include/ops/group/memory/tile/global_to_shared.cuh": "266d97ee9a74b5c7c2706ecb9477e4b13b2ab52f8bc4324397a6c41a182282e0", + "third_party/thunderkittens/include/ops/group/memory/tile/parallel_global_to_global.cuh": "36e0923ba349ab9fefb277685d260f221c5949118c0dcd5d4daffcf46b11eeaa", + "third_party/thunderkittens/include/ops/group/memory/tile/shared_to_register.cuh": "ae047978ed23d2da0f690ae58b6e14c90c1aaa9b3d18fcda7b87f691bd4740e3", + "third_party/thunderkittens/include/ops/group/memory/tile/tensor_to_register.cuh": "f15ababf491e4b04f63d8e14f8cc60df669984853ce71cd1da57246f92d4f6c9", + "third_party/thunderkittens/include/ops/group/memory/tile/tile.cuh": "e0c270ec738a925281733ed9de450b518693d86652c018bba48baec767fe808f", + "third_party/thunderkittens/include/ops/group/memory/tile/tma.cuh": "2892ebb566b922c70f4c16cc7a5555dd329b55b55a5a4987d9e226c142b6b5fa", + "third_party/thunderkittens/include/ops/group/memory/tile/tma_cluster.cuh": "54bbdd60b28f7ad1cbec58867b3e3a0291f1aab3cc386a1938eaf969c0c19a1f", + "third_party/thunderkittens/include/ops/group/memory/vec/global_to_register.cuh": "f784f3e7f43ff4ae299fa660a91ba9b0690538334ca7fa2a7386c2fe0de50fe7", + "third_party/thunderkittens/include/ops/group/memory/vec/global_to_shared.cuh": "b9e0483deeb85e0887f92443d8f886148c2cdf491b44792f13056e94656f60e1", + "third_party/thunderkittens/include/ops/group/memory/vec/shared_to_register.cuh": "55f3546b89e00f8bc84f7805cd150e30cead7d2933e7992cb9ec720974d204e4", + "third_party/thunderkittens/include/ops/group/memory/vec/tma.cuh": "4b93ba91d05f7defae9ea8c66da29ad687d0780ae148d96271f14d60cadc3fbf", + "third_party/thunderkittens/include/ops/group/memory/vec/tma_cluster.cuh": "00ac7bdfc5b5c0ca38b037f4e262b19fa73f4164921a56c303a33e49349e1743", + "third_party/thunderkittens/include/ops/group/memory/vec/vec.cuh": "b7ee8a98c36dc495f2201f91ac9a3be5af88bf5f07a14961b84a304f655d7c3b", + "third_party/thunderkittens/include/ops/group/mma/base/64x112.impl": "a04d5742a0ea37c6a6ca9379f4912dc8ef7424372342add231867d98e96df232", + "third_party/thunderkittens/include/ops/group/mma/base/64x128.impl": "2508e03800c05c778f345f8c5fe96415b6fcae8075fa7110ffe5372d3497ade1", + "third_party/thunderkittens/include/ops/group/mma/base/64x144.impl": "7bcf837979c41e7aca984b94bff07fd13b22eb7a59fe2e48855cc07f2d4a9dbd", + "third_party/thunderkittens/include/ops/group/mma/base/64x16.impl": "25c681bdbb3d3fd4d5bc6fbc9c6070583ff08d593538248e0a0d8f58b18e1042", + "third_party/thunderkittens/include/ops/group/mma/base/64x160.impl": "56378c1989a098599697193a7bc163617cc4488de13a96b76fa7ef9daf4d8472", + "third_party/thunderkittens/include/ops/group/mma/base/64x176.impl": "72e9e4865b5fcf92d67fb37e4e3114c1493f4cdaa8738772021048e1ffc8845b", + "third_party/thunderkittens/include/ops/group/mma/base/64x192.impl": "91f2c3d0ed837be3706daa603ec445d5a6ad3c20112ccd1e8420b05e9a181391", + "third_party/thunderkittens/include/ops/group/mma/base/64x208.impl": "4676ec206611e28b447cae7f0da5e9d032b2e834ab37f0663dce4248f3933128", + "third_party/thunderkittens/include/ops/group/mma/base/64x224.impl": "2bccdb907ededc0ce9ed55dbc4b6c95801aa2f32c9c7edff0742c492f8df5a45", + "third_party/thunderkittens/include/ops/group/mma/base/64x240.impl": "2ef3ef8c0cf374f97962731777a4352f241ec7d964f6f54f3b45e719d16efd7b", + "third_party/thunderkittens/include/ops/group/mma/base/64x256.impl": "0c7515f4b02cf49ab10e4ece84e97cf3b822e5fed93287cde77db5f708c2dee1", + "third_party/thunderkittens/include/ops/group/mma/base/64x32.impl": "bbedf1bc2e24bc9f152c5f9ce029bb178a40c9f7c510077bccd3e5f925361f8e", + "third_party/thunderkittens/include/ops/group/mma/base/64x48.impl": "bf6fab4e0611376788664f47c5e295bcd58e09e3faf42b47a6531bae61cdd9ba", + "third_party/thunderkittens/include/ops/group/mma/base/64x64.impl": "fb5a176e96e579f834750e9d86f80a399b1aa0dd34a4991ff639da0eab1f6e9e", + "third_party/thunderkittens/include/ops/group/mma/base/64x80.impl": "15b0c22259f07fc17fe70bcb2341e5b673ecef55a37dc7aceca1299ed7022888", + "third_party/thunderkittens/include/ops/group/mma/base/64x96.impl": "948e526120979fa31c88f2232d5a433e7c8a0af7168ea1be01738f8d46f8d996", + "third_party/thunderkittens/include/ops/group/mma/base/base.cuh": "0d97d610d412ee7580d199c3180c94a6f467bbe1273f09b9a4dafd641f9353f9", + "third_party/thunderkittens/include/ops/group/mma/mma.cuh": "015245bd39b75b54cb999328ab887c1d0d04541b6570cc9af24dd716d0a8d8f4", + "third_party/thunderkittens/include/ops/group/mma/tcgen05.cuh": "a2cccdfbc9a9843ce10dc04d2a0811dbef9f65d4065433191b5384c27399898b", + "third_party/thunderkittens/include/ops/group/mma/warp.cuh": "426346462672db5bee203a82f1fc98100bb580085db2d5b61868a5131ba854a4", + "third_party/thunderkittens/include/ops/group/mma/warpgroup.cuh": "574a3227339eae9c4b88b2e2ad420ce47b7cf9eab08d49422984b93210cdf3b2", + "third_party/thunderkittens/include/ops/group/register/register.cuh": "d325ffeb850d7c2c942ac90e49e9cea37d4e565897682bfdd1270470c35733f6", + "third_party/thunderkittens/include/ops/group/register/tile/complex_conversions.cuh": "6072843c18334a812d786ba021564986b07568369858b1c6c1c8926479edaa40", + "third_party/thunderkittens/include/ops/group/register/tile/complex_maps.cuh": "44ecc28580def5b189df2265b72fb66169ad1fb959af241f2215dc741ad0301b", + "third_party/thunderkittens/include/ops/group/register/tile/conversions.cuh": "cfc15594b75ffc80383ddf49fe565a99855629b07e969502f5da0c059597f9be", + "third_party/thunderkittens/include/ops/group/register/tile/maps.cuh": "42bda47c15e10c35176ba3868924b00255f033973dd8f07848b7627cea413c57", + "third_party/thunderkittens/include/ops/group/register/tile/reductions.cuh": "e8914eb995b169aa72104d2d6505acd27546812ded7089723404f5c99476c87a", + "third_party/thunderkittens/include/ops/group/register/tile/tile.cuh": "7eae6e436a2dd01cf1bc3dddce26349b4a7cfca509a3f3b841b4b6dcf4b14400", + "third_party/thunderkittens/include/ops/group/register/vec/conversions.cuh": "e34b95719637bf0449781afa2a63b7d1066080af551515e05d70fcc72e50f130", + "third_party/thunderkittens/include/ops/group/register/vec/maps.cuh": "e8caa7feffdf8324bb0a4b6df3a44af44f291bb26aace54c27c8900deb45f486", + "third_party/thunderkittens/include/ops/group/register/vec/reductions.cuh": "395ee946b9288b2add58272851d3031bccd8aec5c8d9de452b42df61ff3f8c2d", + "third_party/thunderkittens/include/ops/group/register/vec/vec.cuh": "c9325ef900273918f92d68f6853c92a51c2c2897ca1d69a8a8cea9892292563a", + "third_party/thunderkittens/include/ops/group/shared/shared.cuh": "d325ffeb850d7c2c942ac90e49e9cea37d4e565897682bfdd1270470c35733f6", + "third_party/thunderkittens/include/ops/group/shared/tile/conversions.cuh": "1464ef684ec8ed7ecb2459551c5968593b56177e60cad7ed1c6f137b6cd95916", + "third_party/thunderkittens/include/ops/group/shared/tile/maps.cuh": "a35a5659cf24e8f11559e26d5fe4d969925f25e9a86b6c8f7087f0523b09adb9", + "third_party/thunderkittens/include/ops/group/shared/tile/reductions.cuh": "fcb9ea4a52056a17619c71a41041ffff2175c3b13308afb92d8fff3ef8cb162e", + "third_party/thunderkittens/include/ops/group/shared/tile/tile.cuh": "2c45011828ef5bf5622e390a56addd502e0f503baf7c8cd07311b23b2ec7a634", + "third_party/thunderkittens/include/ops/group/shared/vec/conversions.cuh": "81bad9ad68323425d809dd6d1acf15614980a7eb8536deb211e19abaf0bd8d0d", + "third_party/thunderkittens/include/ops/group/shared/vec/maps.cuh": "d8304875784cc1d3d5811e8128d783557750bee62e1301b9bbe38aa15c705065", + "third_party/thunderkittens/include/ops/group/shared/vec/reductions.cuh": "d0639cb18a83a69d823003ed2f1b3152ad9aa5916c42b60b62fb6544aee53b95", + "third_party/thunderkittens/include/ops/group/shared/vec/vec.cuh": "595dc8698d47cac6d1c4fabcf29ff80e395b308d6f59dabfe1c45a35dd60e4ac", + "third_party/thunderkittens/include/ops/group/util/sync.cuh": "76e287ae4e52b0d50925ce61317ead2df69774211ee9b13af8ab5d9ae50b3d70", + "third_party/thunderkittens/include/ops/group/util/tma.cuh": "331f5b697b81960f8a0efafefdd26234b9ef4252eadb9f435b4f72c9d4a0a2f1", + "third_party/thunderkittens/include/ops/group/util/tma_cluster.cuh": "0827e3e9a79d890a4d95e31bcf72f8bd6daaf3a40380b3b30cb88c189c5bdf4c", + "third_party/thunderkittens/include/ops/group/util/util.cuh": "9fd98c5b83daab9fb7fd84369e193961b0b6642ec4fe1a69dfa03ba26ee6e562", + "third_party/thunderkittens/include/ops/ops.cuh": "df42b59e04cf73de846c4a440d2de677a34c862e52713643deb764a0d9c13608", + "third_party/thunderkittens/include/ops/thread/memory/memory.cuh": "acdd06fd053038f88b6b142144d2b3596de0129151a4aea51cc061c63c0a1991", + "third_party/thunderkittens/include/ops/thread/memory/tile/shared_to_tensor.cuh": "59feea31e66399111dd97becdbb3581cc3f6a4c8770c7bd1536db941272bcff1", + "third_party/thunderkittens/include/ops/thread/memory/tile/tile.cuh": "1b8f761a88dad695784f69e795d47613c4347d064c63d8ec904b3a7db6b6505c", + "third_party/thunderkittens/include/ops/thread/memory/tile/tma.cuh": "46a7ce8b187aad3fb745e8682013ab2dc86e79b790fc79db536811b8d6a6813f", + "third_party/thunderkittens/include/ops/thread/memory/vec/tma.cuh": "e27158462b970aa82979edab403e827a57d211fe453b8988eaeeb8f88edefcab", + "third_party/thunderkittens/include/ops/thread/memory/vec/vec.cuh": "d17cbcd1e753c4eeaabff904e8e8d47aaff41aacd8b3522b2ece077e54685913", + "third_party/thunderkittens/include/ops/thread/mma/mma.cuh": "f6b00ee91d69e91776f672731a9945876a300dec4747b353401d5a6ff9741463", + "third_party/thunderkittens/include/ops/thread/mma/tcgen05.cuh": "5e62c26bc99fd1ee7e1e501c15c6727fc4841c32f9cd1abeb134e15863ec5dab", + "third_party/thunderkittens/include/ops/thread/thread.cuh": "e0cc0207cb876321dc8207a9b2fb91e4293c9cab0615a6235800270c44795c4e", + "third_party/thunderkittens/include/ops/thread/util/sync.cuh": "b12f7d742679a6ac742ffc3a406768b59f7502ace643dfa9abd8a7c9d7041b08", + "third_party/thunderkittens/include/ops/thread/util/tma.cuh": "a607438cc8ccd3407c0c90876eb0687095dc268375d910f7b823e59f11a670bf", + "third_party/thunderkittens/include/ops/thread/util/util.cuh": "b02f88d1f2d5439ac5ea74450194d1909601e1c765b1b25c1990144a73685c1c", + "third_party/thunderkittens/include/pyutils/broker.cuh": "94c4ed1fe61d037415045ab1eb68667d4b4ece77761383a4e99e142ee49dd398", + "third_party/thunderkittens/include/pyutils/club.cuh": "0683c1d6f5cbcedc649ef82cb83d0872ab02cf31150e26b5c956aac2b566f7ce", + "third_party/thunderkittens/include/pyutils/parallel_tensor.cuh": "0f9d245e9317d7449550aeb097d5302ebcfdfc650d41502311f48174385275d9", + "third_party/thunderkittens/include/pyutils/profiler.cuh": "5878ea5db1a2ff531240bb8c8774aedc8bdbf4c437c95943b1b8baf35a0a8bdc", + "third_party/thunderkittens/include/pyutils/pyutils.cuh": "785ff56292e5f6d5b82a92ec20049a3548d27c0542d3fc7b8a8a3a7350d8a51a", + "third_party/thunderkittens/include/pyutils/torchutils.cuh": "6346609e35d7f3cafbaddea0166b5f733d16153346452a263632740ac8fcae29", + "third_party/thunderkittens/include/types/global/cgl.cuh": "c7bad4a47bcaee615d744accbe2b35ca4d2981b85cc27c57be9ac45209a83466", + "third_party/thunderkittens/include/types/global/gl.cuh": "0f68aad174bc75fdf758e54cecbd855a26046df2bcfdba59d2409e7977983681", + "third_party/thunderkittens/include/types/global/global.cuh": "dfbbece449be8aa0a58ba99525d86335b097e5d20ce9d9750e3d3b35e0252d09", + "third_party/thunderkittens/include/types/global/tma.cuh": "a0f2a5839c036918ced4caabd32e20fc1d3ce277ecaa8ef1458ed2bc043c2323", + "third_party/thunderkittens/include/types/global/util.cuh": "941e4efe9376d73da53f6a6192e9f15a2b51ae574c47e599755e81f4a26eca00", + "third_party/thunderkittens/include/types/register/crt.cuh": "e89d53f1fb3bbe16e21db6b6b9ca72437777078dafd63a29ad9dc11ee2a78dbc", + "third_party/thunderkittens/include/types/register/crv.cuh": "2c6be7504402ac2afd092f2ce97693fe88e01e934729cdf3f07a016f87d58dcb", + "third_party/thunderkittens/include/types/register/register.cuh": "ce6f326f9400a3020b679395c7adc3d7a676e3e91cfecac51a7f23b93cf35680", + "third_party/thunderkittens/include/types/register/rt.cuh": "cbcc742c3f2164dc529b149682f45e44113060e60d5eb4c948e200a9da646470", + "third_party/thunderkittens/include/types/register/rt_base.cuh": "a76083987ac029aaade27af96bb3e01cf54a773eb9a53e4eae2dabca24facdee", + "third_party/thunderkittens/include/types/register/rt_layout.cuh": "0411f109ceb319d1fea7345976e71fcb7b9a1b35a0ae891d01b128eea3c47e44", + "third_party/thunderkittens/include/types/register/rv.cuh": "84171a5201a53e0e8921ebc092fe6dcf5d3017c9c49cb8c48d5e3ee31ae21190", + "third_party/thunderkittens/include/types/register/rv_layout.cuh": "1b472f2dfb093f14f85cf0e52bf5793b11507ff93fb1b73f4d82f26aa3f877a1", + "third_party/thunderkittens/include/types/shared/cst.cuh": "eb4ce7286acbee11edd4d1e4b63f21bbc052b020db937ed83f1ca0763ddd2d7f", + "third_party/thunderkittens/include/types/shared/csv.cuh": "1dae21a98b4201ddfaafe05d38d05a19aa324e2aaa5e6e821622c66609c8b5ee", + "third_party/thunderkittens/include/types/shared/descriptor.cuh": "3865654a0c58cbf243313e02ee7287afdad2f4debe6e5cb40e707cedf883685f", + "third_party/thunderkittens/include/types/shared/shared.cuh": "ce957ac1e9ee5c5d9b0f5bd8d083ab045ed2e939f455a3531138fb4eeb96c55d", + "third_party/thunderkittens/include/types/shared/st.cuh": "3e5bde8898e440630a56059387bcdf6668965e8f22659d237ae14ddb8f1692f8", + "third_party/thunderkittens/include/types/shared/sv.cuh": "00108b12b666c1e26269f5ceb7c2fdaf64dc4296ec4bdbaecd57ed7a1b4f56a7", + "third_party/thunderkittens/include/types/system/ipc.cuh": "a0d7a08f4819511c58c7305f8fa759cf6cb4d960334353aca1e3d91342becd1b", + "third_party/thunderkittens/include/types/system/pgl.cuh": "b16663ef25f1c08fa4a6cd6c212389dd6fe63c95bf5788935ac88b9269157be2", + "third_party/thunderkittens/include/types/system/system.cuh": "434a007f4ac8b4baade8406b56f357db210ef237b0b96489530fc565afcce070", + "third_party/thunderkittens/include/types/system/vmm.cuh": "4f73f2163bd4cbfea6c97ed7e549de760b56b53d1e0993f0a3c75b5347504f51", + "third_party/thunderkittens/include/types/tensor/tensor.cuh": "f5685b3e636558eed83800e678e609a0e9c630a5e66ed93cf4775732b91dda9a", + "third_party/thunderkittens/include/types/tensor/tt.cuh": "b052dfa1b4f13c4d10eac8fa14c6983b79b7a4f7fa514d5babbfab4a7240f133", + "third_party/thunderkittens/include/types/types.cuh": "fe49f38a06f554c33b0e65ddb3af1f9e6026b652276898f6200b72b1b0626cf2", + "third_party/thunderkittens/prototype/common/common.cuh": "be060a035c62798e744bb9458facd2831b4d9b6cb86f7ea6a8cbcc9170a8e0af", + "third_party/thunderkittens/prototype/common/templates.cuh": "2ab8e744f1dca706d7fcc20e2b256267cff7bc0391d5838fc7a4f8ea09dcfcc4", + "third_party/thunderkittens/prototype/common/util.cuh": "6edf07303b948f1d686f772cb9042a4b0ff353e7063dd52f35db00ab59066b20", + "third_party/thunderkittens/prototype/interpreter/interpreter.cuh": "6292e108902bb321c1d94438f194166c14eb45ba977c365a03525020d4a27101", + "third_party/thunderkittens/prototype/interpreter/templates.cuh": "eb208b256d0152748b8cce2a75558f41321c3fec34ab6da6b1ce35f5083e602d", + "third_party/thunderkittens/prototype/lcf/lcf.cuh": "87fc236c57f6c3dfaceefb0b4faf8090c4ac7ef17eeb16268ce8d5cbf6b60206", + "third_party/thunderkittens/prototype/lcf/templates.cuh": "c212337f72fb110b54788a80ba81ff0bf016918fd94f192963650919661fb374", + "third_party/thunderkittens/prototype/lcsc/lcsc.cuh": "e49ca81f6869ff8e6402bf33fae78471c0fddc90c6005c2ec1dd668414beac33", + "third_party/thunderkittens/prototype/lcsf/lcsf.cuh": "4b2d4b9f850ad98637942874fe3b316b438be2821888a17464326cf8de0ee9a8", + "third_party/thunderkittens/prototype/lcsf/templates.cuh": "8658ca4356bcfe5c10c7073d975c6e682f40a9694381b24d8e016f1148489829", + "third_party/thunderkittens/prototype/prototype.cuh": "2dca0d01f438d5cb8f5509cbe7824e8e84b8e14ccec092e292344002a368c65b" + }, + "third_party": { + "thunderkittens": { + "repository": "https://github.com/HazyResearch/ThunderKittens", + "revision_kind": "content-addressed-source-snapshot", + "license": "third_party/thunderkittens/LICENSE", + "source_tree_sha256": "b351a87ad998da66d814ea5922829761dff89d2dd62db053f9fd58e834f2ccec" + }, + "thunderkittens-v4-prefill": { + "repository": "https://github.com/HazyResearch/ThunderKittens", + "revision_kind": "content-addressed-source-snapshot", + "license": "third_party/thunderkittens-v4-prefill/LICENSE", + "source_tree_sha256": "21177be331191de12f37900819f59b34e9321637612077ee344d96ab8011f34a" + } + } +} diff --git a/deploy/nemotron_native/kernels/stage0/b1/bb_router_selection.cuh b/deploy/nemotron_native/kernels/stage0/b1/bb_router_selection.cuh new file mode 100644 index 0000000..d7aa000 --- /dev/null +++ b/deploy/nemotron_native/kernels/stage0/b1/bb_router_selection.cuh @@ -0,0 +1,39 @@ +#pragma once +__device__ __forceinline__ bool bb_router_select_parallel( + const float *score, const bf16 *gate_bias, float *out) { + static_assert(BB_EXPERTS == 24 && BB_TOP_K == 4); + constexpr unsigned full = 0xffffffffu; + const int lane = threadIdx.x & 31; + const float unbiased = lane < BB_EXPERTS ? score[lane] : 0.f; + float remaining = lane < BB_EXPERTS + ? unbiased + __bfloat162float(gate_bias[lane]) : -CUDART_INF_F; + const bool finite = lane >= BB_EXPERTS || + (((__float_as_uint(unbiased) & 0x7f800000u) != 0x7f800000u) && + ((__float_as_uint(remaining) & 0x7f800000u) != 0x7f800000u)); + if (!__all_sync(full, finite)) return false; + bool selected = false; + float total = 0.f; + #pragma unroll + for (int k = 0; k < BB_TOP_K; ++k) { + float best_score = remaining; + int best_id = lane < BB_EXPERTS ? lane : 0x7fffffff; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + const float other_score = __shfl_xor_sync(full, best_score, off); + const int other_id = __shfl_xor_sync(full, best_id, off); + if (other_score > best_score || + (other_score == best_score && other_id < best_id)) { + best_score = other_score; + best_id = other_id; + } + } + total += __shfl_sync(full, unbiased, best_id); + if (lane == best_id) { + selected = true; + remaining = -CUDART_INF_F; + } + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + if (lane < BB_EXPERTS) out[lane] = selected ? unbiased * norm : 0.f; + return true; +} diff --git a/deploy/nemotron_native/kernels/stage0/b1/mk2_abi.cuh b/deploy/nemotron_native/kernels/stage0/b1/mk2_abi.cuh new file mode 100644 index 0000000..bf1e415 --- /dev/null +++ b/deploy/nemotron_native/kernels/stage0/b1/mk2_abi.cuh @@ -0,0 +1,85 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_set_native_predictor_active_mask(MK2Predictor *p, const unsigned char *mask); +int mk2_native_startup_skip_abi_version(); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *, int grid); diff --git a/deploy/nemotron_native/kernels/stage0/b1/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/stage0/b1/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/stage0/b1/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/stage0/b1/native_prefill_impl.inc b/deploy/nemotron_native/kernels/stage0/b1/native_prefill_impl.inc new file mode 100644 index 0000000..a254308 --- /dev/null +++ b/deploy/nemotron_native/kernels/stage0/b1/native_prefill_impl.inc @@ -0,0 +1,198 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 16 || BATCH == 32 || BATCH == 64)) + throw std::runtime_error("native prefill requires matching batch 1/16/32/64"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (padded > INT_MAX / ((BB_EXPERTS + 1) * D_MODEL)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + float *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(float))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(float))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} diff --git a/deploy/nemotron_native/kernels/stage0/b1/predictor_mk2.cu b/deploy/nemotron_native/kernels/stage0/b1/predictor_mk2.cu new file mode 100644 index 0000000..79a468c --- /dev/null +++ b/deploy/nemotron_native/kernels/stage0/b1/predictor_mk2.cu @@ -0,0 +1,3910 @@ +#include "../../inflight/progress.cuh" +#if MK2_INFLIGHT +#include +#endif +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef MK2_PARALLEL_ROUTER +#define MK2_PARALLEL_ROUTER 0 +#endif +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +constexpr int BATCH_PAD = ((BATCH + 15) / 16) * 16; +static_assert(BATCH >= 1 && BATCH <= BATCH_PAD); +static_assert(BATCH % 16 != 0 || BATCH_PAD == BATCH); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +constexpr int THREADS = 256; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_ACC_F32 }; +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + volatile unsigned int *generation = state + 1; + const unsigned int seen = *generation; + if (atomicAdd(state, 1u) == gridDim.x - 1) { + atomicExch(state, 0u); + atomicAdd(state + 1, 1u); + } else { + while (*generation == seen) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / THREADS; + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const bf16 *in = src + (size_t)row * D_MODEL; + float v[PER_THREAD]; + float sum = 0.f, sq = 0.f; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + v[j] = __bfloat162float(in[threadIdx.x + j * THREADS]); + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) { scratch[warp] = sum; scratch[WARPS + warp] = sq; } + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f, q = 0.f; + for (int w = 0; w < WARPS; ++w) { s += scratch[w]; q += scratch[WARPS + w]; } + scratch[2 * WARPS] = s / D_MODEL; + scratch[2 * WARPS + 1] = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + } + __syncthreads(); + const float mean = scratch[2 * WARPS], inv = scratch[2 * WARPS + 1]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * weight[i]); + } + __syncthreads(); + } +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH) { + constexpr int PER_THREAD = D_MODEL / THREADS; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const size_t base = (size_t)row * D_MODEL; + float v[PER_THREAD]; + float sq = 0.f; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + v[j] = __bfloat162float(h[base + i]) + __bfloat162float(r[base + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) scratch[warp] = sq; + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f; + for (int w = 0; w < WARPS; ++w) s += scratch[w]; + scratch[WARPS] = rsqrtf(s / D_MODEL + RMS_EPS); + } + __syncthreads(); + const float inv = scratch[WARPS]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + out[base + i] = __float2bfloat16( + __bfloat162float(weight[i]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + } + __syncthreads(); + } +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map MK2_PROGRESS_PARAMETER(progress_mask)) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + MK2_PROGRESS_GUARD(progress_mask, row); + const int c = idx - row * MB_CONV_DIM; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + const float raw = __bfloat162float( + zxbcdt[(size_t)row * MB_IN_PROJ + MB_INNER + c]); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + xbc_act[(size_t)row * MB_CONV_DIM + c] = __float2bfloat16(acc); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map MK2_PROGRESS_PARAMETER(progress_mask)) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + const int pairs = BATCH * MB_HEADS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_HEADS * split); + MK2_PROGRESS_GUARD(progress_mask, row); + const int head = (pair / split) % MB_HEADS; + const int first_channel = (pair % split) * channels; + const int group = head / MB_HEADS_PER_GROUP; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + L.dt_bias[head]); + const float dA = __expf(L.A[head] * dt); + const float Dh = L.D[head]; + float *st = L.ssm_state + (size_t)slot_map[row] * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + float prior[channels][MB_STATE / 32]; + float input[channels]; + #pragma unroll + for (int c = 0; c < channels; ++c) { + const int p = first_channel + c; + input[c] = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) + prior[c][j] = st[(size_t)p * MB_STATE + lane + j * 32]; + } + #pragma unroll + for (int c = 0; c < channels; ++c) { + const int p = first_channel + c; + const float xv = input[c]; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + const size_t at = (size_t)p * MB_STATE + n; + const float s = prior[c][j] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_CONV_DIM; + const int c = idx - t * MB_CONV_DIM; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v; +#ifdef PF_CONV_ZERO_STATE + if (src < 0) v = 0.f; +#else + if (src < 0) v = __bfloat162float(state[pos + k]); +#endif + else v = __bfloat162float( + zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]); + acc += __bfloat162float(w[k]) * v; + } +#ifndef PF_CONV_NO_SILU + acc = acc / (1.f + __expf(-acc)); +#endif + xbc_act[(size_t)t * MB_CONV_DIM + c] = __float2bfloat16(acc); + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + const int units = n_seq * MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int seq = unit / (MB_HEADS * MB_HEAD_DIM); + const int rem = unit - seq * (MB_HEADS * MB_HEAD_DIM); + const int head = rem / MB_HEAD_DIM; + const int p = rem - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p * MB_STATE; + float sreg[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) sreg[j] = st[lane + j * 32]; + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const bf16 *xa = xbc_act + t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + const bf16 *dtrow = zxbcdt + t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + sreg[j] = sreg[j] * dA + __bfloat162float(Bsec[n]) * dt * xv; + acc += sreg[j] * __bfloat162float(Csec[n]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) st[lane + j * 32] = sreg[j]; + } + return; + } +#if PF_SSM_SPLIT_P + const int units = MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; unit += warps_total) { + const int head = unit / MB_HEAD_DIM; + const int p_only = unit - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; +#else + for (int head = blockIdx.x * WARPS + warp; head < MB_HEADS; + head += warps_total) { + const int group = head / MB_HEADS_PER_GROUP; +#endif + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *st = L.ssm_state + (size_t)slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE +#if PF_SSM_SPLIT_P + + (size_t)p_only * MB_STATE +#endif + ; + for (int t = 0; t < n_tokens; ++t) { + const bf16 *xa = xbc_act + (size_t)t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = + zxbcdt + (size_t)t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } +#if PF_SSM_SPLIT_P + { + const int p = p_only; +#else + for (int p = 0; p < MB_HEAD_DIM; ++p) { +#endif + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { +#if PF_SSM_SPLIT_P + const size_t at = (size_t)lane + j * 32; +#else + const size_t at = (size_t)p * MB_STATE + lane + j * 32; +#endif + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } + } +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16( + held[j] * inv * __bfloat162float(L.gnorm_w[idx])); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len MK2_PROGRESS_PARAMETER(progress_mask)) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = BATCH * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / BB_ATT_HEADS; + MK2_PROGRESS_GUARD(progress_mask, row); + const int head = pair % BB_ATT_HEADS; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const int kv = head / kv_per_q; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *q = base + head * BB_HEAD_DIM; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (head % kv_per_q == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + for (int tpos = 0; tpos < T; ++tpos) { +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new + : kv_slot(L, row, 0, tpos) + kv_offset(L, tpos, kv); +#endif + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m, s); + const float rescale = __expf(m - m_new); + const float w = __expf(s - m_new); + denom = denom * rescale + w; +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *vt = kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#else + const bf16 *vt = cur_token + ? v_new + : kv_slot(L, row, 1, tpos) + kv_offset(L, tpos, kv); +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + w * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * BB_KV_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / BB_KV_HEADS; + const int kv = idx - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv); + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv); + for (int i = 0; i < BB_HEAD_DIM; ++i) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + for (int s = 0; s <= pos; ++s) { + const bf16 *kt = kv_slot(L, seq, 0, s) + kv_offset(L, s, kv); + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + const bf16 *vt = kv_slot(L, seq, 1, s) + kv_offset(L, s, kv); + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + wt * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +#include "bb_router_selection.cuh" +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH) { + __shared__ float score[BB_EXPERTS]; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + for (int e = warp; e < BB_EXPERTS; e += WARPS) { + const bf16 *wr = gate_w + (size_t)e * D_MODEL; + float acc = 0.f; + for (int i = lane; i < D_MODEL; i += 32) + acc += __bfloat162float(normed[(size_t)row * D_MODEL + i]) * + __bfloat162float(wr[i]); + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) score[e] = 1.f / (1.f + __expf(-acc)); + } + __syncthreads(); + bool parallel_done = false; +#if MK2_PARALLEL_ROUTER && MK2_BATCH == 1 + if (n_rows == 1 && threadIdx.x < 32) + parallel_done = bb_router_select_parallel( + score, gate_bias, moe_w + row * BB_EXPERTS); +#endif + if (threadIdx.x == 0 && !parallel_done) { + float remaining[BB_EXPERTS]; + int pick[BB_TOP_K]; + for (int e = 0; e < BB_EXPERTS; ++e) { + moe_w[row * BB_EXPERTS + e] = 0.f; + remaining[e] = score[e] + __bfloat162float(gate_bias[e]); + } + float total = 0.f; + for (int k = 0; k < BB_TOP_K; ++k) { + int best = 0; + for (int e = 1; e < BB_EXPERTS; ++e) + if (remaining[e] > remaining[best]) best = e; + pick[k] = best; + total += score[best]; + remaining[best] = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + for (int k = 0; k < BB_TOP_K; ++k) + moe_w[row * BB_EXPERTS + pick[k]] = score[pick[k]] * norm; + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert <= BB_EXPERTS; ++expert) + value = __fadd_rn(value, parts[base + (size_t)expert * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +#ifndef MK2_PARALLEL_SAMPLE +#define MK2_PARALLEL_SAMPLE 0 +#endif +#ifndef MK2_PARALLEL_PHONEME +#define MK2_PARALLEL_PHONEME 0 +#endif +__device__ __forceinline__ void mk2_argmax_merge(float &v, int &rank, float other, int other_rank) { + if (other > v || (other == v && other_rank < rank)) { + v = other; rank = other_rank; + } +} +__device__ __forceinline__ int mk2_block_argmax(float value, int rank) { + __shared__ float warp_value[WARPS]; + __shared__ int warp_rank[WARPS]; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffff, value, off); + const int other_rank = __shfl_down_sync(0xffffffff, rank, off); + mk2_argmax_merge(value, rank, other, other_rank); + } + if (lane == 0) { warp_value[warp] = value; warp_rank[warp] = rank; } + __syncthreads(); + if (warp == 0) { + value = lane < WARPS ? warp_value[lane] : -CUDART_INF_F; + rank = lane < WARPS ? warp_rank[lane] : INT_MAX; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffff, value, off); + const int other_rank = __shfl_down_sync(0xffffffff, rank, off); + mk2_argmax_merge(value, rank, other, other_rank); + } + } + return rank; +} +__device__ __forceinline__ int mk2_eos_insertion_rank(const uint64_t *keys, float eos) { + int lo = 0, hi = TOP_K; + while (lo < hi) { + const int mid = (lo + hi) >> 1; + if (eos > ordered_to_float((uint32_t)(keys[mid] >> 32))) hi = mid; + else lo = mid + 1; + } + return lo; +} +__device__ __forceinline__ void mk2_sample_topk_serial( + const uint64_t *keys, float eos, const float *noise, int *out) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float v; + int n; + if (!eos_used && eos > av) { + v = eos; + n = AUDIO_EOS; + eos_used = true; + } else { + v = av; + n = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = v + noise[j]; + if (noised > best_noised) { best_noised = noised; best_code = n; } + } + *out = best_code; +} +__device__ __forceinline__ void mk2_sample_topk_parallel( + const uint64_t *keys, float eos, const float *noise, int *out) { + const int j = threadIdx.x; + const bool sorted_nan = j < TOP_K && isnan(ordered_to_float((uint32_t)(keys[j] >> 32))); + if (__syncthreads_or(sorted_nan)) { + if (threadIdx.x == 0) mk2_sample_topk_serial(keys, eos, noise, out); + return; + } + const int lane = threadIdx.x & 31; + int insertion = lane == 0 ? mk2_eos_insertion_rank(keys, eos) : 0; + insertion = __shfl_sync(0xffffffff, insertion, 0); + float value = -CUDART_INF_F; + int rank = INT_MAX; + if (j < TOP_K) { + const int a = j - (j > insertion); + const float v = j == insertion ? eos : ordered_to_float((uint32_t)(keys[a] >> 32)); + const float noised = v + noise[j]; + if (noised > -CUDART_INF_F) { value = noised; rank = j; } + } + const int winner = mk2_block_argmax(value, rank); + if (threadIdx.x == 0) { + if (winner == INT_MAX) *out = 0; + else if (winner == insertion) *out = AUDIO_EOS; + else { + const int a = winner - (winner > insertion); + *out = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + } + } +} +__device__ __forceinline__ void mk2_phoneme_argmax_parallel(const bf16 *logits, int *out) { + float best = -CUDART_INF_F; + int rank = INT_MAX; + for (int n = threadIdx.x; n < BB_PHONEME; n += THREADS) { + const float v = __bfloat162float(logits[n]); + if (v > best) { best = v; rank = n; } + } + const int winner = mk2_block_argmax(best, rank); + if (threadIdx.x == 0) *out = winner == INT_MAX ? 0 : winner; +} +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + uint64_t *keys, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int n = threadIdx.x; n < SORT_N; n += THREADS) { + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + keys[n] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + __syncthreads(); + for (int k = 2; k <= SORT_N; k <<= 1) { + if (k >= 64) __syncthreads(); + for (int j = k >> 1; j > 0; j >>= 1) { + for (int i = threadIdx.x; i < SORT_N; i += THREADS) { + const int partner = i ^ j; + if (partner > i) { + const bool descending = ((i & k) == 0); + if ((keys[i] < keys[partner]) == descending) { + const uint64_t t = keys[i]; + keys[i] = keys[partner]; + keys[partner] = t; + } + } + } + if (j >= 32) __syncthreads(); + else __syncwarp(); + } + } + __syncthreads(); +#if MK2_PARALLEL_SAMPLE + mk2_sample_topk_parallel(keys, eos_v, + gumbel + ((size_t)row * N_POS + pos) * TOP_K, out_codes + row); +#else + if (threadIdx.x == 0) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float v; + int n; + if (!eos_used && eos_v > av) { + v = eos_v; + n = AUDIO_EOS; + eos_used = true; + } else { + v = av; + n = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = + v + gumbel[((size_t)row * N_POS + pos) * TOP_K + j]; + if (noised > best_noised) { best_noised = noised; best_code = n; } + } + out_codes[row] = best_code; + } +#endif + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES], act_tile (&a_s)[STAGES], out_tile &o_s, + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0) { + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + const int iters = k_total / KCHUNK; + int tiles_here = 0; + for (int t = blockIdx.x; t < n_tiles; t += gridDim.x) ++tiles_here; + const int iter_base = iter_io; + iter_io += tiles_here * iters * n_tok; + const int my_tiles = (n_tiles - blockIdx.x + gridDim.x - 1) / gridDim.x; + if (warpgroup::groupid() == 1) { + if (warpgroup::laneid() == 0) { + int iter = iter_base; + for (int work = blockIdx.x; work < n_tiles; work += gridDim.x) { + const int tile_id = mk2_expert_tile(work, active_mask, compact_tiles); + for (int tok = 0; tok < n_tok; ++tok) + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(drained[s], phase); + tma::expect_bytes(filled[s], + sizeof(weight_tile) + sizeof(act_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s], w, {tile_id, i}, filled[s]); + const int a_col = out_tiles_per_expert + ? (tile_id / out_tiles_per_expert) * act_k_tiles_per_expert + i + : i; + tma::load_async<2, cache_policy::EVICT_LAST>( + a_s[s], a, {tok, a_col}, filled[s]); + } + } + } + } else { + int iter = iter_base; + for (int work = blockIdx.x; work < n_tiles; work += gridDim.x) + for (int tok = 0; tok < n_tok; ++tok) { + const int tile_id = mk2_expert_tile(work, active_mask, compact_tiles); + rt_fl<16, BATCH_PAD> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s], a_s[s]); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s]); + } + const int base = tile_id * OUT_ATOM; + const int grow0 = tok * BATCH_PAD; + if (epi == EPI_BIAS) { + warpgroup::store(o_sf, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + dest[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(o_sf[int2{col, row}] + + __bfloat162float(bias[base + col])); + } + warpgroup::sync(1); + continue; + } + warpgroup::store(o_s, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_s[int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_w[grow * BB_EXPERTS + expert]; + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * D_MODEL + local] = + wt != 0.f ? __fmul_rn(wt, v) : 0.f; + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1); + } + (void)my_tiles; + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr, + const unsigned char *native_predictor_active_mask = nullptr + MK2_PROGRESS_PARAMETER(native_progress_mask)) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&w_s)[STAGES] = alloc.allocate(); + act_tile (&a_s)[STAGES] = alloc.allocate(); + out_tile (&o_s) = alloc.allocate(); + out_tile_f (&o_sf) = alloc.allocate(); + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * WARPS + 2]; + __shared__ uint64_t sort_keys[SORT_N]; + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS, 0); + for (int c = 0; c < CONSUMER_WARPS; ++c) arrive(drained[s]); + } + } + int gemm_iter = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + mamba_gated_norm(L, bb->zxbcdt, bb->gated); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, nullptr, 1, BATCH, + bb->moe_w, BB_MOE_INTER / OUT_ATOM); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + 1, BATCH, bb->moe_w, D_MODEL / OUT_ATOM); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_deterministic(bb->moe_out, idx); + grid_barrier(g.barrier); + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int n_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = n_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + bb->pf_h[idx] = bb->prompt[idx]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = bb->pf_n_seq > 0 ? bb->pf_n_seq : 1; + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_gated, + T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, T); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * (BB_EXPERTS + 1) * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->pf_normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->pf_moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = bb_moe_fold_deterministic(bb->pf_moe_out, idx); + grid_barrier(g.barrier); + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, bb->pf_final, ln_scratch, + n_total); + grid_barrier(g.barrier); + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, bb->phoneme_b); + grid_barrier(g.barrier); + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + MK2_PROGRESS_GUARD(native_progress_mask, row); +#if MK2_PARALLEL_PHONEME + mk2_phoneme_argmax_parallel( + bb->phoneme_logits + (size_t)row * BB_PHON_PADDED, bb->phoneme_out + row); +#else + if (threadIdx.x == 0) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = 0; n < BB_PHONEME; ++n) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + bb->phoneme_out[row] = arg; + } +#endif + __syncthreads(); + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + __shared__ int native_run_predictor; + if (native_predictor_active_mask && bb && g.native_codes) { + if (threadIdx.x == 0) { + int any_active = 0; + for (int row = 0; row < BATCH; ++row) + any_active |= int(native_predictor_active_mask[row] != 0 && + MK2_PROGRESS_ENABLED(native_progress_mask, row)); + native_run_predictor = any_active; + } + __syncthreads(); + } + const bool skip_native_predictor = native_predictor_active_mask && bb && + g.native_codes && native_run_predictor == 0; + if (skip_native_predictor) { + for (int index = blockIdx.x * THREADS + threadIdx.x; + index < BATCH * N_POS; index += gridDim.x * THREADS) { + MK2_PROGRESS_GUARD(native_progress_mask, index / N_POS); + g.sampled[index] = 0; + g.native_codes[index] = 0; + } + grid_barrier(g.barrier); + } else { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + for (int i = lane; i < D_HEAD; i += 32) { + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + for (int t = 0; t <= pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + for (int i = lane; i < D_HEAD; i += 32) { + float acc = 0.f; + for (int t = 0; t <= pos; ++t) + acc += sc[t] * __bfloat162float( + g.vcache[cache_base + t * D_HEAD + i]); + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + i] = + __float2bfloat16(acc); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, g.sampled_pos, + sort_keys, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { +#ifdef ABLATE_SAMPLE + __syncthreads(); +#endif + } else { + grid_barrier(g.barrier); + } + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x < BATCH && blockIdx.x == 0) { + const int code = g.sampled_pos[threadIdx.x]; + g.sampled[threadIdx.x * N_POS + pos] = code; + if (g.native_codes && MK2_PROGRESS_ENABLED(native_progress_mask, threadIdx.x)) { +#if MK2_INFLIGHT + g.native_codes[threadIdx.x * N_POS + pos] = + native_predictor_active_mask && !native_predictor_active_mask[threadIdx.x] + ? 0 : (long long)code; +#else + g.native_codes[threadIdx.x * N_POS + pos] = (long long)code; +#endif + } + } + grid_barrier(g.barrier); + } + } + if (g.native_seq_len && blockIdx.x == 0 && threadIdx.x < BATCH && + MK2_PROGRESS_ENABLED(native_progress_mask, threadIdx.x)) + g.native_seq_len[threadIdx.x] += 1; + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 2 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * (sizeof(weight_tile) + sizeof(act_tile)) + + sizeof(out_tile) + sizeof(out_tile_f) + 4096; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; + const unsigned char *native_predictor_active_mask = nullptr; + MK2_PROGRESS_MEMBER + int maximum_grid = 0; + bool grid_reset_pending = false; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * (sizeof(weight_tile) + sizeof(act_tile)) + + sizeof(out_tile) + sizeof(out_tile_f) + 4096), 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = std::min(out->sms * out->blocks_per_sm, D_FFN / OUT_ATOM); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 2 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 2 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->maximum_grid = info.sms * info.blocks_per_sm; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.maximum_grid = pred->maximum_grid; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS, nullptr, + native ? st->pred.native_predictor_active_mask : nullptr + MK2_PROGRESS_ARGUMENT(native ? st->pred.native_progress_mask : nullptr)); + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" int mk2_set_native_predictor_active_mask(MK2Predictor *p, + const unsigned char *mask) { + try { + if (!p || reinterpret_cast(p)->kind != 1) + throw std::runtime_error("startup mask requires a full Stage0 handle"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && st->pred.native_predictor_active_mask != mask) + throw std::runtime_error("startup mask pointer cannot change after first Stage0 enqueue"); + if (mask) + mk2_stage0_device_pointer(mask, st->pred.device, "native_predictor_active_mask"); + st->pred.native_predictor_active_mask = mask; + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown startup mask binding failure"; return 1; } +} +extern "C" int mk2_native_startup_skip_abi_version() { return 1; } +#include "../../inflight/runtime.cuh" +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "native_prefill_impl.inc" +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *p, int grid) { + if (!p) { g_last_error = "grid setter requires full Stage0 handle"; return 1; } + auto *st = reinterpret_cast(p); + if (st->kind != 1 || grid < 1 || grid > st->maximum_grid) { + g_last_error = "requested Stage0 grid outside independently resident capacity"; + return 1; + } + if (st->grid != grid) st->grid_reset_pending = true; + st->grid = grid; + g_last_error.clear(); + return 0; +} diff --git a/deploy/nemotron_native/kernels/stage0/b16/mk2_abi.cuh b/deploy/nemotron_native/kernels/stage0/b16/mk2_abi.cuh new file mode 100644 index 0000000..bf1e415 --- /dev/null +++ b/deploy/nemotron_native/kernels/stage0/b16/mk2_abi.cuh @@ -0,0 +1,85 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_set_native_predictor_active_mask(MK2Predictor *p, const unsigned char *mask); +int mk2_native_startup_skip_abi_version(); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *, int grid); diff --git a/deploy/nemotron_native/kernels/stage0/b16/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/stage0/b16/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/stage0/b16/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/stage0/b16/native_prefill_impl.inc b/deploy/nemotron_native/kernels/stage0/b16/native_prefill_impl.inc new file mode 100644 index 0000000..a254308 --- /dev/null +++ b/deploy/nemotron_native/kernels/stage0/b16/native_prefill_impl.inc @@ -0,0 +1,198 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 16 || BATCH == 32 || BATCH == 64)) + throw std::runtime_error("native prefill requires matching batch 1/16/32/64"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (padded > INT_MAX / ((BB_EXPERTS + 1) * D_MODEL)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + float *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(float))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(float))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} diff --git a/deploy/nemotron_native/kernels/stage0/b16/predictor_mk2.cu b/deploy/nemotron_native/kernels/stage0/b16/predictor_mk2.cu new file mode 100644 index 0000000..83810cb --- /dev/null +++ b/deploy/nemotron_native/kernels/stage0/b16/predictor_mk2.cu @@ -0,0 +1,4073 @@ +#include "../../inflight/progress.cuh" +#if MK2_INFLIGHT +#include +#endif +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +constexpr int BATCH_PAD = ((BATCH + 15) / 16) * 16; +static_assert(BATCH >= 1 && BATCH <= BATCH_PAD); +static_assert(BATCH % 16 != 0 || BATCH_PAD == BATCH); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +constexpr int THREADS = 256; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_ACC_F32 }; +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +#ifndef MK2_SPARSE_MOE +#define MK2_SPARSE_MOE 0 +#endif +constexpr int SPARSE_N = 16; +using sparse_act_tile = st_bf; +using sparse_out_tile = st_bf; +using sparse_act_gl = gl; +struct SparseMoe { + bf16 *gathered; + int *rows; + int *counts; + sparse_act_gl gathered_v; + sparse_act_gl hidden_v; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + const SparseMoe *sparse_moe = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +#define BARRIER_GEN_WORD 32 +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + auto *wide_state = reinterpret_cast(state); + const unsigned long long arrival = atomicAdd(wide_state, 1ull); + const unsigned int rank = (unsigned int)(arrival % gridDim.x); + if (rank != gridDim.x - 1u) { + const unsigned long long target = arrival - rank + gridDim.x; + volatile unsigned long long *counter = wide_state; + while (*counter < target) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / THREADS; + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const bf16 *in = src + (size_t)row * D_MODEL; + float v[PER_THREAD], wv[PER_THREAD]; + float sum = 0.f, sq = 0.f; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + v[j] = __bfloat162float(in[threadIdx.x + j * THREADS]); + wv[j] = weight[threadIdx.x + j * THREADS]; + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) { scratch[warp] = sum; scratch[WARPS + warp] = sq; } + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f, q = 0.f; + for (int w = 0; w < WARPS; ++w) { s += scratch[w]; q += scratch[WARPS + w]; } + scratch[2 * WARPS] = s / D_MODEL; + scratch[2 * WARPS + 1] = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + } + __syncthreads(); + const float mean = scratch[2 * WARPS], inv = scratch[2 * WARPS + 1]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * wv[j]); + } + __syncthreads(); + } +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, bf16 *snapshot = nullptr, int snapshot_layer = -1, + float *snapshot_fp32 = nullptr, float *out_fp32 = nullptr) { + constexpr int PER_THREAD = D_MODEL / THREADS; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const size_t base = (size_t)row * D_MODEL; + float v[PER_THREAD]; + bf16 wv[PER_THREAD]; + float sq = 0.f; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + wv[j] = weight[i]; + const bf16 hv = h[base + i]; + if (snapshot) + snapshot[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = hv; + if (snapshot_fp32) + snapshot_fp32[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = + __bfloat162float(hv); + v[j] = __bfloat162float(hv) + __bfloat162float(r[base + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) scratch[warp] = sq; + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f; + for (int w = 0; w < WARPS; ++w) s += scratch[w]; + scratch[WARPS] = rsqrtf(s / D_MODEL + RMS_EPS); + } + __syncthreads(); + const float inv = scratch[WARPS]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + const bf16 normed = __float2bfloat16( + __bfloat162float(wv[j]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + out[base + i] = normed; + if (out_fp32) out_fp32[base + i] = __bfloat162float(normed); + } + __syncthreads(); + } +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map MK2_PROGRESS_PARAMETER(progress_mask)) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + MK2_PROGRESS_GUARD(progress_mask, row); + const int c = idx - row * MB_CONV_DIM; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + const float raw = __bfloat162float( + zxbcdt[(size_t)row * MB_IN_PROJ + MB_INNER + c]); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + xbc_act[(size_t)row * MB_CONV_DIM + c] = __float2bfloat16(acc); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map MK2_PROGRESS_PARAMETER(progress_mask)) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + const int pairs = BATCH * MB_HEADS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_HEADS * split); + MK2_PROGRESS_GUARD(progress_mask, row); + const int head = (pair / split) % MB_HEADS; + const int first_channel = (pair % split) * channels; + const int group = head / MB_HEADS_PER_GROUP; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + L.dt_bias[head]); + const float dA = __expf(L.A[head] * dt); + const float Dh = L.D[head]; + const size_t ssm_at = (size_t)slot_map[row] * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float *st = L.ssm_state + ssm_at; + const float *st_in = st; + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + constexpr int NSTATE = MB_STATE / 32; + constexpr int SSM_PREFETCH = channels < 8 ? channels : 8; + static_assert(channels % SSM_PREFETCH == 0); + for (int g0 = 0; g0 < channels; g0 += SSM_PREFETCH) { + float sv[SSM_PREFETCH][NSTATE]; + float xvr[SSM_PREFETCH]; + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const size_t prow = (size_t)(first_channel + g0 + gi) * MB_STATE; + xvr[gi] = __bfloat162float( + xa[head * MB_HEAD_DIM + first_channel + g0 + gi]); + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = st_in[prow + lane + j * 32]; + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const int p = first_channel + g0 + gi; + const float xv = xvr[gi]; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < NSTATE; ++j) { + const int n = lane + j * 32; + const size_t at = (size_t)p * MB_STATE + n; + const float s = sv[gi][j] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_CONV_DIM; + const int c = idx - t * MB_CONV_DIM; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v; +#ifdef PF_CONV_ZERO_STATE + if (src < 0) v = 0.f; +#else + if (src < 0) v = __bfloat162float(state[pos + k]); +#endif + else v = __bfloat162float( + zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]); + acc += __bfloat162float(w[k]) * v; + } +#ifndef PF_CONV_NO_SILU + acc = acc / (1.f + __expf(-acc)); +#endif + xbc_act[(size_t)t * MB_CONV_DIM + c] = __float2bfloat16(acc); + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + const int units = n_seq * MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int seq = unit / (MB_HEADS * MB_HEAD_DIM); + const int rem = unit - seq * (MB_HEADS * MB_HEAD_DIM); + const int head = rem / MB_HEAD_DIM; + const int p = rem - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p * MB_STATE; + float sreg[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) sreg[j] = st[lane + j * 32]; + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const bf16 *xa = xbc_act + t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + const bf16 *dtrow = zxbcdt + t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + sreg[j] = sreg[j] * dA + __bfloat162float(Bsec[n]) * dt * xv; + acc += sreg[j] * __bfloat162float(Csec[n]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) st[lane + j * 32] = sreg[j]; + } + return; + } +#if PF_SSM_SPLIT_P + const int units = MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; unit += warps_total) { + const int head = unit / MB_HEAD_DIM; + const int p_only = unit - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; +#else + for (int head = blockIdx.x * WARPS + warp; head < MB_HEADS; + head += warps_total) { + const int group = head / MB_HEADS_PER_GROUP; +#endif + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *st = L.ssm_state + (size_t)slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE +#if PF_SSM_SPLIT_P + + (size_t)p_only * MB_STATE +#endif + ; + for (int t = 0; t < n_tokens; ++t) { + const bf16 *xa = xbc_act + (size_t)t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = + zxbcdt + (size_t)t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } +#if PF_SSM_SPLIT_P + { + const int p = p_only; +#else + for (int p = 0; p < MB_HEAD_DIM; ++p) { +#endif + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { +#if PF_SSM_SPLIT_P + const size_t at = (size_t)lane + j * 32; +#else + const size_t at = (size_t)p * MB_STATE + lane + j * 32; +#endif + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } + } +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32], wg[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + wg[j] = __bfloat162float(L.gnorm_w[idx]); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16(held[j] * inv * wg[j]); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len MK2_PROGRESS_PARAMETER(progress_mask)) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + constexpr int Q_PER_WARP = 2; + constexpr int GROUPS_PER_KV = 2; + const int pairs = BATCH * BB_KV_HEADS * GROUPS_PER_KV; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x + warp * gridDim.x; pair < pairs; + pair += gridDim.x * WARPS) { + const int row = pair / (BB_KV_HEADS * GROUPS_PER_KV); + MK2_PROGRESS_GUARD(progress_mask, row); + const int within_row = pair % (BB_KV_HEADS * GROUPS_PER_KV); + const int kv = within_row / GROUPS_PER_KV; + const int q_group = within_row % GROUPS_PER_KV; + const int q_first = q_group * Q_PER_WARP; + const int q_count = q_group == 0 ? Q_PER_WARP : 1; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (q_group == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[h][j] = 0.f; + } + float m[Q_PER_WARP] = {-CUDART_INF_F, -CUDART_INF_F}; + float denom[Q_PER_WARP] = {0.f, 0.f}; + constexpr int PF = 8; + const int tpb = L.tokens_per_block; + const size_t elems_per_block = + (size_t)tpb * BB_KV_HEADS * BB_HEAD_DIM; + const int *bt_k = L.block_table + (size_t)(row * 2) * L.max_blocks_per_seq; + const int *bt_v = bt_k + L.max_blocks_per_seq; + const int head_base = kv * tpb * BB_HEAD_DIM; + float qr[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const bf16 *q = base + (kv * kv_per_q + q_first + h) * BB_HEAD_DIM; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qr[h][j] = __bfloat162float(q[lane + j * 32]); + } + for (int page0 = 0; page0 < T; page0 += tpb) { + const int k_page = bt_k[page0 / tpb] & 0x7FFFFFFF; + const int v_page = bt_v[page0 / tpb] & 0x7FFFFFFF; + const bf16 *kbase = L.pool_base + + (size_t)k_page * elems_per_block + + head_base; + const bf16 *vbase = L.pool_base + + (size_t)v_page * elems_per_block + + head_base; + const int tend = page0 + tpb < T ? page0 + tpb : T; + for (int t0 = page0; t0 < tend; t0 += PF) { + const int navail = tend - t0 < PF ? tend - t0 : PF; + float kr[PF][PER_LANE], vr[PF][PER_LANE]; + #pragma unroll + for (int u = 0; u < PF; ++u) { + const int tpos = t0 + (u < navail ? u : navail - 1); +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = vbase + (tpos - page0) * BB_HEAD_DIM; +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new : kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = cur_token + ? v_new : vbase + (tpos - page0) * BB_HEAD_DIM; +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + kr[u][j] = __bfloat162float(kt[i]); + vr[u][j] = __bfloat162float(vt[i]); + } + } + #pragma unroll + for (int u = 0; u < PF; ++u) { + if (u >= navail) break; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + dot += qr[h][j] * kr[u][j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m[h], s); + const float rescale = __expf(m[h] - m_new); + const float w = __expf(s - m_new); + denom[h] = denom[h] * rescale + w; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[h][j] = acc[h][j] * rescale + w * vr[u][j]; + m[h] = m_new; + } + } + } + } + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const float inv = 1.f / denom[h]; + const int head = kv * kv_per_q + q_first + h; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[h][j] * inv); + } + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * BB_KV_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / BB_KV_HEADS; + const int kv = idx - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv); + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv); + for (int i = 0; i < BB_HEAD_DIM; ++i) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + for (int s = 0; s <= pos; ++s) { + const bf16 *kt = kv_slot(L, seq, 0, s) + kv_offset(L, s, kv); + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + const bf16 *vt = kv_slot(L, seq, 1, s) + kv_offset(L, s, kv); + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + wt * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH) { + __shared__ float score[BB_EXPERTS]; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + for (int e = warp; e < BB_EXPERTS; e += WARPS) { + const bf16 *wr = gate_w + (size_t)e * D_MODEL; + float acc = 0.f; + for (int i = lane; i < D_MODEL; i += 32) + acc += __bfloat162float(normed[(size_t)row * D_MODEL + i]) * + __bfloat162float(wr[i]); + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) score[e] = 1.f / (1.f + __expf(-acc)); + } + __syncthreads(); + if (threadIdx.x == 0) { + float remaining[BB_EXPERTS]; + int pick[BB_TOP_K]; + for (int e = 0; e < BB_EXPERTS; ++e) { + moe_w[row * BB_EXPERTS + e] = 0.f; + remaining[e] = score[e] + __bfloat162float(gate_bias[e]); + } + float total = 0.f; + for (int k = 0; k < BB_TOP_K; ++k) { + int best = 0; + for (int e = 1; e < BB_EXPERTS; ++e) + if (remaining[e] > remaining[best]) best = e; + pick[k] = best; + total += score[best]; + remaining[best] = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + for (int k = 0; k < BB_TOP_K; ++k) + moe_w[row * BB_EXPERTS + pick[k]] = score[pick[k]] * norm; + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert <= BB_EXPERTS; ++expert) + value = __fadd_rn(value, parts[base + (size_t)expert * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ bf16 bb_moe_fold_selected( + const float *parts, const float *weights, int idx) { + const int row = idx / D_MODEL; + const size_t base = (size_t)row * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert < BB_EXPERTS; ++expert) { + const float part = weights[row * BB_EXPERTS + expert] != 0.f + ? parts[base + (size_t)expert * D_MODEL] : 0.f; + value = __fadd_rn(value, part); + } + value = __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ __forceinline__ void sample_argmax_merge( + float &value, int &rank, float other, int other_rank) { + if (other > value || (other == value && other_rank < rank)) { + value = other; + rank = other_rank; + } +} +__device__ __forceinline__ int sample_block_argmax(float value, int rank) { + __shared__ float warp_value[WARPS]; + __shared__ int warp_rank[WARPS]; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + if (lane == 0) { warp_value[warp] = value; warp_rank[warp] = rank; } + __syncthreads(); + if (warp == 0) { + value = lane < WARPS ? warp_value[lane] : -CUDART_INF_F; + rank = lane < WARPS ? warp_rank[lane] : INT_MAX; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + } + return rank; +} +__device__ __forceinline__ int sample_eos_insertion_rank( + const uint64_t *keys, float eos) { + int lo = 0, hi = TOP_K; + while (lo < hi) { + const int mid = (lo + hi) >> 1; + if (eos > ordered_to_float((uint32_t)(keys[mid] >> 32))) hi = mid; + else lo = mid + 1; + } + return lo; +} +__device__ __forceinline__ void sample_topk_serial( + const uint64_t *keys, float eos, const float *noise, int *out) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float value; + int code; + if (!eos_used && eos > av) { + value = eos; + code = AUDIO_EOS; + eos_used = true; + } else { + value = av; + code = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = value + noise[j]; + if (noised > best_noised) { + best_noised = noised; + best_code = code; + } + } + *out = best_code; +} +__device__ __forceinline__ void sample_topk_parallel( + const uint64_t *keys, float eos, const float *noise, int *out) { + const int j = threadIdx.x; + const bool sorted_nan = j < TOP_K && + isnan(ordered_to_float((uint32_t)(keys[j] >> 32))); + if (__syncthreads_or(sorted_nan)) { + if (threadIdx.x == 0) sample_topk_serial(keys, eos, noise, out); + return; + } + const int lane = threadIdx.x & 31; + int insertion = lane == 0 ? sample_eos_insertion_rank(keys, eos) : 0; + insertion = __shfl_sync(0xffffffffu, insertion, 0); + float value = -CUDART_INF_F; + int rank = INT_MAX; + if (j < TOP_K) { + const int a = j - (j > insertion); + const float v = j == insertion + ? eos : ordered_to_float((uint32_t)(keys[a] >> 32)); + const float noised = v + noise[j]; + if (noised > -CUDART_INF_F) { value = noised; rank = j; } + } + const int winner = sample_block_argmax(value, rank); + if (threadIdx.x == 0) { + if (winner == INT_MAX) *out = 0; + else if (winner == insertion) *out = AUDIO_EOS; + else { + const int a = winner - (winner > insertion); + *out = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + } + } +} +constexpr int SORT_ITEMS_PER_THREAD = SORT_N / THREADS; +struct sampler_storage { + uint64_t keys[2][SORT_N]; +}; +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + sampler_storage &storage, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + uint64_t thread_keys[SORT_ITEMS_PER_THREAD]; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + thread_keys[item] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + int cross_stage = 0; + for (int k = 2; k <= SORT_N; k <<= 1) { + for (int j = k >> 1; j > 0; j >>= 1) { + if (j >= 32) { + const int bank = cross_stage++ & 1; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + storage.keys[bank][n] = thread_keys[item]; + } + __syncthreads(); + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + const uint64_t other = storage.keys[bank][n ^ j]; + const bool take_hi = ((n & k) == 0) == ((n & j) == 0); + const uint64_t mine = thread_keys[item]; + thread_keys[item] = take_hi + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } else { + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + const uint64_t other = __shfl_xor_sync( + 0xffffffffu, thread_keys[item], j); + const bool take_hi = ((n & k) == 0) == ((n & j) == 0); + const uint64_t mine = thread_keys[item]; + thread_keys[item] = take_hi + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } + } + } + const int top_bank = cross_stage & 1; + if (threadIdx.x < TOP_K) + storage.keys[top_bank][threadIdx.x] = thread_keys[0]; + __syncthreads(); + sample_topk_parallel( + storage.keys[top_bank], eos_v, + gumbel + ((size_t)row * N_POS + pos) * TOP_K, + out_codes + row); + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES], act_tile (&a_s)[STAGES], out_tile &o_s, + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + int &tile_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0) { + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + const int iters = k_total / KCHUNK; + int tile_start = (int)blockIdx.x - tile_io % (int)gridDim.x; + if (tile_start < 0) tile_start += (int)gridDim.x; + tile_io += n_tiles; + int tiles_here = 0; + for (int t = tile_start; t < n_tiles; t += gridDim.x) ++tiles_here; + const int iter_base = iter_io; + iter_io += tiles_here * iters * n_tok; + const int my_tiles = (n_tiles - tile_start + gridDim.x - 1) / gridDim.x; + if (warpgroup::groupid() == 1) { + if (warpgroup::laneid() == 0) { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) { + const int tile_id = mk2_expert_tile(work, active_mask, compact_tiles); + for (int tok = 0; tok < n_tok; ++tok) + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(drained[s], phase); + tma::expect_bytes(filled[s], + sizeof(weight_tile) + sizeof(act_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s], w, {tile_id, i}, filled[s]); + const int a_col = out_tiles_per_expert + ? (tile_id / out_tiles_per_expert) * act_k_tiles_per_expert + i + : i; + tma::load_async<2, cache_policy::EVICT_LAST>( + a_s[s], a, {tok, a_col}, filled[s]); + } + } + } + } else { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) + for (int tok = 0; tok < n_tok; ++tok) { + const int tile_id = mk2_expert_tile(work, active_mask, compact_tiles); + rt_fl<16, BATCH_PAD> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s], a_s[s]); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s]); + } + const int base = tile_id * OUT_ATOM; + const int grow0 = tok * BATCH_PAD; + if (epi == EPI_BIAS) { + warpgroup::store(o_sf, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + dest[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(o_sf[int2{col, row}] + + __bfloat162float(bias[base + col])); + } + warpgroup::sync(1); + continue; + } + warpgroup::store(o_s, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_s[int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_w[grow * BB_EXPERTS + expert]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * + D_MODEL + local] = __fmul_rn(wt, v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1); + } + (void)my_tiles; + } +} +#include "sparse_moe_impl.cuh" +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr, + const unsigned char *native_predictor_active_mask = nullptr + MK2_PROGRESS_PARAMETER(native_progress_mask)) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&w_s)[STAGES] = alloc.allocate(); + act_tile (&a_s)[STAGES] = alloc.allocate(); + out_tile (&o_s) = alloc.allocate(); + out_tile_f (&o_sf) = alloc.allocate(); + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * WARPS + 2]; + __shared__ sampler_storage sort_storage; + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS, 0); + for (int c = 0; c < CONSUMER_WARPS; ++c) arrive(drained[s]); + } + } + int gemm_iter = 0; + int gemm_tile = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + mamba_gated_norm(L, bb->zxbcdt, bb->gated); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w); + } + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + } + grid_barrier(g.barrier); + if constexpr (MK2_SPARSE_MOE && BATCH >= 16) { + if (bb->sparse_moe) { + sparse_moe_pack(bb); + grid_barrier(g.barrier); + } + } + if (bb->sparse_moe) { + sparse_moe_gemm(L.exp_up, bb->sparse_moe->gathered_v, + *bb->sparse_moe, bb->moe_h, bb->moe_w, bb->moe_out, + w_s, a_s, o_s, filled, drained, gemm_iter, gemm_tile); + } else { + gemm(L.exp_up, bb->normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, nullptr, 1, BATCH, + bb->moe_w, BB_MOE_INTER / OUT_ATOM); + } + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + if (bb->sparse_moe) { + sparse_moe_gemm(L.exp_down, bb->sparse_moe->hidden_v, + *bb->sparse_moe, bb->moe_h, bb->moe_w, bb->moe_out, + w_s, a_s, o_s, filled, drained, gemm_iter, gemm_tile); + } else { + gemm(L.exp_down, bb->moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + 1, BATCH, bb->moe_w, D_MODEL / OUT_ATOM); + } + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_selected(bb->moe_out, bb->moe_w, idx); + grid_barrier(g.barrier); + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int n_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = n_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + bb->pf_h[idx] = bb->prompt[idx]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = bb->pf_n_seq > 0 ? bb->pf_n_seq : 1; + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_gated, + T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, T); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * (BB_EXPERTS + 1) * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->pf_normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->pf_moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = bb_moe_fold_deterministic(bb->pf_moe_out, idx); + grid_barrier(g.barrier); + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, bb->pf_final, ln_scratch, + n_total); + grid_barrier(g.barrier); + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, bb->phoneme_b); + grid_barrier(g.barrier); + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + MK2_PROGRESS_GUARD(native_progress_mask, row); + if (threadIdx.x == 0) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = 0; n < BB_PHONEME; ++n) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + bb->phoneme_out[row] = arg; + } + __syncthreads(); + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + __shared__ int native_run_predictor; + if (native_predictor_active_mask && bb && g.native_codes) { + if (threadIdx.x == 0) { + int any_active = 0; + for (int row = 0; row < BATCH; ++row) + any_active |= int(native_predictor_active_mask[row] != 0 && + MK2_PROGRESS_ENABLED(native_progress_mask, row)); + native_run_predictor = any_active; + } + __syncthreads(); + } + const bool skip_native_predictor = native_predictor_active_mask && bb && + g.native_codes && native_run_predictor == 0; + if (skip_native_predictor) { + for (int index = blockIdx.x * THREADS + threadIdx.x; + index < BATCH * N_POS; index += gridDim.x * THREADS) { + MK2_PROGRESS_GUARD(native_progress_mask, index / N_POS); + g.sampled[index] = 0; + g.native_codes[index] = 0; + } + grid_barrier(g.barrier); + } else { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + constexpr int DPL = D_HEAD / 32; + float qv[DPL], kn[DPL], vn[DPL]; + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + qv[j] = __bfloat162float(q[i]); + kn[j] = __bfloat162float(k_new[i]); + vn[j] = __bfloat162float(v_new[i]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += qv[j] * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + { + float acc = 0.f; + #pragma unroll + for (int j = 0; j < DPL; ++j) acc += qv[j] * kn[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[pos] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + float ov[DPL] = {}; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *vt = g.vcache + cache_base + t * D_HEAD; + const float s = sc[t]; + #pragma unroll + for (int j = 0; j < DPL; ++j) + ov[j] += s * __bfloat162float(vt[lane + j * 32]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + ov[j] += sc[pos] * vn[j]; + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + + lane + j * 32] = __float2bfloat16(ov[j]); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, + g.sampled_pos, sort_storage, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif +#ifdef ABLATE_SAMPLE + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { + __syncthreads(); + } else { + grid_barrier(g.barrier); + } +#endif + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x == 0) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + g.sampled[row * N_POS + pos] = code; + if (g.native_codes && MK2_PROGRESS_ENABLED(native_progress_mask, row)) { +#if MK2_INFLIGHT + g.native_codes[row * N_POS + pos] = + native_predictor_active_mask && !native_predictor_active_mask[row] + ? 0 : (long long)code; +#else + g.native_codes[row * N_POS + pos] = (long long)code; +#endif + } + } + } +#ifdef ABLATE_SAMPLE + grid_barrier(g.barrier); +#else + if (pos + 1 == n_pos) grid_barrier(g.barrier); +#endif + } + } + if (g.native_seq_len && blockIdx.x == 0 && threadIdx.x < BATCH && + MK2_PROGRESS_ENABLED(native_progress_mask, threadIdx.x)) + g.native_seq_len[threadIdx.x] += 1; + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 64 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * (sizeof(weight_tile) + sizeof(act_tile)) + + sizeof(out_tile) + sizeof(out_tile_f) + 4096; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; + const unsigned char *native_predictor_active_mask = nullptr; + MK2_PROGRESS_MEMBER + int maximum_grid = 0; + bool grid_reset_pending = false; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * (sizeof(weight_tile) + sizeof(act_tile)) + + sizeof(out_tile) + sizeof(out_tile_f) + 4096), 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = std::min(out->sms * out->blocks_per_sm, 4096); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 64 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->maximum_grid = info.sms * info.blocks_per_sm; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + if constexpr (MK2_SPARSE_MOE && BATCH >= 16) { + bf16 *gathered = nullptr; + int *rows = nullptr, *counts = nullptr; + SparseMoe *device_sparse = nullptr; + mk2_check(owned.alloc(&gathered, + (size_t)BB_EXPERTS * BATCH_PAD * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&rows, (size_t)BB_EXPERTS * BATCH_PAD * sizeof(int))); + mk2_check(owned.alloc(&counts, BB_EXPERTS * sizeof(int))); + SparseMoe sparse{gathered, rows, counts, + sparse_act_gl{gathered, nullptr, nullptr, (size_t)BB_EXPERTS * BATCH_PAD, + (size_t)D_MODEL}, + sparse_act_gl{moe_h, nullptr, nullptr, (size_t)BB_EXPERTS * BATCH_PAD, + (size_t)BB_MOE_INTER}}; + mk2_check(owned.alloc(&device_sparse, sizeof(SparseMoe))); + mk2_check(cudaMemcpy(device_sparse, &sparse, sizeof(sparse), cudaMemcpyHostToDevice)); + hb.sparse_moe = device_sparse; + } + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.maximum_grid = pred->maximum_grid; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(s, &capture_status)); + if (st->pred.grid_reset_pending || capture_status != cudaStreamCaptureStatusNone) { + mk2_check(cudaMemsetAsync(st->pred.g.barrier, 0, 64 * sizeof(unsigned int), s)); + st->pred.grid_reset_pending = false; + } + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS, nullptr, + native ? st->pred.native_predictor_active_mask : nullptr + MK2_PROGRESS_ARGUMENT(native ? st->pred.native_progress_mask : nullptr)); + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" int mk2_set_native_predictor_active_mask(MK2Predictor *p, + const unsigned char *mask) { + try { + if (!p || reinterpret_cast(p)->kind != 1) + throw std::runtime_error("startup mask requires a full Stage0 handle"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && st->pred.native_predictor_active_mask != mask) + throw std::runtime_error("startup mask pointer cannot change after first Stage0 enqueue"); + if (mask) + mk2_stage0_device_pointer(mask, st->pred.device, "native_predictor_active_mask"); + st->pred.native_predictor_active_mask = mask; + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown startup mask binding failure"; return 1; } +} +extern "C" int mk2_native_startup_skip_abi_version() { return 1; } +#include "../../inflight/runtime.cuh" +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "native_prefill_impl.inc" +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *p, int grid) { + if (!p) { g_last_error = "grid setter requires full Stage0 handle"; return 1; } + auto *st = reinterpret_cast(p); + if (st->kind != 1 || grid < 1 || grid > st->maximum_grid) { + g_last_error = "requested Stage0 grid outside independently resident capacity"; + return 1; + } + if (st->grid != grid) st->grid_reset_pending = true; + st->grid = grid; + g_last_error.clear(); + return 0; +} diff --git a/deploy/nemotron_native/kernels/stage0/b16/sparse_moe_impl.cuh b/deploy/nemotron_native/kernels/stage0/b16/sparse_moe_impl.cuh new file mode 100644 index 0000000..056890d --- /dev/null +++ b/deploy/nemotron_native/kernels/stage0/b16/sparse_moe_impl.cuh @@ -0,0 +1,139 @@ +__device__ void sparse_moe_pack(const Backbone *bb) { + const SparseMoe &s = *bb->sparse_moe; + __shared__ int count; + __shared__ int local_rows[BATCH]; + const int experts_per_wave = min(int(gridDim.x), BB_EXPERTS); + const int first_expert = int(blockIdx.x) % experts_per_wave; + const int rank = int(blockIdx.x) / experts_per_wave; + const int team_size = (int(gridDim.x) + experts_per_wave - 1 - first_expert) / + experts_per_wave; + for (int e = first_expert; e < BB_EXPERTS; e += experts_per_wave) { + if (threadIdx.x == 0) { + int n = 0; + for (int row = 0; row < BATCH; ++row) + if (bb->moe_w[row * BB_EXPERTS + e] != 0.f) + local_rows[n++] = row; + count = n; + if (rank == 0) { + s.counts[e] = n; + for (int i = 0; i < n; ++i) + s.rows[e * BATCH_PAD + i] = local_rows[i]; + } + } + __syncthreads(); + constexpr int chunks = D_MODEL / 8; + const int padded = ((count + SPARSE_N - 1) / SPARSE_N) * SPARSE_N; + for (int i = rank * THREADS + threadIdx.x; i < padded * chunks; + i += team_size * THREADS) { + const int row = i / chunks, col = i % chunks; + const int packed = e * BATCH_PAD + row; + auto *dst = reinterpret_cast(s.gathered + + (size_t)packed * D_MODEL + col * 8); + if (row < count) { + const int original = local_rows[row]; + *dst = *reinterpret_cast(bb->normed + + (size_t)original * D_MODEL + col * 8); + } else { + *dst = make_float4(0.f, 0.f, 0.f, 0.f); + } + } + __syncthreads(); + } +} +template +__device__ void sparse_moe_gemm( + const weight_gl &w, const sparse_act_gl &a, const SparseMoe &s, + bf16 *hidden, const float *routes, float *contributions, + weight_tile (&w_s)[STAGES], act_tile (&a_s)[STAGES], out_tile &o_s, + semaphore *filled, semaphore *drained, int &iter_io, int &tile_io) { + constexpr int width = DOWN ? D_MODEL : BB_MOE_INTER; + constexpr int k_total = DOWN ? BB_MOE_INTER : D_MODEL; + constexpr int out_tiles = width / OUT_ATOM; + constexpr int iters = k_total / KCHUNK; + static_assert(BATCH_PAD >= SPARSE_N && BATCH_PAD % SPARSE_N == 0); + static_assert(sizeof(sparse_act_tile) <= sizeof(act_tile)); + static_assert(sizeof(sparse_out_tile) <= sizeof(out_tile)); + __shared__ int prefix[BB_EXPERTS + 1]; + __shared__ int row_tiles[BB_EXPERTS]; + if (threadIdx.x == 0) { + int total = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + prefix[e] = total; + row_tiles[e] = (s.counts[e] + SPARSE_N - 1) / SPARSE_N; + total += row_tiles[e] * out_tiles; + } + prefix[BB_EXPERTS] = total; + } + __syncthreads(); + const int n_tiles = prefix[BB_EXPERTS]; + int start = int(blockIdx.x) - tile_io % int(gridDim.x); + if (start < 0) start += gridDim.x; + tile_io += n_tiles; + const int mine = start < n_tiles ? 1 + (n_tiles - 1 - start) / gridDim.x : 0; + const int iter_base = iter_io; + iter_io += mine * iters; + if (warpgroup::groupid() == 1) { + if (warpgroup::laneid() == 0) { + int iter = iter_base; + for (int f = start; f < n_tiles; f += gridDim.x) { + int e = 0; + while (prefix[e + 1] <= f) ++e; + const int rem = f - prefix[e], nt = row_tiles[e]; + const int feature_tile = rem / nt; + const int token_tile = rem - feature_tile * nt; + const int weight_row = e * out_tiles + feature_tile; + const int act_row = e * (BATCH_PAD / SPARSE_N) + token_tile; + for (int k = 0; k < iters; ++k, ++iter) { + const int stage = iter % STAGES, phase = (iter / STAGES) & 1; + wait(drained[stage], phase); + tma::expect_bytes(filled[stage], sizeof(weight_tile) + sizeof(sparse_act_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[stage], w, {weight_row, k}, filled[stage]); + tma::load_async<2, cache_policy::EVICT_LAST>( + *reinterpret_cast(&a_s[stage]), + a, {act_row, k}, filled[stage]); + } + } + } + } else if (warpgroup::groupid() == 0) { + int iter = iter_base; + for (int f = start; f < n_tiles; f += gridDim.x) { + int e = 0; + while (prefix[e + 1] <= f) ++e; + const int rem = f - prefix[e], nt = row_tiles[e]; + const int feature_tile = rem / nt; + const int token_tile = rem - feature_tile * nt; + rt_fl<16, SPARSE_N> accum; + warp::zero(accum); + for (int k = 0; k < iters; ++k, ++iter) { + const int stage = iter % STAGES, phase = (iter / STAGES) & 1; + wait(filled[stage], phase); + warpgroup::mma_ABt(accum, w_s[stage], + *reinterpret_cast(&a_s[stage])); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[stage]); + } + sparse_out_tile &out = *reinterpret_cast(&o_s); + warpgroup::store(out, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int i = warpgroup::laneid(); i < OUT_ATOM * SPARSE_N; i += 128) { + const int row = i / OUT_ATOM, col = i % OUT_ATOM; + const int local = token_tile * SPARSE_N + row; + const int packed = e * BATCH_PAD + local; + const int output = feature_tile * OUT_ATOM + col; + const float value = float(out[{col, row}]); + if constexpr (!DOWN) { + hidden[(size_t)packed * BB_MOE_INTER + output] = + bf16(value / (1.f + __expf(-value))); + } else if (local < s.counts[e]) { + const int original = s.rows[packed]; + const float weight = routes[original * BB_EXPERTS + e]; + contributions[((size_t)original * (BB_EXPERTS + 1) + e) * + D_MODEL + output] = __fmul_rn(weight, value); + } + } + warpgroup::sync(1); + } + } +} diff --git a/deploy/nemotron_native/kernels/stage0/b64/mk2_abi.cuh b/deploy/nemotron_native/kernels/stage0/b64/mk2_abi.cuh new file mode 100644 index 0000000..bf1e415 --- /dev/null +++ b/deploy/nemotron_native/kernels/stage0/b64/mk2_abi.cuh @@ -0,0 +1,85 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_set_native_predictor_active_mask(MK2Predictor *p, const unsigned char *mask); +int mk2_native_startup_skip_abi_version(); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *, int grid); diff --git a/deploy/nemotron_native/kernels/stage0/b64/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/stage0/b64/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/stage0/b64/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/stage0/b64/native_prefill_impl.inc b/deploy/nemotron_native/kernels/stage0/b64/native_prefill_impl.inc new file mode 100644 index 0000000..a254308 --- /dev/null +++ b/deploy/nemotron_native/kernels/stage0/b64/native_prefill_impl.inc @@ -0,0 +1,198 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 16 || BATCH == 32 || BATCH == 64)) + throw std::runtime_error("native prefill requires matching batch 1/16/32/64"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (padded > INT_MAX / ((BB_EXPERTS + 1) * D_MODEL)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + float *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(float))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(float))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} diff --git a/deploy/nemotron_native/kernels/stage0/b64/predictor_mk2.cu b/deploy/nemotron_native/kernels/stage0/b64/predictor_mk2.cu new file mode 100644 index 0000000..a4f947e --- /dev/null +++ b/deploy/nemotron_native/kernels/stage0/b64/predictor_mk2.cu @@ -0,0 +1,4022 @@ +#include "../../inflight/progress.cuh" +#if MK2_INFLIGHT +#include +#endif +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +constexpr int BATCH_PAD = ((BATCH + 15) / 16) * 16; +static_assert(BATCH >= 1 && BATCH <= BATCH_PAD); +static_assert(BATCH % 16 != 0 || BATCH_PAD == BATCH); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +constexpr int THREADS = 256; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_ACC_F32 }; +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +#define BARRIER_GEN_WORD 32 +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + auto *wide_state = reinterpret_cast(state); + const unsigned long long arrival = atomicAdd(wide_state, 1ull); + const unsigned int rank = (unsigned int)(arrival % gridDim.x); + if (rank != gridDim.x - 1u) { + const unsigned long long target = arrival - rank + gridDim.x; + volatile unsigned long long *counter = wide_state; + while (*counter < target) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / THREADS; + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const bf16 *in = src + (size_t)row * D_MODEL; + float v[PER_THREAD], wv[PER_THREAD]; + float sum = 0.f, sq = 0.f; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + v[j] = __bfloat162float(in[threadIdx.x + j * THREADS]); + wv[j] = weight[threadIdx.x + j * THREADS]; + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) { scratch[warp] = sum; scratch[WARPS + warp] = sq; } + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f, q = 0.f; + for (int w = 0; w < WARPS; ++w) { s += scratch[w]; q += scratch[WARPS + w]; } + scratch[2 * WARPS] = s / D_MODEL; + scratch[2 * WARPS + 1] = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + } + __syncthreads(); + const float mean = scratch[2 * WARPS], inv = scratch[2 * WARPS + 1]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * wv[j]); + } + __syncthreads(); + } +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, bf16 *snapshot = nullptr, int snapshot_layer = -1, + float *snapshot_fp32 = nullptr, float *out_fp32 = nullptr) { + constexpr int PER_THREAD = D_MODEL / THREADS; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const size_t base = (size_t)row * D_MODEL; + float v[PER_THREAD]; + bf16 wv[PER_THREAD]; + float sq = 0.f; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + wv[j] = weight[i]; + const bf16 hv = h[base + i]; + if (snapshot) + snapshot[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = hv; + if (snapshot_fp32) + snapshot_fp32[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = + __bfloat162float(hv); + v[j] = __bfloat162float(hv) + __bfloat162float(r[base + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) scratch[warp] = sq; + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f; + for (int w = 0; w < WARPS; ++w) s += scratch[w]; + scratch[WARPS] = rsqrtf(s / D_MODEL + RMS_EPS); + } + __syncthreads(); + const float inv = scratch[WARPS]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + const bf16 normed = __float2bfloat16( + __bfloat162float(wv[j]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + out[base + i] = normed; + if (out_fp32) out_fp32[base + i] = __bfloat162float(normed); + } + __syncthreads(); + } +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map MK2_PROGRESS_PARAMETER(progress_mask)) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + MK2_PROGRESS_GUARD(progress_mask, row); + const int c = idx - row * MB_CONV_DIM; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + const float raw = __bfloat162float( + zxbcdt[(size_t)row * MB_IN_PROJ + MB_INNER + c]); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + xbc_act[(size_t)row * MB_CONV_DIM + c] = __float2bfloat16(acc); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map MK2_PROGRESS_PARAMETER(progress_mask)) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + const int pairs = BATCH * MB_HEADS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_HEADS * split); + MK2_PROGRESS_GUARD(progress_mask, row); + const int head = (pair / split) % MB_HEADS; + const int first_channel = (pair % split) * channels; + const int group = head / MB_HEADS_PER_GROUP; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + L.dt_bias[head]); + const float dA = __expf(L.A[head] * dt); + const float Dh = L.D[head]; + const size_t ssm_at = (size_t)slot_map[row] * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float *st = L.ssm_state + ssm_at; + const float *st_in = st; + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + constexpr int NSTATE = MB_STATE / 32; + constexpr int SSM_PREFETCH = channels < 8 ? channels : 8; + static_assert(channels % SSM_PREFETCH == 0); + for (int g0 = 0; g0 < channels; g0 += SSM_PREFETCH) { + float sv[SSM_PREFETCH][NSTATE]; + float xvr[SSM_PREFETCH]; + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const size_t prow = (size_t)(first_channel + g0 + gi) * MB_STATE; + xvr[gi] = __bfloat162float( + xa[head * MB_HEAD_DIM + first_channel + g0 + gi]); + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = st_in[prow + lane + j * 32]; + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const int p = first_channel + g0 + gi; + const float xv = xvr[gi]; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < NSTATE; ++j) { + const int n = lane + j * 32; + const size_t at = (size_t)p * MB_STATE + n; + const float s = sv[gi][j] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_CONV_DIM; + const int c = idx - t * MB_CONV_DIM; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v; +#ifdef PF_CONV_ZERO_STATE + if (src < 0) v = 0.f; +#else + if (src < 0) v = __bfloat162float(state[pos + k]); +#endif + else v = __bfloat162float( + zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]); + acc += __bfloat162float(w[k]) * v; + } +#ifndef PF_CONV_NO_SILU + acc = acc / (1.f + __expf(-acc)); +#endif + xbc_act[(size_t)t * MB_CONV_DIM + c] = __float2bfloat16(acc); + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + const int units = n_seq * MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int seq = unit / (MB_HEADS * MB_HEAD_DIM); + const int rem = unit - seq * (MB_HEADS * MB_HEAD_DIM); + const int head = rem / MB_HEAD_DIM; + const int p = rem - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p * MB_STATE; + float sreg[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) sreg[j] = st[lane + j * 32]; + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const bf16 *xa = xbc_act + t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + const bf16 *dtrow = zxbcdt + t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + sreg[j] = sreg[j] * dA + __bfloat162float(Bsec[n]) * dt * xv; + acc += sreg[j] * __bfloat162float(Csec[n]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) st[lane + j * 32] = sreg[j]; + } + return; + } +#if PF_SSM_SPLIT_P + const int units = MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; unit += warps_total) { + const int head = unit / MB_HEAD_DIM; + const int p_only = unit - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; +#else + for (int head = blockIdx.x * WARPS + warp; head < MB_HEADS; + head += warps_total) { + const int group = head / MB_HEADS_PER_GROUP; +#endif + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *st = L.ssm_state + (size_t)slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE +#if PF_SSM_SPLIT_P + + (size_t)p_only * MB_STATE +#endif + ; + for (int t = 0; t < n_tokens; ++t) { + const bf16 *xa = xbc_act + (size_t)t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = + zxbcdt + (size_t)t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } +#if PF_SSM_SPLIT_P + { + const int p = p_only; +#else + for (int p = 0; p < MB_HEAD_DIM; ++p) { +#endif + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { +#if PF_SSM_SPLIT_P + const size_t at = (size_t)lane + j * 32; +#else + const size_t at = (size_t)p * MB_STATE + lane + j * 32; +#endif + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } + } +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32], wg[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + wg[j] = __bfloat162float(L.gnorm_w[idx]); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16(held[j] * inv * wg[j]); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len MK2_PROGRESS_PARAMETER(progress_mask)) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + constexpr int Q_PER_WARP = 2; + constexpr int GROUPS_PER_KV = 2; + const int pairs = BATCH * BB_KV_HEADS * GROUPS_PER_KV; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x + warp * gridDim.x; pair < pairs; + pair += gridDim.x * WARPS) { + const int row = pair / (BB_KV_HEADS * GROUPS_PER_KV); + MK2_PROGRESS_GUARD(progress_mask, row); + const int within_row = pair % (BB_KV_HEADS * GROUPS_PER_KV); + const int kv = within_row / GROUPS_PER_KV; + const int q_group = within_row % GROUPS_PER_KV; + const int q_first = q_group * Q_PER_WARP; + const int q_count = q_group == 0 ? Q_PER_WARP : 1; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (q_group == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[h][j] = 0.f; + } + float m[Q_PER_WARP] = {-CUDART_INF_F, -CUDART_INF_F}; + float denom[Q_PER_WARP] = {0.f, 0.f}; + constexpr int PF = 8; + const int tpb = L.tokens_per_block; + const size_t elems_per_block = + (size_t)tpb * BB_KV_HEADS * BB_HEAD_DIM; + const int *bt_k = L.block_table + (size_t)(row * 2) * L.max_blocks_per_seq; + const int *bt_v = bt_k + L.max_blocks_per_seq; + const int head_base = kv * tpb * BB_HEAD_DIM; + float qr[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const bf16 *q = base + (kv * kv_per_q + q_first + h) * BB_HEAD_DIM; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qr[h][j] = __bfloat162float(q[lane + j * 32]); + } + for (int page0 = 0; page0 < T; page0 += tpb) { + const int k_page = bt_k[page0 / tpb] & 0x7FFFFFFF; + const int v_page = bt_v[page0 / tpb] & 0x7FFFFFFF; + const bf16 *kbase = L.pool_base + + (size_t)k_page * elems_per_block + + head_base; + const bf16 *vbase = L.pool_base + + (size_t)v_page * elems_per_block + + head_base; + const int tend = page0 + tpb < T ? page0 + tpb : T; + for (int t0 = page0; t0 < tend; t0 += PF) { + const int navail = tend - t0 < PF ? tend - t0 : PF; + float kr[PF][PER_LANE], vr[PF][PER_LANE]; + #pragma unroll + for (int u = 0; u < PF; ++u) { + const int tpos = t0 + (u < navail ? u : navail - 1); +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = vbase + (tpos - page0) * BB_HEAD_DIM; +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new : kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = cur_token + ? v_new : vbase + (tpos - page0) * BB_HEAD_DIM; +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + kr[u][j] = __bfloat162float(kt[i]); + vr[u][j] = __bfloat162float(vt[i]); + } + } + #pragma unroll + for (int u = 0; u < PF; ++u) { + if (u >= navail) break; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + dot += qr[h][j] * kr[u][j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m[h], s); + const float rescale = __expf(m[h] - m_new); + const float w = __expf(s - m_new); + denom[h] = denom[h] * rescale + w; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[h][j] = acc[h][j] * rescale + w * vr[u][j]; + m[h] = m_new; + } + } + } + } + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const float inv = 1.f / denom[h]; + const int head = kv * kv_per_q + q_first + h; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[h][j] * inv); + } + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * BB_KV_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / BB_KV_HEADS; + const int kv = idx - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv); + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv); + for (int i = 0; i < BB_HEAD_DIM; ++i) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + for (int s = 0; s <= pos; ++s) { + const bf16 *kt = kv_slot(L, seq, 0, s) + kv_offset(L, s, kv); + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + const bf16 *vt = kv_slot(L, seq, 1, s) + kv_offset(L, s, kv); + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + wt * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH) { + __shared__ float score[BB_EXPERTS]; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + for (int e = warp; e < BB_EXPERTS; e += WARPS) { + const bf16 *wr = gate_w + (size_t)e * D_MODEL; + float acc = 0.f; + for (int i = lane; i < D_MODEL; i += 32) + acc += __bfloat162float(normed[(size_t)row * D_MODEL + i]) * + __bfloat162float(wr[i]); + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) score[e] = 1.f / (1.f + __expf(-acc)); + } + __syncthreads(); + if (threadIdx.x == 0) { + float remaining[BB_EXPERTS]; + int pick[BB_TOP_K]; + for (int e = 0; e < BB_EXPERTS; ++e) { + moe_w[row * BB_EXPERTS + e] = 0.f; + remaining[e] = score[e] + __bfloat162float(gate_bias[e]); + } + float total = 0.f; + for (int k = 0; k < BB_TOP_K; ++k) { + int best = 0; + for (int e = 1; e < BB_EXPERTS; ++e) + if (remaining[e] > remaining[best]) best = e; + pick[k] = best; + total += score[best]; + remaining[best] = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + for (int k = 0; k < BB_TOP_K; ++k) + moe_w[row * BB_EXPERTS + pick[k]] = score[pick[k]] * norm; + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert <= BB_EXPERTS; ++expert) + value = __fadd_rn(value, parts[base + (size_t)expert * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ bf16 bb_moe_fold_selected( + const float *parts, const float *weights, int idx) { + const int row = idx / D_MODEL; + const size_t base = (size_t)row * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert < BB_EXPERTS; ++expert) { + const float part = weights[row * BB_EXPERTS + expert] != 0.f + ? parts[base + (size_t)expert * D_MODEL] : 0.f; + value = __fadd_rn(value, part); + } + value = __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ __forceinline__ void sample_argmax_merge( + float &value, int &rank, float other, int other_rank) { + if (other > value || (other == value && other_rank < rank)) { + value = other; + rank = other_rank; + } +} +__device__ __forceinline__ int sample_block_argmax(float value, int rank) { + __shared__ float warp_value[WARPS]; + __shared__ int warp_rank[WARPS]; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + if (lane == 0) { warp_value[warp] = value; warp_rank[warp] = rank; } + __syncthreads(); + if (warp == 0) { + value = lane < WARPS ? warp_value[lane] : -CUDART_INF_F; + rank = lane < WARPS ? warp_rank[lane] : INT_MAX; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + } + return rank; +} +__device__ __forceinline__ int sample_eos_insertion_rank( + const uint64_t *keys, float eos) { + int lo = 0, hi = TOP_K; + while (lo < hi) { + const int mid = (lo + hi) >> 1; + if (eos > ordered_to_float((uint32_t)(keys[mid] >> 32))) hi = mid; + else lo = mid + 1; + } + return lo; +} +__device__ __forceinline__ void sample_topk_serial( + const uint64_t *keys, float eos, const float *noise, int *out) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float value; + int code; + if (!eos_used && eos > av) { + value = eos; + code = AUDIO_EOS; + eos_used = true; + } else { + value = av; + code = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = value + noise[j]; + if (noised > best_noised) { + best_noised = noised; + best_code = code; + } + } + *out = best_code; +} +__device__ __forceinline__ void sample_topk_parallel( + const uint64_t *keys, float eos, const float *noise, int *out) { + const int j = threadIdx.x; + const bool sorted_nan = j < TOP_K && + isnan(ordered_to_float((uint32_t)(keys[j] >> 32))); + if (__syncthreads_or(sorted_nan)) { + if (threadIdx.x == 0) sample_topk_serial(keys, eos, noise, out); + return; + } + const int lane = threadIdx.x & 31; + int insertion = lane == 0 ? sample_eos_insertion_rank(keys, eos) : 0; + insertion = __shfl_sync(0xffffffffu, insertion, 0); + float value = -CUDART_INF_F; + int rank = INT_MAX; + if (j < TOP_K) { + const int a = j - (j > insertion); + const float v = j == insertion + ? eos : ordered_to_float((uint32_t)(keys[a] >> 32)); + const float noised = v + noise[j]; + if (noised > -CUDART_INF_F) { value = noised; rank = j; } + } + const int winner = sample_block_argmax(value, rank); + if (threadIdx.x == 0) { + if (winner == INT_MAX) *out = 0; + else if (winner == insertion) *out = AUDIO_EOS; + else { + const int a = winner - (winner > insertion); + *out = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + } + } +} +constexpr int SORT_ITEMS_PER_THREAD = SORT_N / THREADS; +struct sampler_storage { + uint64_t keys[2][SORT_N]; +}; +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + sampler_storage &storage, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + uint64_t thread_keys[SORT_ITEMS_PER_THREAD]; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + thread_keys[item] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + int cross_stage = 0; + for (int k = 2; k <= SORT_N; k <<= 1) { + for (int j = k >> 1; j > 0; j >>= 1) { + if (j >= 32) { + const int bank = cross_stage++ & 1; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + storage.keys[bank][n] = thread_keys[item]; + } + __syncthreads(); + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + const uint64_t other = storage.keys[bank][n ^ j]; + const bool take_hi = ((n & k) == 0) == ((n & j) == 0); + const uint64_t mine = thread_keys[item]; + thread_keys[item] = take_hi + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } else { + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + const uint64_t other = __shfl_xor_sync( + 0xffffffffu, thread_keys[item], j); + const bool take_hi = ((n & k) == 0) == ((n & j) == 0); + const uint64_t mine = thread_keys[item]; + thread_keys[item] = take_hi + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } + } + } + const int top_bank = cross_stage & 1; + if (threadIdx.x < TOP_K) + storage.keys[top_bank][threadIdx.x] = thread_keys[0]; + __syncthreads(); + sample_topk_parallel( + storage.keys[top_bank], eos_v, + gumbel + ((size_t)row * N_POS + pos) * TOP_K, + out_codes + row); + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES], act_tile (&a_s)[STAGES], out_tile &o_s, + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + int &tile_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0) { + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + const int iters = k_total / KCHUNK; + int tile_start = (int)blockIdx.x - tile_io % (int)gridDim.x; + if (tile_start < 0) tile_start += (int)gridDim.x; + tile_io += n_tiles; + int tiles_here = 0; + for (int t = tile_start; t < n_tiles; t += gridDim.x) ++tiles_here; + const int iter_base = iter_io; + iter_io += tiles_here * iters * n_tok; + const int my_tiles = (n_tiles - tile_start + gridDim.x - 1) / gridDim.x; + if (warpgroup::groupid() == 1) { + if (warpgroup::laneid() == 0) { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) { + const int tile_id = mk2_expert_tile(work, active_mask, compact_tiles); + for (int tok = 0; tok < n_tok; ++tok) + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(drained[s], phase); + tma::expect_bytes(filled[s], + sizeof(weight_tile) + sizeof(act_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s], w, {tile_id, i}, filled[s]); + const int a_col = out_tiles_per_expert + ? (tile_id / out_tiles_per_expert) * act_k_tiles_per_expert + i + : i; + tma::load_async<2, cache_policy::EVICT_LAST>( + a_s[s], a, {tok, a_col}, filled[s]); + } + } + } + } else { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) + for (int tok = 0; tok < n_tok; ++tok) { + const int tile_id = mk2_expert_tile(work, active_mask, compact_tiles); + rt_fl<16, BATCH_PAD> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s], a_s[s]); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s]); + } + const int base = tile_id * OUT_ATOM; + const int grow0 = tok * BATCH_PAD; + if (epi == EPI_BIAS) { + warpgroup::store(o_sf, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + dest[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(o_sf[int2{col, row}] + + __bfloat162float(bias[base + col])); + } + warpgroup::sync(1); + continue; + } + warpgroup::store(o_s, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_s[int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_w[grow * BB_EXPERTS + expert]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * + D_MODEL + local] = __fmul_rn(wt, v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1); + } + (void)my_tiles; + } +} +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr, + const unsigned char *native_predictor_active_mask = nullptr + MK2_PROGRESS_PARAMETER(native_progress_mask)) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&w_s)[STAGES] = alloc.allocate(); + act_tile (&a_s)[STAGES] = alloc.allocate(); + out_tile (&o_s) = alloc.allocate(); + out_tile_f (&o_sf) = alloc.allocate(); + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * WARPS + 2]; + __shared__ sampler_storage sort_storage; + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS, 0); + for (int c = 0; c < CONSUMER_WARPS; ++c) arrive(drained[s]); + } + } + int gemm_iter = 0; + int gemm_tile = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + mamba_gated_norm(L, bb->zxbcdt, bb->gated); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w); + } + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + } + grid_barrier(g.barrier); + gemm(L.exp_up, bb->normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, nullptr, 1, BATCH, + bb->moe_w, BB_MOE_INTER / OUT_ATOM); + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + 1, BATCH, bb->moe_w, D_MODEL / OUT_ATOM); + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_selected(bb->moe_out, bb->moe_w, idx); + grid_barrier(g.barrier); + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int n_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = n_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + bb->pf_h[idx] = bb->prompt[idx]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = bb->pf_n_seq > 0 ? bb->pf_n_seq : 1; + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_gated, + T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, T); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * (BB_EXPERTS + 1) * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->pf_normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->pf_moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = bb_moe_fold_deterministic(bb->pf_moe_out, idx); + grid_barrier(g.barrier); + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, bb->pf_final, ln_scratch, + n_total); + grid_barrier(g.barrier); + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, bb->phoneme_b); + grid_barrier(g.barrier); + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + MK2_PROGRESS_GUARD(native_progress_mask, row); + if (threadIdx.x == 0) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = 0; n < BB_PHONEME; ++n) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + bb->phoneme_out[row] = arg; + } + __syncthreads(); + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + __shared__ int native_run_predictor; + if (native_predictor_active_mask && bb && g.native_codes) { + if (threadIdx.x == 0) { + int any_active = 0; + for (int row = 0; row < BATCH; ++row) + any_active |= int(native_predictor_active_mask[row] != 0 && + MK2_PROGRESS_ENABLED(native_progress_mask, row)); + native_run_predictor = any_active; + } + __syncthreads(); + } + const bool skip_native_predictor = native_predictor_active_mask && bb && + g.native_codes && native_run_predictor == 0; + if (skip_native_predictor) { + for (int index = blockIdx.x * THREADS + threadIdx.x; + index < BATCH * N_POS; index += gridDim.x * THREADS) { + MK2_PROGRESS_GUARD(native_progress_mask, index / N_POS); + g.sampled[index] = 0; + g.native_codes[index] = 0; + } + grid_barrier(g.barrier); + } else { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + constexpr int DPL = D_HEAD / 32; + float qv[DPL], kn[DPL], vn[DPL]; + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + qv[j] = __bfloat162float(q[i]); + kn[j] = __bfloat162float(k_new[i]); + vn[j] = __bfloat162float(v_new[i]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += qv[j] * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + { + float acc = 0.f; + #pragma unroll + for (int j = 0; j < DPL; ++j) acc += qv[j] * kn[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[pos] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + float ov[DPL] = {}; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *vt = g.vcache + cache_base + t * D_HEAD; + const float s = sc[t]; + #pragma unroll + for (int j = 0; j < DPL; ++j) + ov[j] += s * __bfloat162float(vt[lane + j * 32]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + ov[j] += sc[pos] * vn[j]; + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + + lane + j * 32] = __float2bfloat16(ov[j]); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, + g.sampled_pos, sort_storage, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif +#ifdef ABLATE_SAMPLE + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { + __syncthreads(); + } else { + grid_barrier(g.barrier); + } +#endif + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x == 0) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + g.sampled[row * N_POS + pos] = code; + if (g.native_codes && MK2_PROGRESS_ENABLED(native_progress_mask, row)) { +#if MK2_INFLIGHT + g.native_codes[row * N_POS + pos] = + native_predictor_active_mask && !native_predictor_active_mask[row] + ? 0 : (long long)code; +#else + g.native_codes[row * N_POS + pos] = (long long)code; +#endif + } + } + } +#ifdef ABLATE_SAMPLE + grid_barrier(g.barrier); +#else + if (pos + 1 == n_pos) grid_barrier(g.barrier); +#endif + } + } + if (g.native_seq_len && blockIdx.x == 0 && threadIdx.x < BATCH && + MK2_PROGRESS_ENABLED(native_progress_mask, threadIdx.x)) + g.native_seq_len[threadIdx.x] += 1; + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 64 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * (sizeof(weight_tile) + sizeof(act_tile)) + + sizeof(out_tile) + sizeof(out_tile_f) + 4096; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; + const unsigned char *native_predictor_active_mask = nullptr; + MK2_PROGRESS_MEMBER + int maximum_grid = 0; + bool grid_reset_pending = false; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * (sizeof(weight_tile) + sizeof(act_tile)) + + sizeof(out_tile) + sizeof(out_tile_f) + 4096), 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = std::min(out->sms * out->blocks_per_sm, 4096); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 64 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->maximum_grid = info.sms * info.blocks_per_sm; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.maximum_grid = pred->maximum_grid; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(s, &capture_status)); + if (st->pred.grid_reset_pending || capture_status != cudaStreamCaptureStatusNone) { + mk2_check(cudaMemsetAsync(st->pred.g.barrier, 0, 64 * sizeof(unsigned int), s)); + st->pred.grid_reset_pending = false; + } + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS, nullptr, + native ? st->pred.native_predictor_active_mask : nullptr + MK2_PROGRESS_ARGUMENT(native ? st->pred.native_progress_mask : nullptr)); + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" int mk2_set_native_predictor_active_mask(MK2Predictor *p, + const unsigned char *mask) { + try { + if (!p || reinterpret_cast(p)->kind != 1) + throw std::runtime_error("startup mask requires a full Stage0 handle"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && st->pred.native_predictor_active_mask != mask) + throw std::runtime_error("startup mask pointer cannot change after first Stage0 enqueue"); + if (mask) + mk2_stage0_device_pointer(mask, st->pred.device, "native_predictor_active_mask"); + st->pred.native_predictor_active_mask = mask; + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown startup mask binding failure"; return 1; } +} +extern "C" int mk2_native_startup_skip_abi_version() { return 1; } +#include "../../inflight/runtime.cuh" +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "native_prefill_impl.inc" +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *p, int grid) { + if (!p) { g_last_error = "grid setter requires full Stage0 handle"; return 1; } + auto *st = reinterpret_cast(p); + if (st->kind != 1 || grid < 1 || grid > st->maximum_grid) { + g_last_error = "requested Stage0 grid outside independently resident capacity"; + return 1; + } + if (st->grid != grid) st->grid_reset_pending = true; + st->grid = grid; + g_last_error.clear(); + return 0; +} diff --git a/deploy/nemotron_native/kernels/stage0/small/mk2_abi.cuh b/deploy/nemotron_native/kernels/stage0/small/mk2_abi.cuh new file mode 100644 index 0000000..bf1e415 --- /dev/null +++ b/deploy/nemotron_native/kernels/stage0/small/mk2_abi.cuh @@ -0,0 +1,85 @@ +#pragma once +#include +extern "C" { +struct MK2Predictor; +struct MK2Weights { + const void *qkv[3]; + const void *o[3]; + const void *ff_in[3]; + const void *ff_out[3]; + const void *norm_self[3]; + const void *norm_ff[3]; + const void *pos_emb; + const void *audio_emb[16]; + const void *head_w[16]; + const void *head_b[16]; +}; +MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch); +MK2Predictor *mk2_create(const char *raw_dir, int batch); +int mk2_run(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, void *codes_int64, void *stream); +int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, const void *gumbel_fp32, + float temperature, const void *audio_valid_bool, + void *codes_int64, void *stream); +void mk2_destroy(MK2Predictor *p); +struct MK2BuildInfo { + int version, batch, batch_pad, stages, kchunk; + int direct_input, b1_local_barriers, skip_hidden_snapshot; + int threads, dynamic_smem, static_smem, registers; + int device, sms, blocks_per_sm, grid; +}; +int mk2_build_info(MK2BuildInfo *out); +struct MK2Stage0Bindings { + void *conv_state[14]; + int conv_slot_stride[14]; + void *ssm_state[14]; + int ssm_slot_stride[14]; + const void *slot_map; + void *kv_pool_base[3]; + const void *kv_block_table[3]; + int kv_max_blocks_per_seq; + int kv_tokens_per_block; + const void *seq_len; + const void *embedding; + const void *gumbel; + float temperature; + const void *moe_w_inject_debug; + void *codes_out; + void *phoneme_out; +}; +struct MK2BackboneWeights { + const void *norm[31]; + const void *norm_f; + const void *mb_in_proj[14]; + const void *mb_out_proj[14]; + const void *mb_conv_w[14]; + const void *mb_conv_b[14]; + const void *mb_A[14]; + const void *mb_D[14]; + const void *mb_dt_bias[14]; + const void *mb_gnorm[14]; + const void *at_qkv[3]; + const void *at_o[3]; + const void *moe_exp_up[14]; + const void *moe_exp_down[14]; + const void *moe_sh_up[14]; + const void *moe_sh_down[14]; + const void *moe_gate[14]; + const void *moe_gate_bias[14]; + const void *phoneme_w; + const void *phoneme_b; + const char *pattern; +}; +MK2Predictor *mk2_create_stage0(const MK2Weights *predictor, + const MK2BackboneWeights *backbone, int batch); +int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, void *stream); +int mk2_set_native_predictor_active_mask(MK2Predictor *p, const unsigned char *mask); +int mk2_native_startup_skip_abi_version(); +unsigned long long mk2_stage0_launches(MK2Predictor *p); +int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out); +int mk2_batch(const MK2Predictor *p); +int mk2_positions(const MK2Predictor *p); +const char *mk2_last_error(void); +} +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *, int grid); diff --git a/deploy/nemotron_native/kernels/stage0/small/native_prefill_abi.cuh b/deploy/nemotron_native/kernels/stage0/small/native_prefill_abi.cuh new file mode 100644 index 0000000..3731e5f --- /dev/null +++ b/deploy/nemotron_native/kernels/stage0/small/native_prefill_abi.cuh @@ -0,0 +1,33 @@ +#pragma once +#include "mk2_abi.cuh" +#include +extern "C" { +struct MK2NativePrefill; +struct MK2NativePrefillConfig { + int batch; + int prompt_positions; + int max_sequence_tokens; + int tokens_per_block; +}; +struct MK2NativePrefillInfo { + int version; + int batch; + int prompt_positions; + int total_prompt_rows; + int padded_prompt_rows; + int max_sequence_tokens; + int tokens_per_block; + int blocks_per_sequence; + int conv_slot_elements; + int ssm_slot_elements; + std::uint64_t kv_pool_bytes_per_layer; + std::uint64_t prefill_scratch_bytes; +}; +MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *predictor, const MK2BackboneWeights *backbone, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden); +int mk2_native_prefill_run(MK2NativePrefill *p, void *stream); +int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info); +void mk2_native_prefill_destroy(MK2NativePrefill *p); +} diff --git a/deploy/nemotron_native/kernels/stage0/small/native_prefill_impl.inc b/deploy/nemotron_native/kernels/stage0/small/native_prefill_impl.inc new file mode 100644 index 0000000..dae1557 --- /dev/null +++ b/deploy/nemotron_native/kernels/stage0/small/native_prefill_impl.inc @@ -0,0 +1,198 @@ +#include "native_prefill_abi.cuh" +struct MK2NativePrefill { + MK2Predictor *base; + MK2Stage0Bindings state; + MK2NativePrefillInfo info; +}; +namespace { +struct NativePrefillBaseDeleter { + void operator()(MK2Predictor *p) const noexcept { mk2_destroy(p); } +}; +static void native_prefill_check_config(const MK2NativePrefillConfig &c) { + if (c.batch != BATCH || !(BATCH == 1 || BATCH == 2 || BATCH == 4 || BATCH == 8 || BATCH == 16 || BATCH == 32 || BATCH == 64)) + throw std::runtime_error("native prefill requires a matching supported physical batch"); + if (c.prompt_positions <= 0 || c.max_sequence_tokens <= c.prompt_positions || + c.max_sequence_tokens > INT_MAX - 31 || c.tokens_per_block != 32) + throw std::runtime_error("native prefill requires positive prompt, decode capacity and 32-token pages"); + const std::uint64_t rows = std::uint64_t(BATCH) * c.prompt_positions; + const std::uint64_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + if (padded > INT_MAX / ((BB_EXPERTS + 1) * D_MODEL)) + throw std::runtime_error("native prefill token rows exceed safe device loop bounds"); + const std::uint64_t pages = (std::uint64_t(c.max_sequence_tokens) + 31) / 32; + if (pages * BATCH * 2 > INT_MAX) + throw std::runtime_error("native prefill primary page indices exceed int32"); +#if PF_SKIP_MAMBA || PF_SKIP_ATTN || PF_SKIP_MOE || BB_SKIP_MAMBA || BB_SKIP_ATTN || BB_SKIP_MOE || BB_ATTN_SKIP_KV_WRITE || BB_ATTN_LEGACY_CUR_KV + throw std::runtime_error("native prefill refuses arithmetic ablation builds"); +#endif +#if defined(PF_CONV_ZERO_STATE) || defined(PF_CONV_NO_SILU) || defined(PF_CONV_STATE_NO_SHIFT) + throw std::runtime_error("native prefill refuses convolution ablation builds"); +#endif +} +} +extern "C" MK2NativePrefill *mk2_native_prefill_create( + const MK2Weights *pw, const MK2BackboneWeights *bw, + const MK2NativePrefillConfig *config, const void *prompt, void *final_hidden) { + try { + if (!config || !pw || !bw) throw std::runtime_error("null native prefill config/weights"); + native_prefill_check_config(*config); + int device = -1; + mk2_check(cudaGetDevice(&device)); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(nullptr, &capture)); + if (capture != cudaStreamCaptureStatusNone) + throw std::runtime_error("native prefill construction must occur outside capture"); + mk2_stage0_device_pointer(prompt, device, "native prompt"); + mk2_stage0_device_pointer(final_hidden, device, "native final_hidden"); + const size_t rows = size_t(BATCH) * config->prompt_positions; + const size_t padded = ((rows + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; + const size_t output_bytes = rows * D_MODEL * sizeof(bf16); + const auto input_address = reinterpret_cast(prompt); + const auto output_address = reinterpret_cast(final_hidden); + const auto distance = input_address < output_address + ? output_address - input_address : input_address - output_address; + if (distance < output_bytes) + throw std::runtime_error("native prompt/final_hidden buffers must not overlap"); + std::unique_ptr base( + mk2_create_stage0(pw, bw, BATCH)); + if (!base) return nullptr; + auto *st = reinterpret_cast(base.get()); + MK2Allocations owned; + MK2Stage0Bindings binding{}; + const int conv_elements = MB_CONV_DIM * (MB_CONV_K - 1); + const int ssm_elements = MB_HEADS * MB_HEAD_DIM * MB_STATE; + const int blocks = (config->max_sequence_tokens + 31) / 32; + const size_t page_elements = size_t(32) * BB_KV_HEADS * BB_HEAD_DIM; + const size_t pool_bytes = size_t(BATCH) * 2 * blocks * page_elements * sizeof(bf16); + binding.kv_max_blocks_per_seq = blocks; + binding.kv_tokens_per_block = 32; + for (int i = 0; i < 14; ++i) { + bf16 *conv = nullptr; + float *ssm = nullptr; + mk2_check(owned.alloc(&conv, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(owned.alloc(&ssm, size_t(BATCH) * ssm_elements * sizeof(float))); + mk2_check(cudaMemset(conv, 0, size_t(BATCH) * conv_elements * sizeof(bf16))); + mk2_check(cudaMemset(ssm, 0, size_t(BATCH) * ssm_elements * sizeof(float))); + binding.conv_state[i] = conv; + binding.conv_slot_stride[i] = conv_elements; + binding.ssm_state[i] = ssm; + binding.ssm_slot_stride[i] = ssm_elements; + st->mb_host[i].conv_state = conv; + st->mb_host[i].conv_slot_stride = conv_elements; + st->mb_host[i].ssm_state = ssm; + st->mb_host[i].ssm_slot_stride = ssm_elements; + } + std::vector table(size_t(BATCH) * 2 * blocks); + for (size_t i = 0; i < table.size(); ++i) table[i] = static_cast(i); + for (int i = 0; i < 3; ++i) { + bf16 *pool = nullptr; + int *page_table = nullptr; + mk2_check(owned.alloc(&pool, pool_bytes)); + mk2_check(cudaMemset(pool, 0, pool_bytes)); + mk2_check(owned.alloc(&page_table, table.size() * sizeof(int))); + mk2_check(cudaMemcpy(page_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + binding.kv_pool_base[i] = pool; + binding.kv_block_table[i] = page_table; + st->at_host[i].pool_base = pool; + st->at_host[i].block_table = page_table; + st->at_host[i].max_blocks_per_seq = blocks; + st->at_host[i].tokens_per_block = 32; + } + std::vector slots(BATCH), lengths(BATCH, config->prompt_positions); + for (int i = 0; i < BATCH; ++i) slots[i] = i; + int *device_slots = nullptr, *device_lengths = nullptr; + mk2_check(owned.alloc(&device_slots, BATCH * sizeof(int))); + mk2_check(owned.alloc(&device_lengths, BATCH * sizeof(int))); + mk2_check(cudaMemcpy(device_slots, slots.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(device_lengths, lengths.data(), BATCH * sizeof(int), cudaMemcpyHostToDevice)); + binding.slot_map = device_slots; + binding.seq_len = device_lengths; + Backbone &hb = st->host_bb; + size_t scratch_bytes = 0; + auto bf = [&](bf16 **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(bf16); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + auto fp = [&](float **p, size_t columns) { + const size_t bytes = padded * columns * sizeof(float); + mk2_check(owned.alloc(p, bytes)); + mk2_check(cudaMemset(*p, 0, bytes)); + scratch_bytes += bytes; + }; + bf(&hb.pf_h, D_MODEL); bf(&hb.pf_r, D_MODEL); bf(&hb.pf_normed, D_MODEL); + bf(&hb.pf_zxbcdt, MB_IN_PROJ); bf(&hb.pf_xbc_act, MB_CONV_DIM); + bf(&hb.pf_gated, MB_INNER); bf(&hb.pf_qkv, BB_QKV); + bf(&hb.pf_moe_h, BB_EXPERTS * BB_MOE_INTER); + bf(&hb.pf_moe_sh, BB_SHARED_INTER); + fp(&hb.pf_moe_out, (BB_EXPERTS + 1) * D_MODEL); + fp(&hb.pf_moe_w, BB_EXPERTS); + hb.pf_final = static_cast(final_hidden); + hb.pf_normed_v = act_gl{hb.pf_normed, nullptr, nullptr, padded, size_t(D_MODEL)}; + hb.pf_gated_v = act_gl{hb.pf_gated, nullptr, nullptr, padded, size_t(MB_INNER)}; + hb.pf_qkv_v = act_gl{hb.pf_qkv, nullptr, nullptr, padded, size_t(BB_QKV)}; + hb.pf_moe_h_v = act_gl{hb.pf_moe_h, nullptr, nullptr, padded, + size_t(BB_EXPERTS * BB_MOE_INTER)}; + hb.pf_moe_sh_v = act_gl{hb.pf_moe_sh, nullptr, nullptr, padded, size_t(BB_SHARED_INTER)}; + hb.n_prefill_tokens = static_cast(rows); + hb.prefill_slot = 0; + hb.pf_n_seq = BATCH; + hb.pf_fuse_decode = 0; + hb.prompt = static_cast(prompt); + hb.pf_seq_slot = device_slots; + hb.slot_map = device_slots; + hb.seq_len = device_lengths; + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + hb.native_prefill_reset = 1; + mk2_check(cudaMemcpy(const_cast(hb.mamba), st->mb_host.data(), + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(const_cast(hb.attn), st->at_host.data(), + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice)); + mk2_check(cudaMemcpy(st->d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + MK2NativePrefillInfo info{1, BATCH, config->prompt_positions, + static_cast(rows), static_cast(padded), config->max_sequence_tokens, + 32, blocks, conv_elements, ssm_elements, pool_bytes, scratch_bytes}; + std::unique_ptr handle(new MK2NativePrefill{base.get(), binding, info}); + st->owned.reserve(st->owned.size() + owned.ptrs.size()); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + base.release(); + g_last_error.clear(); + return handle.release(); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown native prefill construction failure"; return nullptr; } +} +extern "C" int mk2_native_prefill_run(MK2NativePrefill *p, void *stream) { + try { + if (!p || !p->base) throw std::runtime_error("null native prefill handle"); + auto *st = reinterpret_cast(p->base); + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != st->pred.device || st->pred.batch != BATCH || + st->pred.kind != 1 || !st->host_bb.native_prefill_reset) + throw std::runtime_error("native prefill handle device or mode mismatch"); + predictor_mk2<<pred.grid, THREADS, st->pred.smem, + static_cast(stream)>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS); + mk2_check(cudaGetLastError()); + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown native prefill enqueue failure"; return 1; } +} +extern "C" int mk2_native_prefill_get_state( + const MK2NativePrefill *p, MK2Stage0Bindings *out, MK2NativePrefillInfo *info) { + if (!p || !out) { g_last_error = "null native prefill handle/state output"; return 1; } + *out = p->state; + if (info) *info = p->info; + g_last_error.clear(); + return 0; +} +extern "C" void mk2_native_prefill_destroy(MK2NativePrefill *p) { + if (!p) return; + mk2_destroy(p->base); + delete p; +} diff --git a/deploy/nemotron_native/kernels/stage0/small/predictor_mk2.cu b/deploy/nemotron_native/kernels/stage0/small/predictor_mk2.cu new file mode 100644 index 0000000..5689f8a --- /dev/null +++ b/deploy/nemotron_native/kernels/stage0/small/predictor_mk2.cu @@ -0,0 +1,4074 @@ +#include "../../inflight/progress.cuh" +#if MK2_INFLIGHT +#include +#endif +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +using namespace kittens; +constexpr int D_MODEL = 1536; +constexpr int D_FFN = 6144; +constexpr int N_HEADS = 12; +constexpr int D_HEAD = 128; +constexpr int N_LAYERS = 3; +constexpr int N_POS = 16; +#ifndef MK2_BATCH +#define MK2_BATCH 32 +#endif +constexpr int BATCH = MK2_BATCH; +#ifndef MK2_ACTIVE_EXPERTS +#define MK2_ACTIVE_EXPERTS 0 +#endif +#ifndef MK2_SSM_SPLITS +#define MK2_SSM_SPLITS 1 +#endif +static_assert(MK2_SSM_SPLITS > 0 && 24 % MK2_SSM_SPLITS == 0); +#ifndef MK2_DIRECT_INPUT +#define MK2_DIRECT_INPUT 0 +#endif +#ifndef MK2_B1_LOCAL_BARRIERS +#define MK2_B1_LOCAL_BARRIERS 0 +#endif +#ifndef MK2_SKIP_HIDDEN_SNAPSHOT +#define MK2_SKIP_HIDDEN_SNAPSHOT 0 +#endif +constexpr int BATCH_PAD = ((BATCH + 15) / 16) * 16; +static_assert(BATCH >= 1 && BATCH <= BATCH_PAD); +static_assert(BATCH == 2 || BATCH == 4 || BATCH == 8); +static_assert(BATCH % 16 != 0 || BATCH_PAD == BATCH); +constexpr int CODEBOOK = 1032; +constexpr int VOCAB_PADDED = 1088; +constexpr int TOP_K = 80; +constexpr int BB_LAYERS = 31; +constexpr float RMS_EPS = 1e-5f; +constexpr int MB_HEADS = 64; +constexpr int MB_HEAD_DIM = 24; +constexpr int MB_STATE = 128; +constexpr int MB_GROUPS = 8; +constexpr int MB_CONV_K = 4; +constexpr int MB_INNER = MB_HEADS * MB_HEAD_DIM; +constexpr int MB_CONV_DIM = MB_INNER + 2 * MB_GROUPS * MB_STATE; +constexpr int MB_IN_PROJ = 2 * MB_INNER + 2 * MB_GROUPS * MB_STATE + MB_HEADS; +constexpr int MB_GROUP_SIZE = MB_INNER / MB_GROUPS; +constexpr int MB_HEADS_PER_GROUP = MB_HEADS / MB_GROUPS; +constexpr int BB_ATT_HEADS = 12; +constexpr int BB_KV_HEADS = 4; +constexpr int BB_HEAD_DIM = 128; +constexpr int BB_Q_SIZE = BB_ATT_HEADS * BB_HEAD_DIM; +constexpr int BB_KV_SIZE = BB_KV_HEADS * BB_HEAD_DIM; +constexpr int BB_QKV = BB_Q_SIZE + 2 * BB_KV_SIZE; +constexpr int BB_T_MAX = 8; +constexpr int BB_EXPERTS = 24; +constexpr int BB_TOP_K = 4; +constexpr int BB_MOE_INTER = 768; +constexpr int BB_SHARED_INTER = 2048; +constexpr float BB_ROUTED_SCALE = 2.5f; +constexpr int BB_PHONEME = 2051; +constexpr int BB_PHON_PADDED = 2112; +#ifndef MK2_PREFILL_MAX_TOKENS +#define MK2_PREFILL_MAX_TOKENS 4288 +#endif +constexpr int PREFILL_MAX_TOKENS = MK2_PREFILL_MAX_TOKENS; +constexpr int PREFILL_STORAGE_TOKENS = + ((PREFILL_MAX_TOKENS + BATCH_PAD - 1) / BATCH_PAD) * BATCH_PAD; +#ifndef PF_SSM_SPLIT_P +#define PF_SSM_SPLIT_P 1 +#endif +#ifndef PF_SKIP_MAMBA +#define PF_SKIP_MAMBA 0 +#endif +#ifndef PF_SKIP_ATTN +#define PF_SKIP_ATTN 0 +#endif +#ifndef PF_SKIP_MOE +#define PF_SKIP_MOE 0 +#endif +constexpr int SORT_N = 1024; +constexpr int AUDIO_EOS = 1025; +constexpr int LN_EPS_NUM = 1; +constexpr float LN_EPS = 1e-5f; +constexpr int OUT_ATOM = 64; +#ifndef KCHUNK +#define KCHUNK 256 +#endif +#ifdef ABLATE_BB_MAMBA +#define BB_SKIP_MAMBA 1 +#else +#define BB_SKIP_MAMBA 0 +#endif +#ifdef ABLATE_BB_MOE +#define BB_SKIP_MOE 1 +#else +#define BB_SKIP_MOE 0 +#endif +#ifdef ABLATE_BB_ATTN +#define BB_SKIP_ATTN 1 +#else +#define BB_SKIP_ATTN 0 +#endif +#ifndef BB_ATTN_SKIP_KV_WRITE +#define BB_ATTN_SKIP_KV_WRITE 0 +#endif +#ifndef BB_ATTN_LEGACY_CUR_KV +#define BB_ATTN_LEGACY_CUR_KV 0 +#endif +#ifndef STAGES +#define STAGES 3 +#endif +constexpr int THREADS = 256; +constexpr int CONSUMER_WARPS = 4; +constexpr int WARPS = THREADS / 32; +using weight_tile = st_bf; +using act_tile = st_bf; +using out_tile = st_bf; +using out_tile_f = st; +using weight_gl = gl; +using act_gl = gl; +enum Epilogue { EPI_STORE, EPI_RESIDUAL, EPI_GELU, EPI_BIAS, + EPI_SILU, EPI_MOE_ACC, EPI_ACC_F32 }; +struct MambaLayer { + weight_gl in_proj; + weight_gl out_proj; + const bf16 *conv_w; + const bf16 *conv_b; + const float *A; + const float *D; + const float *dt_bias; + const bf16 *gnorm_w; + bf16 *conv_state; + int conv_slot_stride; + float *ssm_state; + int ssm_slot_stride; +}; +struct Globals { + weight_gl qkv[N_LAYERS], o[N_LAYERS], ff_in[N_LAYERS], ff_out[N_LAYERS]; + weight_gl head[N_POS]; + act_gl x_v; + act_gl normed_v; + act_gl attn_v; + act_gl ffh_v; + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed; + bf16 *logits; + const bf16 *head_bias[N_POS]; + const int8_t *forbidden; + const float *gumbel; + int *sampled; + int *sampled_pos; + float temperature; + const float *norm_self[N_LAYERS], *norm_ff[N_LAYERS]; + const bf16 *pos_emb; + const bf16 *audio_emb[N_POS]; + bf16 *kcache, *vcache, *hidden_rows; + const int *codes; + unsigned int *barrier; + const bf16 *native_embedding = nullptr; + long long *native_codes = nullptr; + int *native_seq_len = nullptr; +}; +struct AttnLayer { + weight_gl qkv; + weight_gl o_proj; + bf16 *pool_base; + const int *block_table; + int max_blocks_per_seq; + int tokens_per_block; +}; +__device__ __forceinline__ bf16 *kv_slot( + const AttnLayer &L, int seq, int kv_sel, int pos) { + const int *row = L.block_table + + ((size_t)seq * 2 + kv_sel) * L.max_blocks_per_seq; + const int blk = row[pos / L.tokens_per_block] & 0x7FFFFFFF; + const size_t elems_per_block = + (size_t)L.tokens_per_block * BB_KV_HEADS * BB_HEAD_DIM; + return L.pool_base + (size_t)blk * elems_per_block; +} +__device__ __forceinline__ int kv_offset(const AttnLayer &L, int pos, int head) { + return head * L.tokens_per_block * BB_HEAD_DIM + + (pos % L.tokens_per_block) * BB_HEAD_DIM; +} +struct MoeLayer { + weight_gl exp_up; + weight_gl exp_down; + weight_gl sh_up; + weight_gl sh_down; + const bf16 *gate_w; + const bf16 *gate_bias; +}; +#ifndef MK2_SPARSE_MOE +#define MK2_SPARSE_MOE 0 +#endif +constexpr int SPARSE_N = 16; +using sparse_act_tile = st_bf; +using sparse_out_tile = st_bf; +using sparse_act_gl = gl; +struct SparseMoe { + bf16 *gathered; + int *rows; + int *counts; + sparse_act_gl gathered_v; + sparse_act_gl hidden_v; +}; +struct Backbone { + act_gl normed_v; + act_gl gated_v; + const MambaLayer *mamba; + const bf16 *norm_f_w; + const AttnLayer *attn; + const MoeLayer *moe; + const bf16 *norm_w[BB_LAYERS]; + int mamba_of_layer[BB_LAYERS]; + int attn_of_layer[BB_LAYERS]; + int moe_of_layer[BB_LAYERS]; + int past_len; + bf16 *h; + bf16 *r; + bf16 *normed; + bf16 *zxbcdt; + bf16 *gated; + bf16 *layer_out; + bf16 *qkv; + bf16 *moe_h; + bf16 *moe_sh; + float *moe_out; + float *moe_w; + const float *moe_w_inject; + float *moe_w_audit; + bf16 *dec_snapshot; + bf16 *xbc_act; + const int *slot_map; + int n_prefill_tokens; + int prefill_slot; + int pf_n_seq; + int pf_fuse_decode; + const bf16 *prompt; + const int *pf_seq_slot; + bf16 *pf_h, *pf_r, *pf_normed; + bf16 *pf_zxbcdt; + bf16 *pf_xbc_act; + bf16 *pf_gated; + bf16 *pf_qkv; + bf16 *pf_moe_h, *pf_moe_sh; + float *pf_moe_out, *pf_moe_w; + bf16 *pf_layer_out; + bf16 *pf_final; + act_gl pf_normed_v, pf_gated_v, pf_qkv_v, pf_moe_h_v, pf_moe_sh_v; + weight_gl phoneme_w; + const bf16 *phoneme_b; + bf16 *phoneme_logits; + int *phoneme_out; + unsigned long long *launches; + const int *seq_len; + act_gl qkv_v; + act_gl moe_h_v; + act_gl moe_sh_v; + int native_prefill_reset = 0; + const SparseMoe *sparse_moe = nullptr; +}; +#ifndef BARRIER_BACKOFF_NS +#define BARRIER_BACKOFF_NS 32 +#endif +#define BARRIER_GEN_WORD 32 +__device__ __forceinline__ void grid_barrier(unsigned int *state) { + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + auto *wide_state = reinterpret_cast(state); + const unsigned long long arrival = atomicAdd(wide_state, 1ull); + const unsigned int rank = (unsigned int)(arrival % gridDim.x); + if (rank != gridDim.x - 1u) { + const unsigned long long target = arrival - rank + gridDim.x; + volatile unsigned long long *counter = wide_state; + while (*counter < target) __nanosleep(BARRIER_BACKOFF_NS); + } + } + __syncthreads(); +} +__device__ __forceinline__ float gelu_tanh_fp32(float v) { + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + return 0.5f * v * (1.f + tanhf(inner)); +} +__device__ void layernorm_rows( + const bf16 *src, const float *weight, bf16 *dst, float *scratch) { + constexpr int PER_THREAD = D_MODEL / THREADS; + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const bf16 *in = src + (size_t)row * D_MODEL; + float v[PER_THREAD], wv[PER_THREAD]; + float sum = 0.f, sq = 0.f; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + v[j] = __bfloat162float(in[threadIdx.x + j * THREADS]); + wv[j] = weight[threadIdx.x + j * THREADS]; + sum += v[j]; + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sq += __shfl_xor_sync(0xffffffff, sq, off); + } + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) { scratch[warp] = sum; scratch[WARPS + warp] = sq; } + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f, q = 0.f; + for (int w = 0; w < WARPS; ++w) { s += scratch[w]; q += scratch[WARPS + w]; } + scratch[2 * WARPS] = s / D_MODEL; + scratch[2 * WARPS + 1] = + rsqrtf(q / D_MODEL - (s / D_MODEL) * (s / D_MODEL) + LN_EPS); + } + __syncthreads(); + const float mean = scratch[2 * WARPS], inv = scratch[2 * WARPS + 1]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + dst[(size_t)row * D_MODEL + i] = + __float2bfloat16((v[j] - mean) * inv * wv[j]); + } + __syncthreads(); + } +} +__device__ void bb_norm_residual( + bf16 *h, bf16 *r, const bf16 *weight, bf16 *out, float *scratch, + int n_rows = BATCH, bf16 *snapshot = nullptr, int snapshot_layer = -1, + float *snapshot_fp32 = nullptr, float *out_fp32 = nullptr) { + constexpr int PER_THREAD = D_MODEL / THREADS; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const size_t base = (size_t)row * D_MODEL; + float v[PER_THREAD]; + bf16 wv[PER_THREAD]; + float sq = 0.f; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + wv[j] = weight[i]; + const bf16 hv = h[base + i]; + if (snapshot) + snapshot[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = hv; + if (snapshot_fp32) + snapshot_fp32[((size_t)snapshot_layer * BATCH + row) * D_MODEL + i] = + __bfloat162float(hv); + v[j] = __bfloat162float(hv) + __bfloat162float(r[base + i]); + r[base + i] = __float2bfloat16(v[j]); + sq += v[j] * v[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + if (lane == 0) scratch[warp] = sq; + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f; + for (int w = 0; w < WARPS; ++w) s += scratch[w]; + scratch[WARPS] = rsqrtf(s / D_MODEL + RMS_EPS); + } + __syncthreads(); + const float inv = scratch[WARPS]; + #pragma unroll + for (int j = 0; j < PER_THREAD; ++j) { + const int i = threadIdx.x + j * THREADS; + const bf16 normed = __float2bfloat16( + __bfloat162float(wv[j]) * __bfloat162float(__float2bfloat16(v[j] * inv))); + out[base + i] = normed; + if (out_fp32) out_fp32[base + i] = __bfloat162float(normed); + } + __syncthreads(); + } +} +__device__ void mamba_conv( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, const int *slot_map MK2_PROGRESS_PARAMETER(progress_mask)) { + const int total = BATCH * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int row = idx / MB_CONV_DIM; + MK2_PROGRESS_GUARD(progress_mask, row); + const int c = idx - row * MB_CONV_DIM; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + bf16 *state = L.conv_state + (size_t)slot_map[row] * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + const float raw = __bfloat162float( + zxbcdt[(size_t)row * MB_IN_PROJ + MB_INNER + c]); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) + acc += __bfloat162float(w[k]) * __bfloat162float(state[k]); + acc += __bfloat162float(w[MB_CONV_K - 1]) * raw; + acc = acc / (1.f + __expf(-acc)); + xbc_act[(size_t)row * MB_CONV_DIM + c] = __float2bfloat16(acc); + #pragma unroll + for (int k = 0; k < MB_CONV_K - 2; ++k) state[k] = state[k + 1]; + state[MB_CONV_K - 2] = __float2bfloat16(raw); + } +} +__device__ __forceinline__ float mamba_softplus(float x) { + return x > 20.f ? x : log1pf(__expf(x)); +} +__device__ void mamba_ssm( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + const int *slot_map MK2_PROGRESS_PARAMETER(progress_mask)) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + constexpr int split = MK2_SSM_SPLITS; + constexpr int channels = MB_HEAD_DIM / split; + const int pairs = BATCH * MB_HEADS * split; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / (MB_HEADS * split); + MK2_PROGRESS_GUARD(progress_mask, row); + const int head = (pair / split) % MB_HEADS; + const int first_channel = (pair % split) * channels; + const int group = head / MB_HEADS_PER_GROUP; + const bf16 *xa = xbc_act + (size_t)row * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = zxbcdt + (size_t)row * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + L.dt_bias[head]); + const float dA = __expf(L.A[head] * dt); + const float Dh = L.D[head]; + const size_t ssm_at = (size_t)slot_map[row] * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE; + float *st = L.ssm_state + ssm_at; + const float *st_in = st; + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } + constexpr int NSTATE = MB_STATE / 32; + constexpr int SSM_PREFETCH = channels < 8 ? channels : 8; + static_assert(channels % SSM_PREFETCH == 0); + for (int g0 = 0; g0 < channels; g0 += SSM_PREFETCH) { + float sv[SSM_PREFETCH][NSTATE]; + float xvr[SSM_PREFETCH]; + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const size_t prow = (size_t)(first_channel + g0 + gi) * MB_STATE; + xvr[gi] = __bfloat162float( + xa[head * MB_HEAD_DIM + first_channel + g0 + gi]); + #pragma unroll + for (int j = 0; j < NSTATE; ++j) + sv[gi][j] = st_in[prow + lane + j * 32]; + } + #pragma unroll + for (int gi = 0; gi < SSM_PREFETCH; ++gi) { + const int p = first_channel + g0 + gi; + const float xv = xvr[gi]; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < NSTATE; ++j) { + const int n = lane + j * 32; + const size_t at = (size_t)p * MB_STATE + n; + const float s = sv[gi][j] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)row * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } + } +} +__device__ void mamba_conv_prefill( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *xbc_act, int n_tokens, + int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * MB_CONV_DIM; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / MB_CONV_DIM; + const int c = idx - t * MB_CONV_DIM; + const int pos = t / n_seq; + const int seq = t - pos * n_seq; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const bf16 *w = L.conv_w + (size_t)c * MB_CONV_K; + const bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + float acc = __bfloat162float(L.conv_b[c]); + #pragma unroll + for (int k = 0; k < MB_CONV_K; ++k) { + const int back = (MB_CONV_K - 1) - k; + const int src = (pos - back >= 0) ? t - back * n_seq : -1; + float v; +#ifdef PF_CONV_ZERO_STATE + if (src < 0) v = 0.f; +#else + if (src < 0) v = __bfloat162float(state[pos + k]); +#endif + else v = __bfloat162float( + zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]); + acc += __bfloat162float(w[k]) * v; + } +#ifndef PF_CONV_NO_SILU + acc = acc / (1.f + __expf(-acc)); +#endif + xbc_act[(size_t)t * MB_CONV_DIM + c] = __float2bfloat16(acc); + } +} +__device__ void mamba_conv_prefill_state( + const MambaLayer &L, const bf16 *zxbcdt, int n_tokens, int slot, + int n_seq = 1, const int *seq_slot = nullptr) { + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < n_seq * MB_CONV_DIM; idx += gridDim.x * THREADS) { + const int seq = idx / MB_CONV_DIM; + const int c = idx - seq * MB_CONV_DIM; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + bf16 *state = L.conv_state + (size_t)my_slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int p = n_pos - (MB_CONV_K - 1) + k; + state[k] = p >= 0 + ? zxbcdt[((size_t)p * n_seq + seq) * MB_IN_PROJ + MB_INNER + c] + : old[k + n_pos]; + } + } + return; + } + for (int c = blockIdx.x * THREADS + threadIdx.x; c < MB_CONV_DIM; + c += gridDim.x * THREADS) { + bf16 *state = L.conv_state + (size_t)slot * L.conv_slot_stride + + (size_t)c * (MB_CONV_K - 1); + bf16 old[MB_CONV_K - 1]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) old[k] = state[k]; + #pragma unroll + for (int k = 0; k < MB_CONV_K - 1; ++k) { + const int src = n_tokens - (MB_CONV_K - 1) + k; +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else +#ifdef PF_CONV_STATE_NO_SHIFT + if (src >= 0) + state[k] = zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c]; +#else + state[k] = src >= 0 + ? zxbcdt[(size_t)src * MB_IN_PROJ + MB_INNER + c] + : old[k + n_tokens]; +#endif +#endif +#endif + } + } +} +__device__ void mamba_ssm_prefill( + const MambaLayer &L, const bf16 *zxbcdt, const bf16 *xbc_act, bf16 *gated, + int n_tokens, int slot, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + if (n_seq > 1) { + const int n_pos = n_tokens / n_seq; + const int units = n_seq * MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; + unit += warps_total) { + const int seq = unit / (MB_HEADS * MB_HEAD_DIM); + const int rem = unit - seq * (MB_HEADS * MB_HEAD_DIM); + const int head = rem / MB_HEAD_DIM; + const int p = rem - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; + const int my_slot = seq_slot ? seq_slot[seq] : slot; + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *const st = L.ssm_state + (size_t)my_slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE + + (size_t)p * MB_STATE; + float sreg[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) sreg[j] = st[lane + j * 32]; + for (int pos = 0; pos < n_pos; ++pos) { + const size_t t = (size_t)pos * n_seq + seq; + const bf16 *xa = xbc_act + t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + + group * MB_STATE; + const bf16 *dtrow = zxbcdt + t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + const int n = lane + j * 32; + sreg[j] = sreg[j] * dA + __bfloat162float(Bsec[n]) * dt * xv; + acc += sreg[j] * __bfloat162float(Csec[n]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) st[lane + j * 32] = sreg[j]; + } + return; + } +#if PF_SSM_SPLIT_P + const int units = MB_HEADS * MB_HEAD_DIM; + for (int unit = blockIdx.x * WARPS + warp; unit < units; unit += warps_total) { + const int head = unit / MB_HEAD_DIM; + const int p_only = unit - head * MB_HEAD_DIM; + const int group = head / MB_HEADS_PER_GROUP; +#else + for (int head = blockIdx.x * WARPS + warp; head < MB_HEADS; + head += warps_total) { + const int group = head / MB_HEADS_PER_GROUP; +#endif + const float Ah = L.A[head], Dh = L.D[head], dtb = L.dt_bias[head]; + float *st = L.ssm_state + (size_t)slot * L.ssm_slot_stride + + (size_t)head * MB_HEAD_DIM * MB_STATE +#if PF_SSM_SPLIT_P + + (size_t)p_only * MB_STATE +#endif + ; + for (int t = 0; t < n_tokens; ++t) { + const bf16 *xa = xbc_act + (size_t)t * MB_CONV_DIM; + const bf16 *Bsec = xa + MB_INNER + group * MB_STATE; + const bf16 *Csec = xa + MB_INNER + MB_GROUPS * MB_STATE + group * MB_STATE; + const bf16 *dtrow = + zxbcdt + (size_t)t * MB_IN_PROJ + MB_INNER + MB_CONV_DIM; + const float dt = mamba_softplus( + __bfloat162float(dtrow[head]) + dtb); + const float dA = __expf(Ah * dt); + float Bv[MB_STATE / 32], Cv[MB_STATE / 32]; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { + Bv[j] = __bfloat162float(Bsec[lane + j * 32]) * dt; + Cv[j] = __bfloat162float(Csec[lane + j * 32]); + } +#if PF_SSM_SPLIT_P + { + const int p = p_only; +#else + for (int p = 0; p < MB_HEAD_DIM; ++p) { +#endif + const float xv = __bfloat162float(xa[head * MB_HEAD_DIM + p]); + float acc = 0.f; + #pragma unroll + for (int j = 0; j < MB_STATE / 32; ++j) { +#if PF_SSM_SPLIT_P + const size_t at = (size_t)lane + j * 32; +#else + const size_t at = (size_t)p * MB_STATE + lane + j * 32; +#endif + const float s = st[at] * dA + Bv[j] * xv; + st[at] = s; + acc += s * Cv[j]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) + gated[(size_t)t * MB_INNER + head * MB_HEAD_DIM + p] = + __float2bfloat16(acc + xv * Dh); + } + } + } +} +__device__ void mamba_gated_norm( + const MambaLayer &L, const bf16 *zxbcdt, bf16 *gated, int n_rows = BATCH) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_rows * MB_GROUPS; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int row = pair / MB_GROUPS; + const int group = pair - row * MB_GROUPS; + const bf16 *zrow = zxbcdt + (size_t)row * MB_IN_PROJ; + bf16 *grow = gated + (size_t)row * MB_INNER; + float sq = 0.f; + float held[MB_GROUP_SIZE / 32], wg[MB_GROUP_SIZE / 32]; + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + const float z = __bfloat162float(zrow[idx]); + const float g = __bfloat162float(grow[idx]) * (z / (1.f + __expf(-z))); + wg[j] = __bfloat162float(L.gnorm_w[idx]); + held[j] = g; + sq += g * g; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + const float inv = rsqrtf(sq / MB_GROUP_SIZE + RMS_EPS); + #pragma unroll + for (int j = 0; j < MB_GROUP_SIZE / 32; ++j) { + const int idx = group * MB_GROUP_SIZE + lane + j * 32; + grow[idx] = __float2bfloat16(held[j] * inv * wg[j]); + } + } +} +__device__ void bb_attention( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int past_len_scalar, + const int *seq_len MK2_PROGRESS_PARAMETER(progress_mask)) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + constexpr int Q_PER_WARP = 2; + constexpr int GROUPS_PER_KV = 2; + const int pairs = BATCH * BB_KV_HEADS * GROUPS_PER_KV; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + for (int pair = blockIdx.x + warp * gridDim.x; pair < pairs; + pair += gridDim.x * WARPS) { + const int row = pair / (BB_KV_HEADS * GROUPS_PER_KV); + MK2_PROGRESS_GUARD(progress_mask, row); + const int within_row = pair % (BB_KV_HEADS * GROUPS_PER_KV); + const int kv = within_row / GROUPS_PER_KV; + const int q_group = within_row % GROUPS_PER_KV; + const int q_first = q_group * Q_PER_WARP; + const int q_count = q_group == 0 ? Q_PER_WARP : 1; + const int past_len = seq_len ? seq_len[row] : past_len_scalar; + const int T = past_len + 1; + const bf16 *base = qkv + (size_t)row * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; +#if !BB_ATTN_SKIP_KV_WRITE + if (q_group == 0) { + bf16 *kc = kv_slot(L, row, 0, past_len) + kv_offset(L, past_len, kv); + bf16 *vc = kv_slot(L, row, 1, past_len) + kv_offset(L, past_len, kv); + for (int i = lane; i < BB_HEAD_DIM; i += 32) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +#endif + __syncwarp(); + constexpr int PER_LANE = BB_HEAD_DIM / 32; + float acc[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[h][j] = 0.f; + } + float m[Q_PER_WARP] = {-CUDART_INF_F, -CUDART_INF_F}; + float denom[Q_PER_WARP] = {0.f, 0.f}; + constexpr int PF = 8; + const int tpb = L.tokens_per_block; + const size_t elems_per_block = + (size_t)tpb * BB_KV_HEADS * BB_HEAD_DIM; + const int *bt_k = L.block_table + (size_t)(row * 2) * L.max_blocks_per_seq; + const int *bt_v = bt_k + L.max_blocks_per_seq; + const int head_base = kv * tpb * BB_HEAD_DIM; + float qr[Q_PER_WARP][PER_LANE]; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const bf16 *q = base + (kv * kv_per_q + q_first + h) * BB_HEAD_DIM; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + qr[h][j] = __bfloat162float(q[lane + j * 32]); + } + for (int page0 = 0; page0 < T; page0 += tpb) { + const int k_page = bt_k[page0 / tpb] & 0x7FFFFFFF; + const int v_page = bt_v[page0 / tpb] & 0x7FFFFFFF; + const bf16 *kbase = L.pool_base + + (size_t)k_page * elems_per_block + + head_base; + const bf16 *vbase = L.pool_base + + (size_t)v_page * elems_per_block + + head_base; + const int tend = page0 + tpb < T ? page0 + tpb : T; + for (int t0 = page0; t0 < tend; t0 += PF) { + const int navail = tend - t0 < PF ? tend - t0 : PF; + float kr[PF][PER_LANE], vr[PF][PER_LANE]; + #pragma unroll + for (int u = 0; u < PF; ++u) { + const int tpos = t0 + (u < navail ? u : navail - 1); +#if BB_ATTN_LEGACY_CUR_KV + const bf16 *kt = kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = vbase + (tpos - page0) * BB_HEAD_DIM; +#else + const bool cur_token = tpos == past_len; + const bf16 *kt = cur_token + ? k_new : kbase + (tpos - page0) * BB_HEAD_DIM; + const bf16 *vt = cur_token + ? v_new : vbase + (tpos - page0) * BB_HEAD_DIM; +#endif + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + kr[u][j] = __bfloat162float(kt[i]); + vr[u][j] = __bfloat162float(vt[i]); + } + } + #pragma unroll + for (int u = 0; u < PF; ++u) { + if (u >= navail) break; + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + dot += qr[h][j] * kr[u][j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float s = dot * scale; + const float m_new = fmaxf(m[h], s); + const float rescale = __expf(m[h] - m_new); + const float w = __expf(s - m_new); + denom[h] = denom[h] * rescale + w; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) + acc[h][j] = acc[h][j] * rescale + w * vr[u][j]; + m[h] = m_new; + } + } + } + } + #pragma unroll + for (int h = 0; h < Q_PER_WARP; ++h) { + if (h >= q_count) break; + const float inv = 1.f / denom[h]; + const int head = kv * kv_per_q + q_first + h; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)row * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[h][j] * inv); + } + } + } +} +__device__ void bb_attention_prefill_kv( + const AttnLayer &L, const bf16 *qkv, int n_tokens, int seq_base, int tok0, + int n_seq = 1, const int *seq_slot = nullptr) { + const int total = n_tokens * BB_KV_HEADS; + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < total; + idx += gridDim.x * THREADS) { + const int t = idx / BB_KV_HEADS; + const int kv = idx - t * BB_KV_HEADS; + const bf16 *base = qkv + (size_t)t * BB_QKV; + const bf16 *k_new = base + BB_Q_SIZE + kv * BB_HEAD_DIM; + const bf16 *v_new = base + BB_Q_SIZE + BB_KV_SIZE + kv * BB_HEAD_DIM; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + bf16 *kc = kv_slot(L, seq, 0, pos) + kv_offset(L, pos, kv); + bf16 *vc = kv_slot(L, seq, 1, pos) + kv_offset(L, pos, kv); + for (int i = 0; i < BB_HEAD_DIM; ++i) { + kc[i] = k_new[i]; + vc[i] = v_new[i]; + } + } +} +__device__ void bb_attention_prefill( + const AttnLayer &L, const bf16 *qkv, bf16 *out, int n_tokens, int seq_base, + int tok0, int n_seq = 1, const int *seq_slot = nullptr) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const int pairs = n_tokens * BB_ATT_HEADS; + const float scale = rsqrtf((float)BB_HEAD_DIM); + const int kv_per_q = BB_ATT_HEADS / BB_KV_HEADS; + constexpr int PER_LANE = BB_HEAD_DIM / 32; + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; pair += warps_total) { + const int t = pair / BB_ATT_HEADS; + const int head = pair - t * BB_ATT_HEADS; + const int kv = head / kv_per_q; + const bf16 *q = qkv + (size_t)t * BB_QKV + head * BB_HEAD_DIM; + float acc[PER_LANE]; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) acc[j] = 0.f; + float m = -CUDART_INF_F, denom = 0.f; + const int pos = n_seq > 1 ? t / n_seq : tok0 + t; + const int seq = n_seq > 1 + ? (seq_slot ? seq_slot[t - (t / n_seq) * n_seq] : t % n_seq) + : seq_base; + for (int s = 0; s <= pos; ++s) { + const bf16 *kt = kv_slot(L, seq, 0, s) + kv_offset(L, s, kv); + float dot = 0.f; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + dot += __bfloat162float(q[i]) * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + dot += __shfl_xor_sync(0xffffffff, dot, off); + const float sc = dot * scale; + const float m_new = fmaxf(m, sc); + const float rescale = __expf(m - m_new); + const float wt = __expf(sc - m_new); + denom = denom * rescale + wt; + const bf16 *vt = kv_slot(L, seq, 1, s) + kv_offset(L, s, kv); + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + acc[j] = acc[j] * rescale + wt * __bfloat162float(vt[i]); + } + m = m_new; + } + const float inv = 1.f / denom; + #pragma unroll + for (int j = 0; j < PER_LANE; ++j) { + const int i = lane + j * 32; + out[(size_t)t * D_MODEL + head * BB_HEAD_DIM + i] = + __float2bfloat16(acc[j] * inv); + } + } +} +__device__ void bb_moe_route( + const bf16 *normed, const bf16 *gate_w, const bf16 *gate_bias, float *moe_w, + int n_rows = BATCH) { + __shared__ float score[BB_EXPERTS]; + for (int row = blockIdx.x; row < n_rows; row += gridDim.x) { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + for (int e = warp; e < BB_EXPERTS; e += WARPS) { + const bf16 *wr = gate_w + (size_t)e * D_MODEL; + float acc = 0.f; + for (int i = lane; i < D_MODEL; i += 32) + acc += __bfloat162float(normed[(size_t)row * D_MODEL + i]) * + __bfloat162float(wr[i]); + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + if (lane == 0) score[e] = 1.f / (1.f + __expf(-acc)); + } + __syncthreads(); + if (threadIdx.x == 0) { + float remaining[BB_EXPERTS]; + int pick[BB_TOP_K]; + for (int e = 0; e < BB_EXPERTS; ++e) { + moe_w[row * BB_EXPERTS + e] = 0.f; + remaining[e] = score[e] + __bfloat162float(gate_bias[e]); + } + float total = 0.f; + for (int k = 0; k < BB_TOP_K; ++k) { + int best = 0; + for (int e = 1; e < BB_EXPERTS; ++e) + if (remaining[e] > remaining[best]) best = e; + pick[k] = best; + total += score[best]; + remaining[best] = -CUDART_INF_F; + } + const float norm = BB_ROUTED_SCALE / (total + 1e-20f); + for (int k = 0; k < BB_TOP_K; ++k) + moe_w[row * BB_EXPERTS + pick[k]] = score[pick[k]] * norm; + } + __syncthreads(); + } +} +__device__ __forceinline__ bf16 bb_moe_fold_deterministic(const float *parts, int idx) { + const size_t base = (size_t)(idx / D_MODEL) * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert <= BB_EXPERTS; ++expert) + value = __fadd_rn(value, parts[base + (size_t)expert * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ bf16 bb_moe_fold_selected( + const float *parts, const float *weights, int idx) { + const int row = idx / D_MODEL; + const size_t base = (size_t)row * (BB_EXPERTS + 1) * D_MODEL + idx % D_MODEL; + float value = 0.f; + #pragma unroll + for (int expert = 0; expert < BB_EXPERTS; ++expert) { + const float part = weights[row * BB_EXPERTS + expert] != 0.f + ? parts[base + (size_t)expert * D_MODEL] : 0.f; + value = __fadd_rn(value, part); + } + value = __fadd_rn(value, parts[base + (size_t)BB_EXPERTS * D_MODEL]); + return __float2bfloat16(value); +} +__device__ __forceinline__ uint32_t float_to_ordered(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ __forceinline__ float ordered_to_float(uint32_t u) { + return __uint_as_float((u & 0x80000000u) ? (u & 0x7FFFFFFFu) : ~u); +} +__device__ __forceinline__ void sample_argmax_merge( + float &value, int &rank, float other, int other_rank) { + if (other > value || (other == value && other_rank < rank)) { + value = other; + rank = other_rank; + } +} +__device__ __forceinline__ int sample_block_argmax(float value, int rank) { + __shared__ float warp_value[WARPS]; + __shared__ int warp_rank[WARPS]; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + if (lane == 0) { warp_value[warp] = value; warp_rank[warp] = rank; } + __syncthreads(); + if (warp == 0) { + value = lane < WARPS ? warp_value[lane] : -CUDART_INF_F; + rank = lane < WARPS ? warp_rank[lane] : INT_MAX; + #pragma unroll + for (int off = 16; off; off >>= 1) { + const float other = __shfl_down_sync(0xffffffffu, value, off); + const int other_rank = __shfl_down_sync(0xffffffffu, rank, off); + sample_argmax_merge(value, rank, other, other_rank); + } + } + return rank; +} +__device__ __forceinline__ int sample_eos_insertion_rank( + const uint64_t *keys, float eos) { + int lo = 0, hi = TOP_K; + while (lo < hi) { + const int mid = (lo + hi) >> 1; + if (eos > ordered_to_float((uint32_t)(keys[mid] >> 32))) hi = mid; + else lo = mid + 1; + } + return lo; +} +__device__ __forceinline__ void sample_topk_serial( + const uint64_t *keys, float eos, const float *noise, int *out) { + float best_noised = -CUDART_INF_F; + int best_code = 0; + int a = 0; + bool eos_used = false; + for (int j = 0; j < TOP_K; ++j) { + const float av = ordered_to_float((uint32_t)(keys[a] >> 32)); + float value; + int code; + if (!eos_used && eos > av) { + value = eos; + code = AUDIO_EOS; + eos_used = true; + } else { + value = av; + code = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + ++a; + } + const float noised = value + noise[j]; + if (noised > best_noised) { + best_noised = noised; + best_code = code; + } + } + *out = best_code; +} +__device__ __forceinline__ void sample_topk_parallel( + const uint64_t *keys, float eos, const float *noise, int *out) { + const int j = threadIdx.x; + const bool sorted_nan = j < TOP_K && + isnan(ordered_to_float((uint32_t)(keys[j] >> 32))); + if (__syncthreads_or(sorted_nan)) { + if (threadIdx.x == 0) sample_topk_serial(keys, eos, noise, out); + return; + } + const int lane = threadIdx.x & 31; + int insertion = lane == 0 ? sample_eos_insertion_rank(keys, eos) : 0; + insertion = __shfl_sync(0xffffffffu, insertion, 0); + float value = -CUDART_INF_F; + int rank = INT_MAX; + if (j < TOP_K) { + const int a = j - (j > insertion); + const float v = j == insertion + ? eos : ordered_to_float((uint32_t)(keys[a] >> 32)); + const float noised = v + noise[j]; + if (noised > -CUDART_INF_F) { value = noised; rank = j; } + } + const int winner = sample_block_argmax(value, rank); + if (threadIdx.x == 0) { + if (winner == INT_MAX) *out = 0; + else if (winner == insertion) *out = AUDIO_EOS; + else { + const int a = winner - (winner > insertion); + *out = (int)(0xFFFFFFFFu - (uint32_t)(keys[a] & 0xFFFFFFFFu)); + } + } +} +constexpr int SORT_ITEMS_PER_THREAD = SORT_N / THREADS; +struct sampler_storage { + uint64_t keys[2][SORT_N]; +}; +__device__ void sample_codes( + const bf16 *logits, const bf16 *bias_unused, const int8_t *forbidden, + const float *gumbel, int pos, int *out_codes, + sampler_storage &storage, float temperature) { +#ifdef HEAD_BF16 + const bf16 *bias = bias_unused; +#else + (void)bias_unused; +#endif + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + uint64_t thread_keys[SORT_ITEMS_PER_THREAD]; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + float v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + n]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[n]); +#endif + v /= temperature; + thread_keys[item] = ((uint64_t)float_to_ordered(v) << 32) | + (uint64_t)(0xFFFFFFFFu - (uint32_t)n); + } + __shared__ float eos_v; + if (threadIdx.x == 0) { + float v = -CUDART_INF_F; + if (!forbidden[AUDIO_EOS]) { + v = __bfloat162float(logits[(size_t)row * VOCAB_PADDED + AUDIO_EOS]); +#ifdef HEAD_BF16 + v += __bfloat162float(bias[AUDIO_EOS]); +#endif + v /= temperature; + } + eos_v = v; + } + int cross_stage = 0; + for (int k = 2; k <= SORT_N; k <<= 1) { + for (int j = k >> 1; j > 0; j >>= 1) { + if (j >= 32) { + const int bank = cross_stage++ & 1; + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + storage.keys[bank][n] = thread_keys[item]; + } + __syncthreads(); + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + const uint64_t other = storage.keys[bank][n ^ j]; + const bool take_hi = ((n & k) == 0) == ((n & j) == 0); + const uint64_t mine = thread_keys[item]; + thread_keys[item] = take_hi + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } else { + #pragma unroll + for (int item = 0; item < SORT_ITEMS_PER_THREAD; ++item) { + const int n = threadIdx.x + item * THREADS; + const uint64_t other = __shfl_xor_sync( + 0xffffffffu, thread_keys[item], j); + const bool take_hi = ((n & k) == 0) == ((n & j) == 0); + const uint64_t mine = thread_keys[item]; + thread_keys[item] = take_hi + ? (mine > other ? mine : other) + : (mine < other ? mine : other); + } + } + } + } + const int top_bank = cross_stage & 1; + if (threadIdx.x < TOP_K) + storage.keys[top_bank][threadIdx.x] = thread_keys[0]; + __syncthreads(); + sample_topk_parallel( + storage.keys[top_bank], eos_v, + gumbel + ((size_t)row * N_POS + pos) * TOP_K, + out_codes + row); + __syncthreads(); + } +} +__device__ __forceinline__ int mk2_expert_tile(int work, unsigned mask, + int tiles_per_expert) { + if (!tiles_per_expert) return work; + int rank = work / tiles_per_expert; + while (rank-- > 0) mask &= mask - 1; + return (__ffs(mask) - 1) * tiles_per_expert + work % tiles_per_expert; +} +__device__ void gemm( + const weight_gl &w, const act_gl &a, int n_out, int k_total, + bf16 *dest, int dest_stride, bf16 *residual, Epilogue epi, + weight_tile (&w_s)[STAGES], act_tile (&a_s)[STAGES], out_tile &o_s, + semaphore *filled, semaphore *drained, out_tile_f &o_sf, int &iter_io, + int &tile_io, + const bf16 *bias = nullptr, + int out_tiles_per_expert = 0, int act_k_tiles_per_expert = 0, + const float *moe_w = nullptr, float *moe_out = nullptr, + int n_tok = 1, + int rows_valid = BATCH, + const float *active_routes = nullptr, int active_expert_tiles = 0) { + unsigned active_mask = 0; + int compact_tiles = 0; + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + if (active_routes && active_expert_tiles) { + for (int expert = 0; expert < BB_EXPERTS; ++expert) + if (active_routes[expert] != 0.f) active_mask |= 1u << expert; + compact_tiles = active_expert_tiles; + } + } + const int n_tiles = compact_tiles ? __popc(active_mask) * compact_tiles + : n_out / OUT_ATOM; + const int iters = k_total / KCHUNK; + int tile_start = (int)blockIdx.x - tile_io % (int)gridDim.x; + if (tile_start < 0) tile_start += (int)gridDim.x; + tile_io += n_tiles; + int tiles_here = 0; + for (int t = tile_start; t < n_tiles; t += gridDim.x) ++tiles_here; + const int iter_base = iter_io; + iter_io += tiles_here * iters * n_tok; + const int my_tiles = (n_tiles - tile_start + gridDim.x - 1) / gridDim.x; + if (warpgroup::groupid() == 1) { + if (warpgroup::laneid() == 0) { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) { + const int tile_id = mk2_expert_tile(work, active_mask, compact_tiles); + for (int tok = 0; tok < n_tok; ++tok) + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(drained[s], phase); + tma::expect_bytes(filled[s], + sizeof(weight_tile) + sizeof(act_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[s], w, {tile_id, i}, filled[s]); + const int a_col = out_tiles_per_expert + ? (tile_id / out_tiles_per_expert) * act_k_tiles_per_expert + i + : i; + tma::load_async<2, cache_policy::EVICT_LAST>( + a_s[s], a, {tok, a_col}, filled[s]); + } + } + } + } else { + int iter = iter_base; + for (int work = tile_start; work < n_tiles; work += gridDim.x) + for (int tok = 0; tok < n_tok; ++tok) { + const int tile_id = mk2_expert_tile(work, active_mask, compact_tiles); + rt_fl<16, BATCH_PAD> accum; + warp::zero(accum); + for (int i = 0; i < iters; ++i, ++iter) { + const int s = iter % STAGES; + const int phase = (iter / STAGES) & 1; + wait(filled[s], phase); + warpgroup::mma_ABt(accum, w_s[s], a_s[s]); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[s]); + } + const int base = tile_id * OUT_ATOM; + const int grow0 = tok * BATCH_PAD; + if (epi == EPI_BIAS) { + warpgroup::store(o_sf, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM, col = idx % OUT_ATOM; + dest[(size_t)(grow0 + row) * dest_stride + base + col] = + __float2bfloat16(o_sf[int2{col, row}] + + __bfloat162float(bias[base + col])); + } + warpgroup::sync(1); + continue; + } + warpgroup::store(o_s, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int idx = warpgroup::laneid(); idx < OUT_ATOM * rows_valid; + idx += 128) { + const int row = idx / OUT_ATOM; + const int col = idx % OUT_ATOM; + const float v = __bfloat162float(o_s[int2{col, row}]); + const int grow = grow0 + row; + const size_t at = (size_t)grow * dest_stride + base + col; + switch (epi) { + case EPI_STORE: dest[at] = __float2bfloat16(v); break; + case EPI_SILU: + dest[at] = __float2bfloat16(v / (1.f + __expf(-v))); + break; + case EPI_MOE_ACC: { + const int expert = tile_id / out_tiles_per_expert; + const int local = (tile_id % out_tiles_per_expert) * OUT_ATOM + + col; + const float wt = moe_w[grow * BB_EXPERTS + expert]; + if (wt != 0.f) + moe_out[((size_t)grow * (BB_EXPERTS + 1) + expert) * + D_MODEL + local] = __fmul_rn(wt, v); + break; + } + case EPI_ACC_F32: + moe_out[((size_t)grow * (BB_EXPERTS + 1) + BB_EXPERTS) * D_MODEL + base + col] = v; + break; + case EPI_BIAS: break; + case EPI_GELU: + dest[at] = __float2bfloat16(gelu_tanh_fp32(v)); + break; + case EPI_RESIDUAL: + residual[at] = + __float2bfloat16(__bfloat162float(residual[at]) + v); + break; + } + } + warpgroup::sync(1); + } + (void)my_tiles; + } +} +#include "sparse_moe_impl.cuh" +__global__ __launch_bounds__(THREADS, 1) void predictor_mk2( + const __grid_constant__ Globals g, const Backbone *bb, int n_layers, + int n_pos, int n_backbone_layers, const int *run_predicate = nullptr, + const unsigned char *native_predictor_active_mask = nullptr + MK2_PROGRESS_PARAMETER(native_progress_mask)) { + if (!bb && run_predicate && *run_predicate == 0) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; + i < BATCH * N_POS; i += gridDim.x * blockDim.x) + g.sampled[i] = 0; + return; + } + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + atomicAdd(bb->launches, 1ull); + unsigned long long t0; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + bb->launches[2] = t0; + } + extern __shared__ int __shm[]; + tma_swizzle_allocator alloc(&__shm[0]); + weight_tile (&w_s)[STAGES] = alloc.allocate(); + act_tile (&a_s)[STAGES] = alloc.allocate(); + out_tile (&o_s) = alloc.allocate(); + out_tile_f (&o_sf) = alloc.allocate(); + __shared__ semaphore filled[STAGES], drained[STAGES]; + __shared__ float logits[N_POS]; + __shared__ float ln_scratch[2 * WARPS + 2]; + __shared__ sampler_storage sort_storage; + if (threadIdx.x == 0) { + for (int s = 0; s < STAGES; ++s) { + init_semaphore(filled[s], 0, 1); + init_semaphore(drained[s], CONSUMER_WARPS, 0); + for (int c = 0; c < CONSUMER_WARPS; ++c) arrive(drained[s]); + } + } + int gemm_iter = 0; + int gemm_tile = 0; + __syncthreads(); + const bool prefilling = (bb && bb->n_prefill_tokens > 0); + if (g.native_embedding) { + for (int i = blockIdx.x * THREADS + threadIdx.x; + i < BATCH * D_MODEL; i += gridDim.x * THREADS) { + bb->h[i] = g.native_embedding[i]; + bb->r[i] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + } + for (int layer = 0; !prefilling && layer < n_backbone_layers; ++layer) { + bb_norm_residual(bb->h, bb->r, bb->norm_w[layer], bb->normed, ln_scratch); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + if (mi >= 0) { + if (BB_SKIP_MAMBA) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->normed_v, MB_IN_PROJ, D_MODEL, + bb->zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + mamba_conv(L, bb->zxbcdt, bb->xbc_act, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + mamba_ssm(L, bb->zxbcdt, bb->xbc_act, bb->gated, bb->slot_map + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + mamba_gated_norm(L, bb->zxbcdt, bb->gated); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->gated_v, D_MODEL, MB_INNER, + bb->h, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (bb->attn_of_layer[layer] >= 0) { + if (BB_SKIP_ATTN) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const AttnLayer &L = bb->attn[bb->attn_of_layer[layer]]; + gemm(L.qkv, bb->normed_v, BB_QKV, D_MODEL, bb->qkv, BB_QKV, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + bb_attention(L, bb->qkv, bb->gated, bb->past_len, bb->seq_len + MK2_PROGRESS_ARGUMENT(native_progress_mask)); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->gated_v, D_MODEL, D_MODEL, bb->h, D_MODEL, + nullptr, EPI_STORE, w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + } else if (BB_SKIP_MOE) { + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + grid_barrier(g.barrier); + } else { + const int mo = bb->moe_of_layer[layer]; + const MoeLayer &L = bb->moe[mo]; + if (bb->moe_w_audit) { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + (size_t)mo * BATCH * BB_EXPERTS); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * BB_EXPERTS; idx += gridDim.x * THREADS) + bb->moe_w[idx] = + bb->moe_w_inject[(size_t)mo * BATCH * BB_EXPERTS + idx]; + } else { + bb_moe_route(bb->normed, L.gate_w, L.gate_bias, bb->moe_w); + } + if constexpr (MK2_ACTIVE_EXPERTS && BATCH == 1) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * (BB_EXPERTS + 1) * D_MODEL; idx += gridDim.x * THREADS) + bb->moe_out[idx] = 0.f; + } + grid_barrier(g.barrier); + if constexpr (MK2_SPARSE_MOE && BATCH >= 2) { + if (bb->sparse_moe) { + sparse_moe_pack(bb); + grid_barrier(g.barrier); + } + } + if (bb->sparse_moe) { + sparse_moe_gemm(L.exp_up, bb->sparse_moe->gathered_v, + *bb->sparse_moe, bb->moe_h, bb->moe_w, bb->moe_out, + w_s, a_s, o_s, filled, drained, gemm_iter, gemm_tile); + } else { + gemm(L.exp_up, bb->normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, nullptr, 1, BATCH, + bb->moe_w, BB_MOE_INTER / OUT_ATOM); + } + gemm(L.sh_up, bb->normed_v, BB_SHARED_INTER, D_MODEL, + bb->moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + if (bb->sparse_moe) { + sparse_moe_gemm(L.exp_down, bb->sparse_moe->hidden_v, + *bb->sparse_moe, bb->moe_h, bb->moe_w, bb->moe_out, + w_s, a_s, o_s, filled, drained, gemm_iter, gemm_tile); + } else { + gemm(L.exp_down, bb->moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, bb->moe_w, bb->moe_out, + 1, BATCH, bb->moe_w, D_MODEL / OUT_ATOM); + } + gemm(L.sh_down, bb->moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->moe_out); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->h[idx] = bb_moe_fold_selected(bb->moe_out, bb->moe_w, idx); + grid_barrier(g.barrier); + } + if (bb->layer_out) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + bb->layer_out[((size_t)layer * BATCH + row) * D_MODEL + i] = + bb->h[(size_t)row * D_MODEL + i]; + grid_barrier(g.barrier); + } + } + if (bb && bb->n_prefill_tokens > 0) { + const int n_total = bb->n_prefill_tokens; + const int slot = bb->prefill_slot; + if (bb->native_prefill_reset) { + constexpr int CV = MB_CONV_DIM * (MB_CONV_K - 1); + constexpr int SS = MB_HEADS * MB_HEAD_DIM * MB_STATE; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * CV; + i += gridDim.x * THREADS) { + const int cell = i % CV; + const int row_layer = i / CV; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.conv_state[(size_t)row * L.conv_slot_stride + cell] = __float2bfloat16(0.f); + } + for (int i = blockIdx.x * THREADS + threadIdx.x; i < 14 * BATCH * SS; + i += gridDim.x * THREADS) { + const int cell = i % SS; + const int row_layer = i / SS; + const int row = row_layer % BATCH; + const MambaLayer &L = bb->mamba[row_layer / BATCH]; + L.ssm_state[(size_t)row * L.ssm_slot_stride + cell] = 0.f; + } + for (int row = blockIdx.x * THREADS + threadIdx.x; row < BATCH; + row += gridDim.x * THREADS) + const_cast(bb->seq_len)[row] = n_total / BATCH; + grid_barrier(g.barrier); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < n_total * D_MODEL; + idx += gridDim.x * THREADS) { + bb->pf_h[idx] = bb->prompt[idx]; + bb->pf_r[idx] = __float2bfloat16(0.f); + } + grid_barrier(g.barrier); + const int n_tok = (n_total + BATCH_PAD - 1) / BATCH_PAD; + const int n_seq = bb->pf_n_seq > 0 ? bb->pf_n_seq : 1; + for (int layer = 0; layer < n_backbone_layers; ++layer) { + const int T = n_total; + const int tok0 = 0; + const size_t off = 0; + { + bb_norm_residual(bb->pf_h + off, bb->pf_r + off, bb->norm_w[layer], + bb->pf_normed, ln_scratch, T); + grid_barrier(g.barrier); + const int mi = bb->mamba_of_layer[layer]; + const int ai = bb->attn_of_layer[layer]; + if (mi >= 0 && PF_SKIP_MAMBA) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else if (mi >= 0) { + const MambaLayer &L = bb->mamba[mi]; + gemm(L.in_proj, bb->pf_normed_v, MB_IN_PROJ, D_MODEL, + bb->pf_zxbcdt, MB_IN_PROJ, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + mamba_conv_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, T, slot, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_ssm_prefill(L, bb->pf_zxbcdt, bb->pf_xbc_act, bb->pf_gated, + T, slot, n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + mamba_conv_prefill_state(L, bb->pf_zxbcdt, T, slot, n_seq, + bb->pf_seq_slot); + mamba_gated_norm(L, bb->pf_zxbcdt, bb->pf_gated, T); + grid_barrier(g.barrier); + gemm(L.out_proj, bb->pf_gated_v, D_MODEL, MB_INNER, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (ai >= 0 && PF_SKIP_ATTN) { + for (int i = 0; i < 4; ++i) grid_barrier(g.barrier); + } else if (ai >= 0) { + const AttnLayer &L = bb->attn[ai]; + gemm(L.qkv, bb->pf_normed_v, BB_QKV, D_MODEL, + bb->pf_qkv, BB_QKV, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + bb_attention_prefill_kv(L, bb->pf_qkv, T, slot, tok0, n_seq, + bb->pf_seq_slot); + grid_barrier(g.barrier); + bb_attention_prefill(L, bb->pf_qkv, bb->pf_gated, T, slot, tok0, + n_seq, bb->pf_seq_slot); + grid_barrier(g.barrier); + gemm(L.o_proj, bb->pf_gated_v, D_MODEL, D_MODEL, + bb->pf_h + off, D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + } else if (PF_SKIP_MOE) { + for (int i = 0; i < 5; ++i) grid_barrier(g.barrier); + } else { + const MoeLayer &L = bb->moe[bb->moe_of_layer[layer]]; + const int mo_audit = bb->moe_of_layer[layer]; + if (bb->moe_w_audit) { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, + bb->moe_w_audit + + ((size_t)mo_audit * PREFILL_MAX_TOKENS + tok0) * + BB_EXPERTS, + T); + grid_barrier(g.barrier); + } + if (bb->moe_w_inject) { + const int mo = bb->moe_of_layer[layer]; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < T * BB_EXPERTS; idx += gridDim.x * THREADS) { + const int tk = idx / BB_EXPERTS; + const int e = idx - tk * BB_EXPERTS; + bb->pf_moe_w[idx] = bb->moe_w_inject[ + ((size_t)mo * bb->n_prefill_tokens + tok0 + tk) * + BB_EXPERTS + e]; + } + } else { + bb_moe_route(bb->pf_normed, L.gate_w, L.gate_bias, bb->pf_moe_w, T); + } + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * (BB_EXPERTS + 1) * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_moe_out[idx] = 0.f; + grid_barrier(g.barrier); + gemm(L.exp_up, bb->pf_normed_v, BB_EXPERTS * BB_MOE_INTER, D_MODEL, + bb->pf_moe_h, BB_EXPERTS * BB_MOE_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + gemm(L.sh_up, bb->pf_normed_v, BB_SHARED_INTER, D_MODEL, + bb->pf_moe_sh, BB_SHARED_INTER, nullptr, EPI_SILU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, + nullptr, 0, 0, nullptr, nullptr, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.exp_down, bb->pf_moe_h_v, BB_EXPERTS * D_MODEL, BB_MOE_INTER, + nullptr, D_MODEL, nullptr, EPI_MOE_ACC, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + D_MODEL / OUT_ATOM, BB_MOE_INTER / KCHUNK, + bb->pf_moe_w, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + gemm(L.sh_down, bb->pf_moe_sh_v, D_MODEL, BB_SHARED_INTER, + nullptr, D_MODEL, nullptr, EPI_ACC_F32, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr, + 0, 0, nullptr, bb->pf_moe_out, n_tok, BATCH_PAD); + grid_barrier(g.barrier); + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_h[off + idx] = bb_moe_fold_deterministic(bb->pf_moe_out, idx); + grid_barrier(g.barrier); + } + if (bb->pf_layer_out && tok0 + T <= PREFILL_MAX_TOKENS) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; idx < T * D_MODEL; + idx += gridDim.x * THREADS) + bb->pf_layer_out[(size_t)layer * PREFILL_MAX_TOKENS * D_MODEL + + off + idx] = bb->pf_h[off + idx]; + grid_barrier(g.barrier); + } + } + } + bb_norm_residual(bb->pf_h, bb->pf_r, bb->norm_f_w, bb->pf_final, ln_scratch, + n_total); + grid_barrier(g.barrier); + if (!bb->pf_fuse_decode) return; + { + const int n_pos = n_total / n_seq; + const size_t last = (size_t)(n_pos - 1) * n_seq; + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) { + const int row = idx / D_MODEL; + const int i = idx - row * D_MODEL; + const size_t src = (last + (row < n_seq ? row : n_seq - 1)) * D_MODEL + i; + bb->h[idx] = bb->pf_h[src]; + bb->r[idx] = bb->pf_r[src]; + } + grid_barrier(g.barrier); + } + } + if (n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) + atomicAdd(bb->launches + 1, 1ull); + if (n_backbone_layers > 0) { + bb_norm_residual(bb->h, bb->r, bb->norm_f_w, g.x, ln_scratch); + grid_barrier(g.barrier); + if (bb->phoneme_out) { + gemm(bb->phoneme_w, g.x_v, BB_PHON_PADDED, D_MODEL, + bb->phoneme_logits, BB_PHON_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, bb->phoneme_b); + grid_barrier(g.barrier); + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + MK2_PROGRESS_GUARD(native_progress_mask, row); + if (threadIdx.x == 0) { + const bf16 *lg = bb->phoneme_logits + (size_t)row * BB_PHON_PADDED; + float best = -CUDART_INF_F; + int arg = 0; + for (int n = 0; n < BB_PHONEME; ++n) { + const float v = __bfloat162float(lg[n]); + if (v > best) { best = v; arg = n; } + } + bb->phoneme_out[row] = arg; + } + __syncthreads(); + } + grid_barrier(g.barrier); + } + if (bb->dec_snapshot) { + for (int idx = blockIdx.x * THREADS + threadIdx.x; + idx < BATCH * D_MODEL; idx += gridDim.x * THREADS) + bb->dec_snapshot[idx] = g.x[idx]; + grid_barrier(g.barrier); + } + } + __shared__ int native_run_predictor; + if (native_predictor_active_mask && bb && g.native_codes) { + if (threadIdx.x == 0) { + int any_active = 0; + for (int row = 0; row < BATCH; ++row) + any_active |= int(native_predictor_active_mask[row] != 0 && + MK2_PROGRESS_ENABLED(native_progress_mask, row)); + native_run_predictor = any_active; + } + __syncthreads(); + } + const bool skip_native_predictor = native_predictor_active_mask && bb && + g.native_codes && native_run_predictor == 0; + if (skip_native_predictor) { + for (int index = blockIdx.x * THREADS + threadIdx.x; + index < BATCH * N_POS; index += gridDim.x * THREADS) { + MK2_PROGRESS_GUARD(native_progress_mask, index / N_POS); + g.sampled[index] = 0; + g.native_codes[index] = 0; + } + grid_barrier(g.barrier); + } else { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) { + const size_t at = (size_t)row * D_MODEL + i; + g.x[at] = __float2bfloat16(__bfloat162float(g.x[at]) + + __bfloat162float(g.pos_emb[i])); + } + } + if constexpr (!(MK2_B1_LOCAL_BARRIERS && BATCH == 1)) + grid_barrier(g.barrier); + for (int pos = 0; pos < n_pos; ++pos) { + for (int layer = 0; layer < n_layers; ++layer) { +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_self[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.qkv[layer], g.normed_v, 3 * D_MODEL, D_MODEL, + g.qkv_out, 3 * D_MODEL, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + const int pairs = BATCH * N_HEADS; +#ifndef ABLATE_ATTN + { + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int warps_total = gridDim.x * WARPS; + const float scale = rsqrtf((float)D_HEAD); + for (int pair = blockIdx.x * WARPS + warp; pair < pairs; + pair += warps_total) { + const int row = pair / N_HEADS; + const int head = pair % N_HEADS; + const bf16 *qkv_row = g.qkv_out + (size_t)row * 3 * D_MODEL; + const bf16 *q = qkv_row + head * D_HEAD; + const bf16 *k_new = qkv_row + D_MODEL + head * D_HEAD; + const bf16 *v_new = qkv_row + 2 * D_MODEL + head * D_HEAD; + const size_t cache_base = + (((size_t)layer * BATCH + row) * N_HEADS + head) * N_POS * + D_HEAD; + constexpr int DPL = D_HEAD / 32; + float qv[DPL], kn[DPL], vn[DPL]; + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + qv[j] = __bfloat162float(q[i]); + kn[j] = __bfloat162float(k_new[i]); + vn[j] = __bfloat162float(v_new[i]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + const int i = lane + j * 32; + g.kcache[cache_base + pos * D_HEAD + i] = k_new[i]; + g.vcache[cache_base + pos * D_HEAD + i] = v_new[i]; + } + __syncwarp(); + float sc[N_POS]; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *kt = g.kcache + cache_base + t * D_HEAD; + float acc = 0.f; + #pragma unroll + for (int j = 0; j < D_HEAD / 32; ++j) { + const int i = lane + j * 32; + acc += qv[j] * __bfloat162float(kt[i]); + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[t] = acc * scale; + } + { + float acc = 0.f; + #pragma unroll + for (int j = 0; j < DPL; ++j) acc += qv[j] * kn[j]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + sc[pos] = acc * scale; + } + float m = sc[0]; + for (int t = 1; t <= pos; ++t) m = fmaxf(m, sc[t]); + float s = 0.f; + for (int t = 0; t <= pos; ++t) { + sc[t] = __expf(sc[t] - m); + s += sc[t]; + } + const float inv = 1.f / s; + for (int t = 0; t <= pos; ++t) sc[t] *= inv; + float ov[DPL] = {}; + #pragma unroll 4 + for (int t = 0; t < pos; ++t) { + const bf16 *vt = g.vcache + cache_base + t * D_HEAD; + const float s = sc[t]; + #pragma unroll + for (int j = 0; j < DPL; ++j) + ov[j] += s * __bfloat162float(vt[lane + j * 32]); + } + #pragma unroll + for (int j = 0; j < DPL; ++j) { + ov[j] += sc[pos] * vn[j]; + g.attn_out[(size_t)row * D_MODEL + head * D_HEAD + + lane + j * 32] = __float2bfloat16(ov[j]); + } + } + } +#endif + grid_barrier(g.barrier); + gemm(g.o[layer], g.attn_v, D_MODEL, D_MODEL, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); +#ifndef ABLATE_NORM + layernorm_rows(g.x, g.norm_ff[layer], g.normed, ln_scratch); +#endif + grid_barrier(g.barrier); + gemm(g.ff_in[layer], g.normed_v, D_FFN, D_MODEL, + g.ff_hidden, D_FFN, nullptr, EPI_GELU, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + gemm(g.ff_out[layer], g.ffh_v, D_MODEL, D_FFN, + nullptr, D_MODEL, g.x, EPI_RESIDUAL, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile); + grid_barrier(g.barrier); + } + if constexpr (!MK2_SKIP_HIDDEN_SNAPSHOT) + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.hidden_rows[((size_t)pos * BATCH + row) * D_MODEL + i] = + g.x[(size_t)row * D_MODEL + i]; + } +#ifdef ABLATE_HEAD + (void)0; +#elif defined(HEAD_BF16) + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_STORE, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, nullptr); +#else + gemm(g.head[pos], g.x_v, VOCAB_PADDED, D_MODEL, + g.logits, VOCAB_PADDED, nullptr, EPI_BIAS, + w_s, a_s, o_s, filled, drained, o_sf, gemm_iter, gemm_tile, g.head_bias[pos]); +#endif + grid_barrier(g.barrier); +#ifndef ABLATE_SAMPLE + sample_codes(g.logits, g.head_bias[pos], g.forbidden, g.gumbel, pos, + g.sampled_pos, sort_storage, g.temperature); +#else + if (threadIdx.x < BATCH && blockIdx.x == 0) + g.sampled_pos[threadIdx.x] = g.codes[threadIdx.x * N_POS + pos]; +#endif +#ifdef ABLATE_SAMPLE + if constexpr (MK2_B1_LOCAL_BARRIERS && BATCH == 1) { + __syncthreads(); + } else { + grid_barrier(g.barrier); + } +#endif + if (pos + 1 < n_pos) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + const bf16 *emb = g.audio_emb[pos] + (size_t)code * D_MODEL; + const bf16 *pe = g.pos_emb + (pos + 1) * D_MODEL; + for (int i = threadIdx.x; i < D_MODEL; i += THREADS) + g.x[(size_t)row * D_MODEL + i] = + __float2bfloat16(__bfloat162float(emb[i]) + + __bfloat162float(pe[i])); + } + } + if (threadIdx.x == 0) { + for (int row = blockIdx.x; row < BATCH; row += gridDim.x) { + const int code = g.sampled_pos[row]; + g.sampled[row * N_POS + pos] = code; + if (g.native_codes && MK2_PROGRESS_ENABLED(native_progress_mask, row)) { +#if MK2_INFLIGHT + g.native_codes[row * N_POS + pos] = + native_predictor_active_mask && !native_predictor_active_mask[row] + ? 0 : (long long)code; +#else + g.native_codes[row * N_POS + pos] = (long long)code; +#endif + } + } + } +#ifdef ABLATE_SAMPLE + grid_barrier(g.barrier); +#else + if (pos + 1 == n_pos) grid_barrier(g.barrier); +#endif + } + } + if (g.native_seq_len && blockIdx.x == 0 && threadIdx.x < BATCH && + MK2_PROGRESS_ENABLED(native_progress_mask, threadIdx.x)) + g.native_seq_len[threadIdx.x] += 1; + if (bb && n_backbone_layers > 0 && bb->launches && blockIdx.x == 0 && + threadIdx.x == 0) { + unsigned long long t1; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); + const unsigned long long t0 = bb->launches[2]; + if (t1 > t0) { + const unsigned long long dt = t1 - t0; + atomicAdd(bb->launches + 3, dt); + atomicAdd(bb->launches + 4, 1ull); + atomicMax(bb->launches + 5, dt); + } + } +} +#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \ + printf("FAIL %s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); exit(1);} } while (0) +__global__ void fp16_to_bf16(const __half *src, bf16 *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); +} +__global__ void fp16_to_bf16_masked(const __half *src, bf16 *dst, int n, + const unsigned char *audio_valid, int *run_predicate) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2bfloat16(__half2float(src[i])); + if (i == 0) { + int run = 0; + for (int row = 0; row < BATCH; ++row) run |= audio_valid[row] != 0; + *run_predicate = run; + } +} +__global__ void codes_to_int64(const int *src, long long *dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = (long long)src[i]; +} +static std::vector read_file(const std::string &path, size_t expected) { + FILE *f = fopen(path.c_str(), "rb"); + if (!f) { printf("FAIL cannot open %s\n", path.c_str()); exit(1); } + fseek(f, 0, SEEK_END); + const size_t bytes = ftell(f); + fseek(f, 0, SEEK_SET); + if (expected && bytes != expected) { + printf("FAIL %s is %zu bytes, expected %zu\n", path.c_str(), bytes, expected); + exit(1); + } + std::vector data(bytes); + if (fread(data.data(), 1, bytes, f) != bytes) { printf("FAIL short read\n"); exit(1); } + fclose(f); + return data; +} +template +static T *upload(const std::string &path, size_t count) { + auto host = read_file(path, count * sizeof(T)); + T *device = nullptr; + CHECK(cudaMalloc(&device, host.size())); + CHECK(cudaMemcpy(device, host.data(), host.size(), cudaMemcpyHostToDevice)); + return device; +} +#ifndef MK2_NO_MAIN +int main(int argc, char **argv) { + setbuf(stdout, nullptr); + const std::string root = argc > 1 ? argv[1] : "/out/predictor_ref_b32"; + bool bench = false, verify = true; + const char *dump_rows = nullptr; + int n_backbone = 0; + int past_len = 7; + const char *dump_moe_w = nullptr; + bool inject_moe = false; + bool stage0 = false; + bool shuffle_slots = false; + bool router_audit = false; + std::string prefill_dir; + int prefill_bench = 0; + int prefill_seq = 1; + int prefill_fuse = 0; + int bench_iters = 200; + bool poison_cur_kv = false; + int decode_steps = 1; + const char *dump_final = nullptr; + int n_layers = N_LAYERS, n_pos = N_POS; + for (int i = 2; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--no-verify") verify = false; + else if (a == "--dump-rows" && i + 1 < argc) dump_rows = argv[++i]; + else if (a == "--backbone-layers" && i + 1 < argc) n_backbone = atoi(argv[++i]); + else if (a == "--past-len" && i + 1 < argc) past_len = atoi(argv[++i]); + else if (a == "--dump-moe-w" && i + 1 < argc) dump_moe_w = argv[++i]; + else if (a == "--inject-moe-w") inject_moe = true; + else if (a == "--stage0") stage0 = true; + else if (a == "--shuffle-slots") shuffle_slots = true; + else if (a == "--router-audit") router_audit = true; + else if (a == "--prefill" && i + 1 < argc) prefill_dir = argv[++i]; + else if (a == "--prefill-bench" && i + 1 < argc) prefill_bench = atoi(argv[++i]); + else if (a == "--prefill-seq" && i + 1 < argc) prefill_seq = atoi(argv[++i]); + else if (a == "--prefill-fuse") prefill_fuse = 1; + else if (a == "--bench-iters" && i + 1 < argc) bench_iters = atoi(argv[++i]); + else if (a == "--poison-cur-kv") poison_cur_kv = true; + else if (a == "--decode-steps" && i + 1 < argc) decode_steps = atoi(argv[++i]); + else if (a == "--dump-final" && i + 1 < argc) dump_final = argv[++i]; + else if (a == "--layers" && i + 1 < argc) n_layers = atoi(argv[++i]); + else if (a == "--positions" && i + 1 < argc) n_pos = atoi(argv[++i]); + } + bf16 *w_qkv[N_LAYERS], *w_o[N_LAYERS], *w_ffin[N_LAYERS], *w_ffout[N_LAYERS]; + const float *ns[N_LAYERS], *nf[N_LAYERS]; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w_qkv[l] = upload(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w_o[l] = upload(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w_ffin[l] = upload(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w_ffout[l] = upload(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + ns[l] = upload(p + "norm_self.bin", D_MODEL); + nf[l] = upload(p + "norm_ff.bin", D_MODEL); + } + const bf16 *pos_emb = upload(root + "/position_embeddings.bin", + (size_t)18 * D_MODEL); + const bf16 *ae[N_POS]; + for (int k = 0; k < N_POS; ++k) + ae[k] = upload(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + bf16 *x = nullptr; + if (stage0) { + CHECK(cudaMalloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(x, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + } else { + x = upload(root + "/dec_hidden.bin", (size_t)BATCH * D_MODEL); + } + bf16 *x_initial = nullptr; + CHECK(cudaMalloc(&x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(x_initial, x, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + const int *codes = upload(root + "/codes.bin", (size_t)BATCH * N_POS); + bf16 *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + unsigned int *barrier; + CHECK(cudaMalloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + CHECK(cudaMalloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + CHECK(cudaMalloc(&kcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&vcache, cache_elems * sizeof(bf16))); + CHECK(cudaMalloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&barrier, 64 * sizeof(unsigned int))); + CHECK(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + auto raw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + CHECK(cudaMalloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(head_w[k], raw.data(), raw.size(), cudaMemcpyHostToDevice)); + auto rb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + bf16 *b_dev = nullptr; + CHECK(cudaMalloc(&b_dev, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemset(b_dev, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(b_dev, rb.data(), rb.size(), cudaMemcpyHostToDevice)); + head_b[k] = b_dev; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[1025] = 0; + int8_t *forbidden = nullptr; + CHECK(cudaMalloc(&forbidden, forb.size())); + CHECK(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + const float *gumbel = upload(root + "/gumbel.bin", + (size_t)BATCH * N_POS * TOP_K); + bf16 *logits_buf = nullptr; + int *sampled = nullptr, *sampled_pos = nullptr; + CHECK(cudaMalloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + CHECK(cudaMalloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + Globals g{ + {weight_gl{w_qkv[0], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[1], nullptr, nullptr, QKV, DM}, + weight_gl{w_qkv[2], nullptr, nullptr, QKV, DM}}, + {weight_gl{w_o[0], nullptr, nullptr, DM, DM}, + weight_gl{w_o[1], nullptr, nullptr, DM, DM}, + weight_gl{w_o[2], nullptr, nullptr, DM, DM}}, + {weight_gl{w_ffin[0], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[1], nullptr, nullptr, DF, DM}, + weight_gl{w_ffin[2], nullptr, nullptr, DF, DM}}, + {weight_gl{w_ffout[0], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[1], nullptr, nullptr, DM, DF}, + weight_gl{w_ffout[2], nullptr, nullptr, DM, DF}}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, gumbel, sampled, sampled_pos, 0.7f, + {ns[0], ns[1], ns[2]}, + {nf[0], nf[1], nf[2]}, + pos_emb, + {ae[0], ae[1], ae[2], ae[3], ae[4], ae[5], ae[6], ae[7], + ae[8], ae[9], ae[10], ae[11], ae[12], ae[13], ae[14], ae[15]}, + kcache, vcache, hidden_rows, codes, barrier, + }; + Backbone *d_bb = nullptr; + bf16 *bb_layer_out_host = nullptr; + float *bb_moe_w_host = nullptr; + bf16 *bb_dec_snapshot_host = nullptr; + int *ph_out_verify = nullptr; + float *moe_w_audit_host = nullptr; + unsigned long long *launch_counter_host = nullptr; + Backbone *bb_saved = nullptr; + std::vector mamba_saved; + bf16 *pf_layer_out_host = nullptr; + if (n_backbone > 0) { + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + static const char PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + std::vector mamba_host; + int mamba_of_layer[BB_LAYERS]; + const bf16 *norm_w[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + norm_w[layer] = upload(lp + "norm.bin", D_MODEL); + if (PATTERN[layer] != 'M') { mamba_of_layer[layer] = -1; continue; } + mamba_of_layer[layer] = (int)mamba_host.size(); + bf16 *in_p = upload(lp + "in_proj.bin", + (size_t)MB_IN_PROJ * D_MODEL); + bf16 *out_p = upload(lp + "out_proj.bin", + (size_t)D_MODEL * D_MODEL); + bf16 *conv_state = nullptr; + float *ssm_state = nullptr; + const size_t cs = (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss = (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE; + CHECK(cudaMalloc(&conv_state, cs * sizeof(bf16))); + CHECK(cudaMalloc(&ssm_state, ss * sizeof(float))); + auto csr = read_file(bb_root + "/conv_state_in_" + + std::to_string(layer) + ".bin", cs * sizeof(bf16)); + auto ssr = read_file(bb_root + "/ssm_state_in_" + + std::to_string(layer) + ".bin", ss * sizeof(float)); + const size_t cs_per = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + const size_t ss_per = (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE; + std::vector csp(csr.size()), ssp(ssr.size()); + for (int i = 0; i < BATCH; ++i) { + const int slot = shuffle_slots ? BATCH - 1 - i : i; + memcpy(csp.data() + (size_t)slot * cs_per * sizeof(bf16), + csr.data() + (size_t)i * cs_per * sizeof(bf16), + cs_per * sizeof(bf16)); + memcpy(ssp.data() + (size_t)slot * ss_per * sizeof(float), + ssr.data() + (size_t)i * ss_per * sizeof(float), + ss_per * sizeof(float)); + } + CHECK(cudaMemcpy(conv_state, csp.data(), csp.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(ssm_state, ssp.data(), ssp.size(), + cudaMemcpyHostToDevice)); + mamba_host.push_back(MambaLayer{ + weight_gl{in_p, nullptr, nullptr, (size_t)MB_IN_PROJ, + (size_t)D_MODEL}, + weight_gl{out_p, nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + upload(lp + "conv_w.bin", (size_t)MB_CONV_DIM * MB_CONV_K), + upload(lp + "conv_b.bin", MB_CONV_DIM), + upload(lp + "A.bin", MB_HEADS), + upload(lp + "D.bin", MB_HEADS), + upload(lp + "dt_bias.bin", MB_HEADS), + upload(lp + "gnorm.bin", MB_INNER), + conv_state, + MB_CONV_DIM * (MB_CONV_K - 1), + ssm_state, + MB_HEADS * MB_HEAD_DIM * MB_STATE, + }); + } + std::vector attn_host; + std::vector moe_host; + int attn_of_layer[BB_LAYERS], moe_of_layer[BB_LAYERS]; + for (int layer = 0; layer < BB_LAYERS; ++layer) { + attn_of_layer[layer] = -1; + moe_of_layer[layer] = -1; + const std::string lp = bb_root + "/l" + std::to_string(layer) + "_"; + if (PATTERN[layer] == '*') { + attn_of_layer[layer] = (int)attn_host.size(); + const int TPB = 32; + const int kv_span = max(BB_T_MAX, PREFILL_MAX_TOKENS); + const int blocks_per_seq = (kv_span + TPB - 1) / TPB; + const size_t elems_per_block = + (size_t)TPB * BB_KV_HEADS * BB_HEAD_DIM; + const size_t kvn = (size_t)BATCH * BB_T_MAX * BB_KV_SIZE; + auto kr = read_file(bb_root + "/k_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + auto vr = read_file(bb_root + "/v_cache_" + std::to_string(layer) + + ".bin", kvn * sizeof(bf16)); + const bf16 *kh = reinterpret_cast(kr.data()); + const bf16 *vh = reinterpret_cast(vr.data()); + const int total_blocks = BATCH * blocks_per_seq * 2; + std::vector pool(total_blocks * elems_per_block, __float2bfloat16(0.f)); + std::vector table((size_t)BATCH * 2 * blocks_per_seq, 0); + for (int s = 0; s < BATCH; ++s) { + for (int b = 0; b < blocks_per_seq; ++b) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = BATCH * blocks_per_seq + s * blocks_per_seq + b; + table[((size_t)s * 2 + 0) * blocks_per_seq + b] = kblk; + table[((size_t)s * 2 + 1) * blocks_per_seq + b] = vblk; + for (int p = b * TPB; p < BB_T_MAX && p < (b + 1) * TPB; ++p) { + for (int h = 0; h < BB_KV_HEADS; ++h) { + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t src = + (((size_t)s * BB_T_MAX + p) * BB_KV_HEADS + h) * + BB_HEAD_DIM + d; + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(p % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = kh[src]; + pool[(size_t)vblk * elems_per_block + off] = vh[src]; + } + } + } + } + } + if (poison_cur_kv && past_len >= 0 && past_len < kv_span) { + const int b = past_len / TPB; + for (int s = 0; s < BATCH; ++s) { + const int kblk = (BATCH - 1 - s) * blocks_per_seq + b; + const int vblk = + BATCH * blocks_per_seq + s * blocks_per_seq + b; + for (int h = 0; h < BB_KV_HEADS; ++h) + for (int d = 0; d < BB_HEAD_DIM; ++d) { + const size_t off = + (size_t)h * TPB * BB_HEAD_DIM + + (size_t)(past_len % TPB) * BB_HEAD_DIM + d; + pool[(size_t)kblk * elems_per_block + off] = + __float2bfloat16(7.0f); + pool[(size_t)vblk * elems_per_block + off] = + __float2bfloat16(-13.0f); + } + } + } + bf16 *d_pool = nullptr; + int *d_table = nullptr; + CHECK(cudaMalloc(&d_pool, pool.size() * sizeof(bf16))); + CHECK(cudaMemcpy(d_pool, pool.data(), pool.size() * sizeof(bf16), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_table, table.size() * sizeof(int))); + CHECK(cudaMemcpy(d_table, table.data(), table.size() * sizeof(int), + cudaMemcpyHostToDevice)); + attn_host.push_back(AttnLayer{ + weight_gl{upload(lp + "qkv.bin", + (size_t)BB_QKV * D_MODEL), + nullptr, nullptr, (size_t)BB_QKV, (size_t)D_MODEL}, + weight_gl{upload(lp + "o_proj.bin", + (size_t)D_MODEL * D_MODEL), + nullptr, nullptr, (size_t)D_MODEL, (size_t)D_MODEL}, + d_pool, d_table, blocks_per_seq, TPB, + }); + } else if (PATTERN[layer] == 'E') { + moe_of_layer[layer] = (int)moe_host.size(); + moe_host.push_back(MoeLayer{ + weight_gl{upload(lp + "exp_up.bin", + (size_t)BB_EXPERTS * BB_MOE_INTER * D_MODEL), + nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{upload(lp + "exp_down.bin", + (size_t)BB_EXPERTS * D_MODEL * BB_MOE_INTER), + nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{upload(lp + "sh_up.bin", + (size_t)BB_SHARED_INTER * D_MODEL), + nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{upload(lp + "sh_down.bin", + (size_t)D_MODEL * BB_SHARED_INTER), + nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + upload(lp + "gate.bin", (size_t)BB_EXPERTS * D_MODEL), + upload(lp + "gate_bias.bin", BB_EXPERTS), + }); + } + } + AttnLayer *d_attn = nullptr; + MoeLayer *d_moe = nullptr; + CHECK(cudaMalloc(&d_attn, attn_host.size() * sizeof(AttnLayer))); + CHECK(cudaMemcpy(d_attn, attn_host.data(), + attn_host.size() * sizeof(AttnLayer), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&d_moe, moe_host.size() * sizeof(MoeLayer))); + CHECK(cudaMemcpy(d_moe, moe_host.data(), moe_host.size() * sizeof(MoeLayer), + cudaMemcpyHostToDevice)); + MambaLayer *d_mamba = nullptr; + CHECK(cudaMalloc(&d_mamba, mamba_host.size() * sizeof(MambaLayer))); + CHECK(cudaMemcpy(d_mamba, mamba_host.data(), + mamba_host.size() * sizeof(MambaLayer), + cudaMemcpyHostToDevice)); + bf16 *bb_h = upload(bb_root + "/hidden_in.bin", + (size_t)BATCH * D_MODEL); + bf16 *bb_r, *bb_normed, *bb_zxbcdt, *bb_gated, *bb_layer_out; + CHECK(cudaMalloc(&bb_r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(bb_r, 0, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&bb_zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&bb_gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_layer_out, + (size_t)BB_LAYERS * BATCH * D_MODEL * sizeof(bf16))); + bb_layer_out_host = bb_layer_out; + bf16 *ph_w_dev = nullptr, *ph_b_dev_nc = nullptr, *ph_logits_dev = nullptr; + int *ph_out_dev = nullptr; + { + auto wr = read_file(bb_root + "/phoneme_w.bin", + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16)); + auto br = read_file(bb_root + "/phoneme_b.bin", + (size_t)BB_PHONEME * sizeof(bf16)); + CHECK(cudaMalloc(&ph_w_dev, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemset(ph_w_dev, 0, + (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + CHECK(cudaMemcpy(ph_w_dev, wr.data(), wr.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_b_dev_nc, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemset(ph_b_dev_nc, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMemcpy(ph_b_dev_nc, br.data(), br.size(), + cudaMemcpyHostToDevice)); + CHECK(cudaMalloc(&ph_logits_dev, + (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + CHECK(cudaMalloc(&ph_out_dev, (size_t)BATCH * sizeof(int))); + } + const bf16 *ph_b_dev = ph_b_dev_nc; + bf16 *bb_qkv, *bb_moe_h, *bb_moe_sh; + float *bb_moe_out, *bb_moe_w; + CHECK(cudaMalloc(&bb_qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_h, + (size_t)BATCH * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_sh, + (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&bb_moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&bb_moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + bb_moe_w_host = bb_moe_w; + bf16 *pf_h, *pf_r, *pf_normed, *pf_zxbcdt, *pf_xbc, *pf_gated_b, *pf_qkv_b; + bf16 *pf_moe_h_b, *pf_moe_sh_b, *pf_layer_out_b, *pf_final_b; + float *pf_moe_out_b, *pf_moe_w_b; + { + const size_t P = PREFILL_STORAGE_TOKENS; + const size_t T1 = PREFILL_STORAGE_TOKENS; + CHECK(cudaMalloc(&pf_h, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_r, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_final_b, P * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_normed, T1 * D_MODEL * sizeof(bf16))); + CHECK(cudaMalloc(&pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16))); + CHECK(cudaMalloc(&pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16))); + CHECK(cudaMalloc(&pf_gated_b, T1 * MB_INNER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_qkv_b, T1 * BB_QKV * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_h_b, + T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16))); + CHECK(cudaMalloc(&pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + CHECK(cudaMalloc(&pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float))); + CHECK(cudaMalloc(&pf_layer_out_b, + (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16))); + for (auto pr : std::vector>{ + {pf_h, P * D_MODEL * sizeof(bf16)}, + {pf_r, P * D_MODEL * sizeof(bf16)}, + {pf_final_b, P * D_MODEL * sizeof(bf16)}, + {pf_normed, T1 * D_MODEL * sizeof(bf16)}, + {pf_zxbcdt, T1 * MB_IN_PROJ * sizeof(bf16)}, + {pf_xbc, T1 * MB_CONV_DIM * sizeof(bf16)}, + {pf_gated_b, T1 * MB_INNER * sizeof(bf16)}, + {pf_qkv_b, T1 * BB_QKV * sizeof(bf16)}, + {pf_moe_h_b, T1 * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16)}, + {pf_moe_sh_b, T1 * BB_SHARED_INTER * sizeof(bf16)}, + {pf_moe_out_b, T1 * (BB_EXPERTS + 1) * D_MODEL * sizeof(float)}, + {pf_moe_w_b, T1 * BB_EXPERTS * sizeof(float)}, + {pf_layer_out_b, (size_t)BB_LAYERS * P * D_MODEL * sizeof(bf16)}}) + CHECK(cudaMemset(pr.first, 0, pr.second)); + } + unsigned long long *bb_launches = nullptr; + CHECK(cudaMalloc(&bb_launches, 6 * sizeof(unsigned long long))); + CHECK(cudaMemset(bb_launches, 0, 6 * sizeof(unsigned long long))); + launch_counter_host = bb_launches; + Backbone host_bb{ + act_gl{bb_normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{bb_gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + d_mamba, + upload(bb_root + "/norm_f.bin", D_MODEL), + d_attn, + d_moe, + {}, + {}, {}, {}, + past_len, + bb_h, bb_r, bb_normed, bb_zxbcdt, bb_gated, bb_layer_out, + bb_qkv, bb_moe_h, bb_moe_sh, bb_moe_out, bb_moe_w, + inject_moe ? upload(bb_root + "/moe_w.bin", + (size_t)14 * BATCH * BB_EXPERTS) + : nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 0, 0, 1, 0, + nullptr, + nullptr, + pf_h, pf_r, pf_normed, pf_zxbcdt, pf_xbc, pf_gated_b, pf_qkv_b, + pf_moe_h_b, pf_moe_sh_b, pf_moe_out_b, pf_moe_w_b, pf_layer_out_b, + pf_final_b, + act_gl{pf_normed, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)D_MODEL}, + act_gl{pf_gated_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)MB_INNER}, + act_gl{pf_qkv_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_QKV}, + act_gl{pf_moe_h_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{pf_moe_sh_b, nullptr, nullptr, (size_t)PREFILL_STORAGE_TOKENS, + (size_t)BB_SHARED_INTER}, + weight_gl{ph_w_dev, nullptr, nullptr, (size_t)BB_PHON_PADDED, + (size_t)D_MODEL}, + ph_b_dev, ph_logits_dev, ph_out_dev, + bb_launches, + nullptr, + act_gl{bb_qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{bb_moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{bb_moe_sh, nullptr, nullptr, (size_t)BATCH, + (size_t)BB_SHARED_INTER}, + }; + CHECK(cudaMalloc(&bb_dec_snapshot_host, + (size_t)BATCH * D_MODEL * sizeof(bf16))); + host_bb.dec_snapshot = bb_dec_snapshot_host; + bf16 *bb_xbc = nullptr; + CHECK(cudaMalloc(&bb_xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + host_bb.xbc_act = bb_xbc; + std::vector slots_host(BATCH); + for (int i = 0; i < BATCH; ++i) slots_host[i] = i; + if (shuffle_slots) { + for (int i = 0; i < BATCH; ++i) slots_host[i] = BATCH - 1 - i; + } + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, BATCH * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots_host.data(), BATCH * sizeof(int), + cudaMemcpyHostToDevice)); + host_bb.slot_map = d_slots; + ph_out_verify = ph_out_dev; + if (router_audit) { + CHECK(cudaMalloc(&moe_w_audit_host, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + CHECK(cudaMemset(moe_w_audit_host, 0, + (size_t)14 * BATCH * BB_EXPERTS * sizeof(float))); + host_bb.moe_w_audit = moe_w_audit_host; + } + for (int i = 0; i < BB_LAYERS; ++i) { + host_bb.norm_w[i] = norm_w[i]; + host_bb.mamba_of_layer[i] = mamba_of_layer[i]; + host_bb.attn_of_layer[i] = attn_of_layer[i]; + host_bb.moe_of_layer[i] = moe_of_layer[i]; + } + CHECK(cudaMalloc(&d_bb, sizeof(Backbone))); + CHECK(cudaMemcpy(d_bb, &host_bb, sizeof(Backbone), cudaMemcpyHostToDevice)); + bb_saved = new Backbone(host_bb); + mamba_saved = mamba_host; + pf_layer_out_host = pf_layer_out_b; + printf(" backbone: %zu Mamba2 layers loaded, %d layers requested\n", + mamba_host.size(), n_backbone); + } + const int smem = STAGES * (sizeof(weight_tile) + sizeof(act_tile)) + + sizeof(out_tile) + sizeof(out_tile_f) + 4096; + CHECK(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + int max_blocks = 0, sms = 0; + CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_blocks, predictor_mk2, THREADS, smem)); + CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0)); + if (max_blocks < 1) { + printf("FAIL no block is resident per SM at %d B shared\n", smem); + return 1; + } + int grid = min(sms * max_blocks, D_FFN / OUT_ATOM); + for (int i = 2; i < argc; ++i) + if (std::string(argv[i]) == "--grid" && i + 1 < argc) grid = atoi(argv[i + 1]); + printf(" grid=%d blocks x %d threads smem=%d B KCHUNK=%d STAGES=%d" + " (%d resident/SM, %d SMs)\n", + grid, THREADS, smem, (int)KCHUNK, (int)STAGES, max_blocks, sms); + if (prefill_bench > 0) { + if (d_bb == nullptr) { + printf(" --prefill-bench needs --backbone-layers > 0\n"); + return 1; + } + const int T = prefill_bench; + if (T > PREFILL_MAX_TOKENS) { + printf(" T=%d exceeds PREFILL_MAX_TOKENS=%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + std::vector prompt_h((size_t)T * D_MODEL); + for (size_t i = 0; i < prompt_h.size(); ++i) + prompt_h[i] = 0x3D00 | (uint16_t)(i & 0xFF); + bf16 *d_prompt = nullptr; + CHECK(cudaMalloc(&d_prompt, prompt_h.size() * sizeof(uint16_t))); + CHECK(cudaMemcpy(d_prompt, prompt_h.data(), + prompt_h.size() * sizeof(uint16_t), cudaMemcpyHostToDevice)); + if (prefill_seq > 1 && T % prefill_seq) { + printf(" T=%d is not a multiple of --prefill-seq %d\n", T, prefill_seq); + return 1; + } + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.pf_n_seq = prefill_seq; + hb.pf_fuse_decode = prefill_fuse; + hb.prompt = d_prompt; + if (prefill_seq > 1) { + std::vector slots(prefill_seq); + for (int i = 0; i < prefill_seq; ++i) slots[i] = i; + int *d_slots = nullptr; + CHECK(cudaMalloc(&d_slots, slots.size() * sizeof(int))); + CHECK(cudaMemcpy(d_slots, slots.data(), slots.size() * sizeof(int), + cudaMemcpyHostToDevice)); + hb.pf_seq_slot = d_slots; + } + hb.pf_layer_out = nullptr; + hb.moe_w_inject = nullptr; + hb.moe_w_audit = nullptr; + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + const int warm = 3, iters = 10; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const int tiles = (T + BATCH_PAD - 1) / BATCH_PAD; + printf(" BENCH T=%d, %d layers, %d tiles, %d seq%s: %.3f ms/launch" + " (%.3f ms/tile)\n", + T, n_backbone, tiles, prefill_seq, + prefill_fuse ? ", FUSED prefill+decode" : "", ms / iters, + ms / iters / tiles); + if (prefill_fuse) { + auto fresh_state = [&]() { + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state, 0, + (size_t)BATCH * MB_CONV_DIM * (MB_CONV_K - 1) * + sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state, 0, + (size_t)BATCH * MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + }; + std::vector fused_codes((size_t)BATCH * N_POS); + std::vector split_codes((size_t)BATCH * N_POS); + fresh_state(); + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(fused_codes.data(), (const void *)g.codes, + fused_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + fresh_state(); + Backbone pb = hb; + pb.pf_fuse_decode = 0; + CHECK(cudaMemcpy(d_bb, &pb, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + const int n_pos_pf = T / prefill_seq; + for (int row = 0; row < BATCH; ++row) { + const int s = row < prefill_seq ? row : prefill_seq - 1; + const size_t src = ((size_t)(n_pos_pf - 1) * prefill_seq + s) * D_MODEL; + CHECK(cudaMemcpy(pb.h + (size_t)row * D_MODEL, pb.pf_h + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + CHECK(cudaMemcpy(pb.r + (size_t)row * D_MODEL, pb.pf_r + src, + D_MODEL * sizeof(bf16), cudaMemcpyDeviceToDevice)); + } + Backbone db = pb; + db.n_prefill_tokens = 0; + CHECK(cudaMemcpy(d_bb, &db, sizeof(Backbone), cudaMemcpyHostToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + CHECK(cudaMemcpy(split_codes.data(), (const void *)g.codes, + split_codes.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + int diff = 0; + for (size_t i = 0; i < fused_codes.size(); ++i) + if (fused_codes[i] != split_codes[i]) ++diff; + printf(" %s fused vs split: %d of %zu codes differ\n", + diff ? "FAIL" : "PASS", diff, fused_codes.size()); + if (diff) return 1; + } + return 0; + } + if (!prefill_dir.empty()) { + if (d_bb == nullptr) { + printf(" --prefill needs --backbone-layers > 0\n"); + return 1; + } + auto meta = read_file(prefill_dir + "/meta.json", 0); + const std::string ms(meta.begin(), meta.end()); + const size_t at = ms.find("\"tokens\""); + const int T = at == std::string::npos ? 0 : atoi(ms.c_str() + ms.find(':', at) + 1); + if (T <= 0 || T > PREFILL_MAX_TOKENS) { + printf(" fixture tokens=%d outside 1..%d\n", T, PREFILL_MAX_TOKENS); + return 1; + } + bf16 *d_prompt = upload(prefill_dir + "/prompt.bin", + (size_t)T * D_MODEL); + Backbone hb = *bb_saved; + hb.n_prefill_tokens = T; + hb.prefill_slot = 0; + hb.prompt = d_prompt; + if (inject_moe) { + const int n_moe = 14; + hb.moe_w_inject = upload( + prefill_dir + "/moe_w.bin", (size_t)n_moe * T * BB_EXPERTS); + printf(" routing: oracle-injected (isolates the MoE math from near-tie " + "flips)\n"); + } else { + hb.moe_w_inject = nullptr; + printf(" routing: kernel's own router\n"); + } + float *pf_audit = nullptr; + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + CHECK(cudaMalloc(&pf_audit, n * sizeof(float))); + CHECK(cudaMemset(pf_audit, 0, n * sizeof(float))); + hb.moe_w_audit = pf_audit; + } + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * sizeof(float))); + } + CHECK(cudaMemcpy(d_bb, &hb, sizeof(Backbone), cudaMemcpyHostToDevice)); + { + std::vector cs(MB_CONV_DIM * (MB_CONV_K - 1)); + CHECK(cudaMemcpy(cs.data(), + mamba_saved[0].conv_state + + (size_t)hb.prefill_slot * mamba_saved[0].conv_slot_stride, + cs.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + int nz = 0; + for (uint16_t v : cs) if (v) ++nz; + if (nz) { + printf(" FAIL conv_state[layer0] slot %d has %d of %zu non-zero\n", + hb.prefill_slot, nz, cs.size()); + return 1; + } + } + static const char PF_PATTERN[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + printf(" prefill: T=%d, %d chunks of %d, %d layers\n", + T, (T + BATCH_PAD - 1) / BATCH_PAD, BATCH_PAD, n_backbone); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + { + const int warm = 3, iters = 20; + for (int i = 0; i < warm; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaDeviceSynchronize()); + cudaEvent_t e0, e1; + CHECK(cudaEventCreate(&e0)); + CHECK(cudaEventCreate(&e1)); + CHECK(cudaEventRecord(e0)); + for (int i = 0; i < iters; ++i) + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, + n_backbone); + CHECK(cudaEventRecord(e1)); + CHECK(cudaDeviceSynchronize()); + float ms = 0.f; + CHECK(cudaEventElapsedTime(&ms, e0, e1)); + const double per = ms / iters; + printf(" TIMING prefill T=%d, %d layers: %.3f ms/launch" + " (%.3f ms/chunk over %d chunks)\n", + T, n_backbone, per, + per / ((T + BATCH_PAD - 1) / BATCH_PAD), + (T + BATCH_PAD - 1) / BATCH_PAD); + for (const MambaLayer &L : mamba_saved) { + CHECK(cudaMemset(L.conv_state + + (size_t)hb.prefill_slot * L.conv_slot_stride, + 0, + (size_t)MB_CONV_DIM * (MB_CONV_K - 1) * sizeof(bf16))); + CHECK(cudaMemset(L.ssm_state + + (size_t)hb.prefill_slot * L.ssm_slot_stride, + 0, + (size_t)MB_HEADS * MB_HEAD_DIM * MB_STATE * + sizeof(float))); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + std::vector got((size_t)BB_LAYERS * PREFILL_MAX_TOKENS * D_MODEL); + CHECK(cudaMemcpy(got.data(), pf_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + int bad_layers = 0; + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d.bin", layer); + auto want_raw = read_file(prefill_dir + name, + (size_t)T * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_cos = 1.0; + int worst_tok = -1; + int bad_tokens = 0; + for (int tok = 0; tok < T; ++tok) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < D_MODEL; ++i) { + const double a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + tok) * D_MODEL + i]); + const double b = want[(size_t)tok * D_MODEL + i]; + dot += a * b; na += a * a; nb += b * b; + } + const double cos = dot / (sqrt(na) * sqrt(nb) + 1e-30); + if (cos < worst_cos) { worst_cos = cos; worst_tok = tok; } + if (cos < 0.995) ++bad_tokens; + } + const bool ok = worst_cos > 0.995; + if (!ok && bad_layers == 0) { + printf(" kernel[0..5]:"); + for (int i = 0; i < 6; ++i) + printf(" %+8.4f", __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i])); + printf("\n oracle[0..5]:"); + for (int i = 0; i < 6; ++i) printf(" %+8.4f", want[i]); + printf("\n ratio :"); + for (int i = 0; i < 6; ++i) { + const float a = __bfloat162float( + got[((size_t)layer * PREFILL_MAX_TOKENS + 0) * D_MODEL + i]); + printf(" %+8.4f", want[i] == 0.f ? 0.f : a / want[i]); + } + printf("\n"); + } + if (!ok) ++bad_layers; + printf(" %5d %c %s worst cos %.6f at token %d (%d/%d tokens bad)\n", + layer, PF_PATTERN[layer], ok ? "PASS" : "FAIL", worst_cos, + worst_tok, bad_tokens, T); + } + { + auto dump = [&](const char *name, const void *dev, size_t bytes) { + std::vector host(bytes); + CHECK(cudaMemcpy(host.data(), dev, bytes, cudaMemcpyDeviceToHost)); + const std::string path = std::string("/out/pf_") + name + ".bin"; + FILE *f = fopen(path.c_str(), "wb"); + if (f) { fwrite(host.data(), 1, bytes, f); fclose(f); } + }; + const size_t rows = (size_t)min(T, BATCH_PAD); + dump("normed", bb_saved->pf_normed, rows * D_MODEL * sizeof(bf16)); + dump("zxbcdt", bb_saved->pf_zxbcdt, rows * MB_IN_PROJ * sizeof(bf16)); + dump("xbc_act", bb_saved->pf_xbc_act, rows * MB_CONV_DIM * sizeof(bf16)); + dump("gated", bb_saved->pf_gated, rows * MB_INNER * sizeof(bf16)); + printf(" dumped last-tile intermediates to /out/pf_*.bin (%zu rows)\n", + rows); + } + { + int bad_state = 0; + double worst_state_cos = 1.0; + int worst_state_layer = -1; + for (int layer = 0; layer < n_backbone; ++layer) { + if (PF_PATTERN[layer] != 'M') continue; + const int mi = bb_saved->mamba_of_layer[layer]; + const size_t n = (size_t)MB_CONV_DIM * (MB_CONV_K - 1); + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), + mamba_saved[mi].conv_state + + (size_t)hb.prefill_slot * + mamba_saved[mi].conv_slot_stride, + n * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + char nm[64]; + snprintf(nm, sizeof(nm), "/conv_state_out_%d.bin", layer); + auto want_raw = read_file(prefill_dir + nm, n * sizeof(uint16_t)); + const uint16_t *want = + reinterpret_cast(want_raw.data()); + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < n; ++i) { + const float a = __bfloat162float(*(const bf16 *)&mine[i]); + const float b = __bfloat162float(*(const bf16 *)&want[i]); + dot += (double)a * b; na += (double)a * a; nb += (double)b * b; + } + const double cs = (na > 0 && nb > 0) ? dot / std::sqrt(na * nb) + : (na == nb ? 1.0 : 0.0); + if (cs < worst_state_cos) { worst_state_cos = cs; worst_state_layer = layer; } + if (cs < 0.999) ++bad_state; + } + printf(" %s conv state out: worst cos %.6f at layer %d (%d bad)\n", + bad_state ? "FAIL" : "PASS", worst_state_cos, worst_state_layer, + bad_state); + if (bad_state) ++bad_layers; + } + { + const size_t n = (size_t)14 * PREFILL_MAX_TOKENS * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), pf_audit, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(prefill_dir + "/moe_w.bin", + (size_t)14 * T * BB_EXPERTS * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_mo = -1, worst_tok = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int tok = 0; tok < T; ++tok) { + const float *a = mine.data() + + ((size_t)mo * PREFILL_MAX_TOKENS + tok) * BB_EXPERTS; + const float *b = want + ((size_t)mo * T + tok) * BB_EXPERTS; + ++total_rows; + bool same = true; + for (int e = 0; e < BB_EXPERTS; ++e) + if ((a[e] != 0.f) != (b[e] != 0.f)) { same = false; break; } + if (same) continue; + ++rows_diff; + double mine_only = 0.0, theirs_only = 0.0; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (a[e] != 0.f && b[e] == 0.f) mine_only = std::max(mine_only, (double)a[e]); + if (b[e] != 0.f && a[e] == 0.f) theirs_only = std::max(theirs_only, (double)b[e]); + } + const double gap = std::abs(mine_only - theirs_only); + if (gap < worst_gap) { worst_gap = gap; worst_mo = mo; worst_tok = tok; } + if (rows_diff <= 3) { + printf(" flip mo=%d tok=%d\n kernel:", mo, tok); + for (int e = 0; e < BB_EXPERTS; ++e) + if (a[e] != 0.f) printf(" e%d=%.4f", e, a[e]); + printf("\n oracle:"); + for (int e = 0; e < BB_EXPERTS; ++e) + if (b[e] != 0.f) printf(" e%d=%.4f", e, b[e]); + printf("\n"); + } + } + } + printf(" router: %d of %d (layer, token) rows pick a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" widest disagreement at MoE layer %d token %d: the two " + "contested experts differ by %.6g\n", + worst_mo, worst_tok, worst_gap); + } + printf("%s prefill: %d/%d layers match the oracle\n", + bad_layers ? "FAIL" : "PASS", n_backbone - bad_layers, n_backbone); + return bad_layers ? 1 : 0; + } + if (decode_steps > 1 && !bb_saved) { + printf("FAIL --decode-steps needs the backbone (--stage0 --backbone-layers)\n"); + return 1; + } + for (int step = 0; step < decode_steps; ++step) { + if (step > 0) { + Backbone nb = *bb_saved; + nb.past_len = past_len + step; + CHECK(cudaMemcpy(d_bb, &nb, sizeof(Backbone), cudaMemcpyHostToDevice)); + } + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaDeviceSynchronize()); + } + if (launch_counter_host) { + unsigned long long lc[6] = {0, 0, 0, 0, 0, 0}; + CHECK(cudaMemcpy(lc, launch_counter_host, sizeof(lc), + cudaMemcpyDeviceToHost)); + const bool ok = lc[0] == (unsigned long long)decode_steps && + lc[1] == (unsigned long long)decode_steps; + printf("%s launch counter: entered=%llu completed=%llu, expected %d\n", + ok ? "PASS" : "FAIL", lc[0], lc[1], decode_steps); + } + bool pass = true; + if (n_backbone > 0 && dump_moe_w) { + std::vector mw((size_t)BATCH * BB_EXPERTS); + CHECK(cudaMemcpy(mw.data(), bb_moe_w_host, mw.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_moe_w, "wb"); + if (f) { fwrite(mw.data(), sizeof(float), mw.size(), f); fclose(f); } + printf(" moe_w -> %s\n", dump_moe_w); + } + if (n_backbone > 0 && verify) { + std::vector got((size_t)BB_LAYERS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), bb_layer_out_host, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + static const char PAT[] = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"; + const std::string bb_root = stage0 ? std::string(root) + : std::string(root) + "/../backbone"; + printf(" layer kind cos max_abs/rms\n"); + for (int layer = 0; layer < n_backbone; ++layer) { + char name[64]; + snprintf(name, sizeof(name), "/layer%02d_%c.bin", layer, PAT[layer]); + auto want_raw = read_file(bb_root + name, + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0, max_abs = 0, sq = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float( + got[(size_t)layer * BATCH * D_MODEL + i]); + const double e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const bool ok = cos > 0.999; + printf(" %5d %c %s %.6f %.4f\n", layer, PAT[layer], + ok ? "PASS" : "FAIL", cos, max_abs / fmax(rms, 1e-9)); + if (!ok) { + printf(" first divergence at layer %d (%c); stopping\n", + layer, PAT[layer]); + pass = false; + break; + } + } + if (moe_w_audit_host) { + const size_t n = (size_t)14 * BATCH * BB_EXPERTS; + std::vector mine(n); + CHECK(cudaMemcpy(mine.data(), moe_w_audit_host, n * sizeof(float), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/moe_w.bin", n * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + int rows_diff = 0, total_rows = 0; + double worst_gap = 1e30; + int worst_layer = -1; + for (int mo = 0; mo < 14; ++mo) { + for (int row = 0; row < BATCH; ++row) { + const size_t off = ((size_t)mo * BATCH + row) * BB_EXPERTS; + ++total_rows; + bool differ = false; + for (int e = 0; e < BB_EXPERTS; ++e) { + const bool a = mine[off + e] != 0.f; + const bool b = want[off + e] != 0.f; + if (a != b) { differ = true; break; } + } + if (!differ) continue; + ++rows_diff; + float mine_only = 0.f, want_only = 0.f; + for (int e = 0; e < BB_EXPERTS; ++e) { + if (mine[off + e] != 0.f && want[off + e] == 0.f) + mine_only = fmaxf(mine_only, mine[off + e]); + if (want[off + e] != 0.f && mine[off + e] == 0.f) + want_only = fmaxf(want_only, want[off + e]); + } + const double rel = fabs(mine_only - want_only) / + fmax(1e-9, (double)fmax(mine_only, want_only)); + printf(" MoE layer %2d row %2d: kernel's disputed pick weighs" + " %.4f, oracle's weighs %.4f (%.2f%% apart)\n", + mo, row, mine_only, want_only, 100.0 * rel); + if (rel < worst_gap) { worst_gap = rel; worst_layer = mo; } + } + } + printf(" router audit: %d of %d rows select a different expert set" + " (%.2f%%)\n", rows_diff, total_rows, + 100.0 * rows_diff / total_rows); + if (rows_diff) + printf(" uniform share for reference: %.4f" + " (routed mass %.1f over top-%d)\n", + BB_ROUTED_SCALE / 4.0, BB_ROUTED_SCALE, 4); + } + if (n_backbone == BB_LAYERS && ph_out_verify) { + std::vector mine(BATCH); + CHECK(cudaMemcpy(mine.data(), ph_out_verify, BATCH * sizeof(int), + cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/phoneme_argmax.bin", + (size_t)BATCH * sizeof(int)); + const int *want = reinterpret_cast(want_raw.data()); + std::vector margin((size_t)BATCH, 1e30f); + bool have_margin = false; + { + const std::string mp = bb_root + "/phoneme_margin.bin"; + if (FILE *probe = fopen(mp.c_str(), "rb")) { + fclose(probe); + auto raw = read_file(mp, (size_t)BATCH * sizeof(float)); + memcpy(margin.data(), raw.data(), raw.size()); + have_margin = true; + } + } + int bad = 0, tie = 0; + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + if (margin[i] <= 1.0f) ++tie; else ++bad; + } + printf("%s phoneme argmax: %d/%d exact", bad == 0 ? "PASS" : "FAIL", + BATCH - bad - tie, BATCH); + if (tie) + printf(", %d near-tie%s not graded", tie, tie == 1 ? "" : "s"); + if (!have_margin && (bad || tie)) + printf(" (no phoneme_margin.bin; every row treated as decidable)"); + printf("\n"); + for (int i = 0; i < BATCH; ++i) { + if (mine[i] == want[i]) continue; + printf(" row %d: oracle %d, kernel %d, oracle top-2 margin " + "%.3fx the bf16 accumulation scale -> %s\n", + i, want[i], mine[i], margin[i], + margin[i] <= 1.0f ? "NEAR-TIE, not a defect" : "DEFECT"); + } + pass = pass && bad == 0; + } + if (n_backbone == BB_LAYERS) { + std::vector dh((size_t)BATCH * D_MODEL); + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + auto want_raw = read_file(bb_root + "/dec_hidden.bin", + (size_t)BATCH * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double dot = 0, aa = 0, ee = 0; + for (int i = 0; i < BATCH * D_MODEL; ++i) { + const double a = __bfloat162float(dh[i]), e = want[i]; + dot += a * e; aa += a * a; ee += e * e; + } + const double cos = dot / fmax(sqrt(aa * ee), 1e-30); + printf(" %s dec_hidden (backbone -> predictor handoff): cos=%.6f\n", + cos > 0.9995 ? "PASS" : "FAIL", cos); + pass = pass && cos > 0.9995; + } + } + { + std::vector mine((size_t)BATCH * N_POS); + CHECK(cudaMemcpy(mine.data(), sampled, mine.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + auto refc = read_file(root + "/sampled_codes.bin", + (size_t)BATCH * N_POS * sizeof(int)); + const int *want_c = reinterpret_cast(refc.data()); + int bad = 0, first_bad_pos = -1; + std::vector per_pos(n_pos, 0); + for (int row = 0; row < BATCH; ++row) + for (int k = 0; k < n_pos; ++k) + if (mine[row * N_POS + k] != want_c[row * N_POS + k]) { + ++bad; + ++per_pos[k]; + if (first_bad_pos < 0) first_bad_pos = k; + } + printf("%s sampled codes: %d/%d exact", + bad == 0 ? "SAMPLE-PASS" : "SAMPLE-FAIL", + BATCH * n_pos - bad, BATCH * n_pos); + if (bad) { + printf(" first divergence at position %d, per-position:", + first_bad_pos); + for (int k = 0; k < n_pos; ++k) printf(" %d", per_pos[k]); + } + printf("\n"); + pass = pass && bad == 0; + } + if (dump_final) { + std::vector dh((size_t)BATCH * D_MODEL); + std::vector cd((size_t)BATCH * N_POS); + if (bb_dec_snapshot_host) + CHECK(cudaMemcpy(dh.data(), bb_dec_snapshot_host, + dh.size() * sizeof(bf16), cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(cd.data(), sampled, cd.size() * sizeof(int), + cudaMemcpyDeviceToHost)); + FILE *f = fopen(dump_final, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_final); return 1; } + fwrite(dh.data(), sizeof(bf16), dh.size(), f); + fwrite(cd.data(), sizeof(int), cd.size(), f); + fclose(f); + printf(" final state -> %s (%zu B dec_hidden + %zu B codes)\n", dump_final, + dh.size() * sizeof(bf16), cd.size() * sizeof(int)); + } + std::vector got((size_t)N_POS * BATCH * D_MODEL); + CHECK(cudaMemcpy(got.data(), g.hidden_rows, got.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + if (dump_rows) { + FILE *f = fopen(dump_rows, "wb"); + if (!f) { printf("FAIL cannot write %s\n", dump_rows); return 1; } + fwrite(got.data(), sizeof(bf16), (size_t)n_pos * BATCH * D_MODEL, f); + fclose(f); + printf(" rows -> %s\n", dump_rows); + } + if (verify && !stage0) { + auto want_raw = read_file(root + "/hidden_rows.bin", + (size_t)BATCH * n_pos * D_MODEL * sizeof(float)); + const float *want = reinterpret_cast(want_raw.data()); + double worst_ratio = 0.0, all_dot = 0.0, all_aa = 0.0, all_ee = 0.0; + int worst_pos = -1; + for (int pos = 0; pos < n_pos; ++pos) { + double max_abs = 0.0, sq = 0.0, dot = 0.0, aa = 0.0, ee = 0.0; + for (int row = 0; row < BATCH; ++row) { + for (int i = 0; i < D_MODEL; ++i) { + const double a = + __bfloat162float(got[((size_t)pos * BATCH + row) * D_MODEL + i]); + const double e = want[((size_t)row * n_pos + pos) * D_MODEL + i]; + max_abs = fmax(max_abs, fabs(a - e)); + sq += e * e; dot += a * e; aa += a * a; ee += e * e; + } + } + const double rms = sqrt(sq / (BATCH * D_MODEL)); + const double ratio = max_abs / fmax(rms, 1e-9); + all_dot += dot; all_aa += aa; all_ee += ee; + if (ratio > worst_ratio) { worst_ratio = ratio; worst_pos = pos; } + printf(" pos %2d max_abs=%.4f rms=%.4f max_abs/rms=%.4f cos=%.6f\n", + pos, max_abs, rms, ratio, dot / fmax(sqrt(aa * ee), 1e-30)); + } + const double cosine_all = all_dot / fmax(sqrt(all_aa * all_ee), 1e-30); + pass = cosine_all > 0.9995; + printf("%s worst max_abs/rms=%.4f (pos %d) cos_all=%.6f\n", + pass ? "PASS" : "FAIL", worst_ratio, worst_pos, cosine_all); + { + std::vector refc = read_file(root + "/ref_codes.bin", + (size_t)BATCH * n_pos * sizeof(int)); + const int *want_codes = reinterpret_cast(refc.data()); + int mismatches = 0, checked_codes = 0; + double worst_margin = 1e30; + int self_mismatch = 0; + double max_flip_gap = 0.0; + double sum_flip_gap = 0.0; + for (int k = 0; k < n_pos; ++k) { + auto pw = read_file(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL * sizeof(bf16)); + auto pb = read_file(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * sizeof(bf16)); + const bf16 *w = reinterpret_cast(pw.data()); + const bf16 *b = reinterpret_cast(pb.data()); + for (int row = 0; row < BATCH; ++row) { + const int ref_code = want_codes[row * n_pos + k]; + const bf16 *h = &got[((size_t)k * BATCH + row) * D_MODEL]; + int best = -1; + float best_v = -1e30f, second = -1e30f, at_ref = 0.f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(h[i]) * __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (n == ref_code) at_ref = acc; + if (acc > best_v) { second = best_v; best_v = acc; best = n; } + else if (acc > second) { second = acc; } + } + worst_margin = fmin(worst_margin, (double)(best_v - second)); + ++checked_codes; + int sbest = -1; + float sbest_v = -1e30f; + for (int n = 0; n < CODEBOOK; ++n) { + float acc = 0.f; + const bf16 *wr = w + (size_t)n * D_MODEL; + for (int i = 0; i < D_MODEL; ++i) + acc += __bfloat162float(__float2bfloat16( + want[((size_t)row * n_pos + k) * D_MODEL + i])) * + __bfloat162float(wr[i]); + acc += __bfloat162float(b[n]); + if (acc > sbest_v) { sbest_v = acc; sbest = n; } + } + if (sbest != ref_code) ++self_mismatch; + if (best != sbest) { + ++mismatches; + const double gap = best_v - at_ref; + max_flip_gap = fmax(max_flip_gap, fabs(gap)); + sum_flip_gap += fabs(gap); + } + } + } + const double agree = 100.0 * (checked_codes - mismatches) / checked_codes; + printf(" host-vs-torch projection noise: %d/%d (%.2f%%) - not the gate\n", + self_mismatch, checked_codes, 100.0 * self_mismatch / checked_codes); + printf("%s sampled codes vs reference rows: %d/%d agree (%.2f%%), %d flip, " + "min argmax margin %.5f\n", + mismatches == 0 ? "GATE-PASS" : "GATE-FAIL", + checked_codes - mismatches, checked_codes, agree, mismatches, + worst_margin); + if (mismatches) + printf(" flipped decisions: mean logit gap %.5f, max %.5f" + " (small gaps = genuine near-ties)\n", + sum_flip_gap / mismatches, max_flip_gap); + pass = pass && mismatches == 0; + } + } + if (bench) { + cudaEvent_t t0, t1; + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + for (int i = 0; i < 3; ++i) { + CHECK(cudaMemcpy(x, x_initial, (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + } + CHECK(cudaDeviceSynchronize()); + const int iters = bench_iters; + std::vector marks(iters + 1); + for (auto &ev : marks) CHECK(cudaEventCreate(&ev)); + CHECK(cudaEventRecord(marks[0])); + for (int i = 0; i < iters; ++i) { + predictor_mk2<<>>(g, d_bb, n_layers, n_pos, n_backbone); + CHECK(cudaEventRecord(marks[i + 1])); + } + CHECK(cudaEventSynchronize(marks[iters])); + std::vector samples; + samples.reserve(iters); + for (int i = 0; i < iters; ++i) { + float one = 0.f; + CHECK(cudaEventElapsedTime(&one, marks[i], marks[i + 1])); + samples.push_back(one); + } + float span = 0.f; + CHECK(cudaEventElapsedTime(&span, marks[0], marks[iters])); + for (auto &ev : marks) CHECK(cudaEventDestroy(ev)); + double total = 0.0; + for (double s : samples) total += s; + std::sort(samples.begin(), samples.end()); + const double per_step = samples[iters / 2]; + printf(" bench: %d iters median %.4f ms mean %.4f ms " + "aggregate %.4f ms min %.4f p05 %.4f p95 %.4f max %.4f " + "IQR %.4f\n", + iters, per_step, total / iters, span / iters, samples.front(), + samples[(size_t)(0.05 * iters)], samples[(size_t)(0.95 * iters)], + samples.back(), + samples[(size_t)(0.75 * iters)] - samples[(size_t)(0.25 * iters)]); + const double weight_bytes = (double)n_pos * n_layers * + ((3.0 * D_MODEL * D_MODEL) + (double)D_MODEL * D_MODEL + + (double)D_FFN * D_MODEL + (double)D_MODEL * D_FFN) * 2.0; + printf(" per decode step: %.3f ms weights %.2f GB %.2f TB/s" + " floor at 3.35 TB/s: %.3f ms\n", + per_step, weight_bytes / 1e9, weight_bytes / (per_step * 1e-3) / 1e12, + weight_bytes / 3.35e12 * 1e3); + printf(" for comparison: mma.sync megakernel 9.1 ms, TensorRT 3.414 ms\n"); + } + return pass ? 0 : 1; +} +#endif +#include "mk2_abi.cuh" +namespace { +thread_local std::string g_last_error; +struct MK2State { + Globals g; + int batch; + int grid; + int smem; + bf16 *hidden_bf16; + int kind; + int device; + std::vector owned; + int *run_predicate = nullptr; + const unsigned char *native_predictor_active_mask = nullptr; + MK2_PROGRESS_MEMBER + int maximum_grid = 0; + bool grid_reset_pending = false; +}; +static void mk2_check(cudaError_t e) { + if (e != cudaSuccess) throw std::runtime_error(cudaGetErrorString(e)); +} +struct MK2Allocations { + std::vector ptrs; + ~MK2Allocations() { for (void *p : ptrs) cudaFree(p); } + template cudaError_t alloc(T **out, size_t n) { + cudaError_t e = cudaMalloc(out, n); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFree(*out); throw; } + } + return e; + } + template T *upload_raw(const std::string &path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot open " + path); + const auto end = f.tellg(); + if (end <= 0 || size_t(end) % sizeof(T) || + (count && size_t(end) != count * sizeof(T))) + throw std::runtime_error("wrong byte count: " + path); + std::vector data((size_t(end))); + f.seekg(0); + if (!f.read(data.data(), data.size())) + throw std::runtime_error("short read: " + path); + T *d = nullptr; + mk2_check(alloc(&d, data.size())); + mk2_check(cudaMemcpy(d, data.data(), data.size(), cudaMemcpyHostToDevice)); + return d; + } +}; +} +extern "C" int mk2_build_info(MK2BuildInfo *out) { + if (!out) { g_last_error = "null build info"; return 1; } + *out = {1, BATCH, BATCH_PAD, STAGES, KCHUNK, MK2_DIRECT_INPUT, + MK2_B1_LOCAL_BARRIERS, MK2_SKIP_HIDDEN_SNAPSHOT, THREADS, + int(STAGES * (sizeof(weight_tile) + sizeof(act_tile)) + + sizeof(out_tile) + sizeof(out_tile_f) + 4096), 0, 0, 0, 0, 0, 0}; + try { + mk2_check(cudaGetDevice(&out->device)); + cudaFuncAttributes attr{}; + mk2_check(cudaFuncGetAttributes(&attr, predictor_mk2)); + out->static_smem = int(attr.sharedSizeBytes); + out->registers = attr.numRegs; + int optin = 0; + mk2_check(cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, + out->device)); + mk2_check(cudaDeviceGetAttribute(&out->sms, cudaDevAttrMultiProcessorCount, + out->device)); + if (out->dynamic_smem + out->static_smem > optin) + throw std::runtime_error("shared memory exceeds device limit: dynamic=" + + std::to_string(out->dynamic_smem) + " static=" + + std::to_string(out->static_smem) + " limit=" + std::to_string(optin)); + mk2_check(cudaFuncSetAttribute(predictor_mk2, + cudaFuncAttributeMaxDynamicSharedMemorySize, out->dynamic_smem)); + mk2_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &out->blocks_per_sm, predictor_mk2, THREADS, out->dynamic_smem)); + if (out->blocks_per_sm < 1) + throw std::runtime_error("no resident block; grid barrier cannot launch"); + out->grid = std::min(out->sms * out->blocks_per_sm, 4096); + if (const char *e = getenv("NEMOTRON_STAGE0_MEGAKERNEL_GRID")) { + const int want = atoi(e); + if (want > 0) out->grid = std::min(out->grid, want); + } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } +} +extern "C" MK2Predictor *mk2_create_from_weights(const MK2Weights *w, int batch) { + if (batch != BATCH) { + g_last_error = "this build supports batch " + std::to_string(BATCH) + + " only; got " + std::to_string(batch); + return nullptr; + } + if (!w) { g_last_error = "null weights"; return nullptr; } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + for (int i = 0; i < N_LAYERS; ++i) + if (!w->qkv[i] || !w->o[i] || !w->ff_in[i] || !w->ff_out[i] || + !w->norm_self[i] || !w->norm_ff[i]) { + g_last_error = "null layer weight"; return nullptr; + } + if (!w->pos_emb) { g_last_error = "null position embedding"; return nullptr; } + for (int i = 0; i < N_POS; ++i) + if (!w->head_w[i] || !w->head_b[i] || !w->audio_emb[i]) { + g_last_error = "null head/audio embedding"; return nullptr; + } + try { + MK2Allocations owned; + bf16 *head_w[N_POS]; + const bf16 *head_b[N_POS]; + for (int k = 0; k < N_POS; ++k) { + mk2_check(owned.alloc(&head_w[k], (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(head_w[k], 0, (size_t)VOCAB_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(head_w[k], w->head_w[k], + (size_t)CODEBOOK * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + bf16 *bd = nullptr; + mk2_check(owned.alloc(&bd, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(bd, 0, (size_t)VOCAB_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(bd, w->head_b[k], (size_t)CODEBOOK * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + head_b[k] = bd; + } + std::vector forb(VOCAB_PADDED, 0); + for (int n = CODEBOOK - 8; n < VOCAB_PADDED; ++n) forb[n] = 1; + forb[AUDIO_EOS] = 0; + int8_t *forbidden = nullptr; + mk2_check(owned.alloc(&forbidden, forb.size())); + mk2_check(cudaMemcpy(forbidden, forb.data(), forb.size(), cudaMemcpyHostToDevice)); + bf16 *x, *qkv_out, *attn_out, *ff_hidden, *normed, *kcache, *vcache, *hidden_rows; + bf16 *logits_buf, *hidden_staging; + int *sampled, *sampled_pos, *codes_unused; + unsigned int *barrier; + int *run_predicate; + mk2_check(owned.alloc(&run_predicate, sizeof(int))); + mk2_check(owned.alloc(&x, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&qkv_out, (size_t)BATCH * 3 * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&attn_out, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&ff_hidden, (size_t)BATCH * D_FFN * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + const size_t cache_elems = (size_t)N_LAYERS * BATCH * N_HEADS * N_POS * D_HEAD; + mk2_check(owned.alloc(&kcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&vcache, cache_elems * sizeof(bf16))); + mk2_check(owned.alloc(&hidden_rows, (size_t)N_POS * BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&logits_buf, (size_t)BATCH * VOCAB_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&sampled, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(owned.alloc(&sampled_pos, (size_t)BATCH * sizeof(int))); + mk2_check(owned.alloc(&barrier, 64 * sizeof(unsigned int))); + mk2_check(cudaMemset(barrier, 0, 64 * sizeof(unsigned int))); + hidden_staging = nullptr; + if constexpr (!MK2_DIRECT_INPUT) + mk2_check(owned.alloc(&hidden_staging, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&codes_unused, (size_t)BATCH * N_POS * sizeof(int))); + mk2_check(cudaMemset(codes_unused, 0, (size_t)BATCH * N_POS * sizeof(int))); + const size_t DM = D_MODEL, DF = D_FFN, B = BATCH, QKV = 3 * D_MODEL; + const size_t VP = VOCAB_PADDED; + auto W = [&](int i, const void *const *arr, size_t rows, size_t cols) { + return weight_gl{(bf16 *)arr[i], nullptr, nullptr, rows, cols}; + }; + Globals g{ + {W(0, w->qkv, QKV, DM), W(1, w->qkv, QKV, DM), W(2, w->qkv, QKV, DM)}, + {W(0, w->o, DM, DM), W(1, w->o, DM, DM), W(2, w->o, DM, DM)}, + {W(0, w->ff_in, DF, DM), W(1, w->ff_in, DF, DM), W(2, w->ff_in, DF, DM)}, + {W(0, w->ff_out, DM, DF), W(1, w->ff_out, DM, DF), W(2, w->ff_out, DM, DF)}, + {weight_gl{head_w[0], nullptr, nullptr, VP, DM}, + weight_gl{head_w[1], nullptr, nullptr, VP, DM}, + weight_gl{head_w[2], nullptr, nullptr, VP, DM}, + weight_gl{head_w[3], nullptr, nullptr, VP, DM}, + weight_gl{head_w[4], nullptr, nullptr, VP, DM}, + weight_gl{head_w[5], nullptr, nullptr, VP, DM}, + weight_gl{head_w[6], nullptr, nullptr, VP, DM}, + weight_gl{head_w[7], nullptr, nullptr, VP, DM}, + weight_gl{head_w[8], nullptr, nullptr, VP, DM}, + weight_gl{head_w[9], nullptr, nullptr, VP, DM}, + weight_gl{head_w[10], nullptr, nullptr, VP, DM}, + weight_gl{head_w[11], nullptr, nullptr, VP, DM}, + weight_gl{head_w[12], nullptr, nullptr, VP, DM}, + weight_gl{head_w[13], nullptr, nullptr, VP, DM}, + weight_gl{head_w[14], nullptr, nullptr, VP, DM}, + weight_gl{head_w[15], nullptr, nullptr, VP, DM}}, + act_gl{x, nullptr, nullptr, B, DM}, + act_gl{normed, nullptr, nullptr, B, DM}, + act_gl{attn_out, nullptr, nullptr, B, DM}, + act_gl{ff_hidden, nullptr, nullptr, B, DF}, + x, qkv_out, attn_out, ff_hidden, normed, + logits_buf, + {head_b[0], head_b[1], head_b[2], head_b[3], head_b[4], head_b[5], + head_b[6], head_b[7], head_b[8], head_b[9], head_b[10], head_b[11], + head_b[12], head_b[13], head_b[14], head_b[15]}, + forbidden, nullptr, sampled, sampled_pos, 0.7f, + {(const float *)w->norm_self[0], (const float *)w->norm_self[1], + (const float *)w->norm_self[2]}, + {(const float *)w->norm_ff[0], (const float *)w->norm_ff[1], + (const float *)w->norm_ff[2]}, + (const bf16 *)w->pos_emb, + {(const bf16 *)w->audio_emb[0], (const bf16 *)w->audio_emb[1], + (const bf16 *)w->audio_emb[2], (const bf16 *)w->audio_emb[3], + (const bf16 *)w->audio_emb[4], (const bf16 *)w->audio_emb[5], + (const bf16 *)w->audio_emb[6], (const bf16 *)w->audio_emb[7], + (const bf16 *)w->audio_emb[8], (const bf16 *)w->audio_emb[9], + (const bf16 *)w->audio_emb[10], (const bf16 *)w->audio_emb[11], + (const bf16 *)w->audio_emb[12], (const bf16 *)w->audio_emb[13], + (const bf16 *)w->audio_emb[14], (const bf16 *)w->audio_emb[15]}, + kcache, vcache, hidden_rows, codes_unused, barrier, + }; + auto *st = new MK2State{g, BATCH, info.grid, info.dynamic_smem, + hidden_staging, 0, info.device, {}, run_predicate}; + st->maximum_grid = info.sms * info.blocks_per_sm; + st->owned.swap(owned.ptrs); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } +} +extern "C" MK2Predictor *mk2_create(const char *raw_dir, int batch) { + if (!raw_dir || batch != BATCH) { + g_last_error = "null raw_dir or wrong batch"; return nullptr; + } + MK2BuildInfo info{}; + if (mk2_build_info(&info)) return nullptr; + MK2Predictor *p = nullptr; + try { + MK2Allocations owned; + const std::string root = raw_dir; + MK2Weights w{}; + for (int l = 0; l < N_LAYERS; ++l) { + const std::string p = root + "/l" + std::to_string(l) + "_"; + w.qkv[l] = owned.upload_raw(p + "qkv.bin", (size_t)3 * D_MODEL * D_MODEL); + w.o[l] = owned.upload_raw(p + "o.bin", (size_t)D_MODEL * D_MODEL); + w.ff_in[l] = owned.upload_raw(p + "ff_in.bin", (size_t)D_FFN * D_MODEL); + w.ff_out[l] = owned.upload_raw(p + "ff_out.bin", (size_t)D_MODEL * D_FFN); + w.norm_self[l] = owned.upload_raw(p + "norm_self.bin", D_MODEL); + w.norm_ff[l] = owned.upload_raw(p + "norm_ff.bin", D_MODEL); + } + w.pos_emb = owned.upload_raw(root + "/position_embeddings.bin", (size_t)18 * D_MODEL); + for (int k = 0; k < N_POS; ++k) { + w.audio_emb[k] = + owned.upload_raw(root + "/audio_emb_" + std::to_string(k) + ".bin", 0); + w.head_w[k] = owned.upload_raw(root + "/out_proj_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK * D_MODEL); + w.head_b[k] = owned.upload_raw(root + "/out_bias_" + std::to_string(k) + ".bin", + (size_t)CODEBOOK); + } + p = mk2_create_from_weights(&w, batch); + if (!p) return nullptr; + auto *st = reinterpret_cast(p); + st->owned.insert(st->owned.end(), owned.ptrs.begin(), owned.ptrs.end()); + owned.ptrs.clear(); + return p; + } catch (const std::exception &e) { + if (p) mk2_destroy(p); + g_last_error = e.what(); return nullptr; + } +} +extern "C" int mk2_run(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + void *codes_int64, void *stream) { + auto *st = reinterpret_cast(p); + if (!st) { g_last_error = "null handle"; return 1; } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>(st->g, nullptr, N_LAYERS, N_POS, 0); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +extern "C" int mk2_run_masked(MK2Predictor *p, const void *hidden_fp16, + const void *gumbel_fp32, float temperature, + const void *audio_valid_bool, void *codes_int64, + void *stream) { + auto *st = reinterpret_cast(p); + if (!st || st->kind != 0 || !hidden_fp16 || !gumbel_fp32 || + !audio_valid_bool || !codes_int64 || !(temperature > 0.f) || + !std::isfinite(temperature)) { + g_last_error = "masked predictor requires predictor handle, non-null device buffers and finite positive temperature"; + return 1; + } + cudaPointerAttributes attributes{}; + const cudaError_t pointer_error = cudaPointerGetAttributes(&attributes, audio_valid_bool); + if (pointer_error != cudaSuccess) { + (void)cudaGetLastError(); + g_last_error = "audio_valid must be a device bool buffer"; + return 1; + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != st->device) { + g_last_error = "audio_valid must be a device bool buffer on the predictor device"; + return 1; + } + cudaStream_t s = static_cast(stream); + const int n_hidden = BATCH * D_MODEL; + fp16_to_bf16_masked<<<(n_hidden + 255) / 256, 256, 0, s>>>( + static_cast(hidden_fp16), + MK2_DIRECT_INPUT ? st->g.x : st->hidden_bf16, n_hidden, + static_cast(audio_valid_bool), st->run_predicate); + if constexpr (!MK2_DIRECT_INPUT) { + const auto copy_error = cudaMemcpyAsync(st->g.x, st->hidden_bf16, + (size_t)n_hidden * sizeof(bf16), cudaMemcpyDeviceToDevice, s); + if (copy_error != cudaSuccess) { + g_last_error = cudaGetErrorString(copy_error); return 1; + } + } + st->g.gumbel = static_cast(gumbel_fp32); + st->g.temperature = temperature; + predictor_mk2<<grid, THREADS, st->smem, s>>>( + st->g, nullptr, N_LAYERS, N_POS, 0, st->run_predicate); + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->g.sampled, static_cast(codes_int64), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + return 0; +} +namespace { +struct MK2Stage0State { + MK2State pred; + Backbone *d_bb; + Backbone host_bb; + std::vector mb_host; + std::vector at_host; + Backbone *bb_stage; + MambaLayer *mb_stage; + AttnLayer *at_stage; + std::vector owned; + std::vector pinned; + MK2Stage0Bindings pinned_bindings{}; + bool bindings_initialized = false; +}; +struct MK2Stage0PinnedAllocations { + std::vector ptrs; + ~MK2Stage0PinnedAllocations() { + for (void *p : ptrs) cudaFreeHost(p); + } + template cudaError_t alloc(T **out, size_t n) { + const cudaError_t e = cudaHostAlloc(out, n, cudaHostAllocDefault); + if (e == cudaSuccess) { + try { ptrs.push_back(*out); } + catch (...) { cudaFreeHost(*out); *out = nullptr; throw; } + } + return e; + } +}; +struct MK2Stage0BaseDeleter { + void operator()(MK2State *st) const noexcept { + if (!st) return; + for (void *p : st->owned) cudaFree(p); + delete st; + } +}; +static void mk2_stage0_device_pointer(const void *p, int device, + const char *name) { + if (!p) throw std::runtime_error(std::string("null Stage 0 binding: ") + name); + cudaPointerAttributes attributes{}; + const cudaError_t e = cudaPointerGetAttributes(&attributes, p); + if (e != cudaSuccess) { + (void)cudaGetLastError(); + throw std::runtime_error(std::string("Stage 0 binding is not a device pointer: ") + name); + } + if (attributes.type != cudaMemoryTypeDevice || attributes.device != device) + throw std::runtime_error(std::string("Stage 0 binding must be on the handle device: ") + name); +} +static bool mk2_stage0_same_pinned_bindings(const MK2Stage0Bindings &a, + const MK2Stage0Bindings &b) { + for (int i = 0; i < 14; ++i) + if (a.conv_state[i] != b.conv_state[i] || + a.conv_slot_stride[i] != b.conv_slot_stride[i] || + a.ssm_state[i] != b.ssm_state[i] || + a.ssm_slot_stride[i] != b.ssm_slot_stride[i]) return false; + for (int i = 0; i < 3; ++i) + if (a.kv_pool_base[i] != b.kv_pool_base[i] || + a.kv_block_table[i] != b.kv_block_table[i]) return false; + return a.slot_map == b.slot_map && a.seq_len == b.seq_len && + a.kv_max_blocks_per_seq == b.kv_max_blocks_per_seq && + a.kv_tokens_per_block == b.kv_tokens_per_block && + a.phoneme_out == b.phoneme_out && + a.moe_w_inject_debug == b.moe_w_inject_debug; +} +template +T *upload_layers(const std::vector &v, MK2Allocations &owned) { + T *d = nullptr; + mk2_check(owned.alloc(&d, v.size() * sizeof(T))); + mk2_check(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} +} +extern "C" MK2Predictor *mk2_create_stage0(const MK2Weights *pw, + const MK2BackboneWeights *bw, + int batch) { + try { + if (!pw || !bw) { g_last_error = "null weights"; return nullptr; } + if (batch != BATCH) { + g_last_error = "full Stage 0 build supports batch " + std::to_string(BATCH); + return nullptr; + } + { + struct Named { const void *p; const char *name; }; + const Named required[] = { + {bw->norm_f, "norm_f"}, + {bw->phoneme_w, "phoneme_w"}, + {bw->phoneme_b, "phoneme_b"}, + }; + for (const auto &r : required) { + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + + "' is null"; + return nullptr; + } + } + for (int l = 0; l < BB_LAYERS; ++l) + if (!bw->norm[l]) { + g_last_error = "backbone weight 'norm[" + std::to_string(l) + + "]' is null"; + return nullptr; + } + for (int i = 0; i < 14; ++i) { + const Named per[] = { + {bw->mb_in_proj[i], "mb_in_proj"}, + {bw->mb_out_proj[i], "mb_out_proj"}, + {bw->mb_conv_w[i], "mb_conv_w"}, + {bw->mb_conv_b[i], "mb_conv_b"}, + {bw->mb_A[i], "mb_A"}, + {bw->mb_D[i], "mb_D"}, + {bw->mb_dt_bias[i], "mb_dt_bias"}, + {bw->mb_gnorm[i], "mb_gnorm"}, + {bw->moe_exp_up[i], "moe_exp_up"}, + {bw->moe_exp_down[i], "moe_exp_down"}, + {bw->moe_sh_up[i], "moe_sh_up"}, + {bw->moe_sh_down[i], "moe_sh_down"}, + {bw->moe_gate[i], "moe_gate"}, + }; + for (const auto &r : per) + if (!r.p) { + g_last_error = std::string("backbone weight '") + r.name + "[" + + std::to_string(i) + "]' is null"; + return nullptr; + } + } + for (int i = 0; i < 3; ++i) { + if (!bw->at_qkv[i] || !bw->at_o[i]) { + g_last_error = "attention weight " + std::to_string(i) + " is null"; + return nullptr; + } + } + } + if (!bw->pattern || strlen(bw->pattern) != BB_LAYERS) { + g_last_error = "pattern must be exactly " + std::to_string(BB_LAYERS) + + " characters of M/E/*"; + return nullptr; + } + int mamba_of[BB_LAYERS], attn_of[BB_LAYERS], moe_of[BB_LAYERS]; + int n_mb = 0, n_at = 0, n_moe = 0; + for (int l = 0; l < BB_LAYERS; ++l) { + mamba_of[l] = attn_of[l] = moe_of[l] = -1; + switch (bw->pattern[l]) { + case 'M': mamba_of[l] = n_mb++; break; + case 'E': moe_of[l] = n_moe++; break; + case '*': attn_of[l] = n_at++; break; + default: + g_last_error = std::string("bad pattern character '") + + bw->pattern[l] + "' at layer " + std::to_string(l); + return nullptr; + } + } + if (n_mb != 14 || n_at != 3 || n_moe != 14) { + g_last_error = "layer pattern mismatch: got " + std::to_string(n_mb) + + " Mamba, " + std::to_string(n_at) + " attention, " + + std::to_string(n_moe) + " MoE; expected 14/3/14"; + return nullptr; + } + std::unique_ptr pred( + reinterpret_cast(mk2_create_from_weights(pw, batch))); + if (!pred) return nullptr; + MK2Allocations owned; + MK2Stage0PinnedAllocations pinned; + std::vector mb; + for (int i = 0; i < 14; ++i) { + mb.push_back(MambaLayer{ + weight_gl{(bf16 *)bw->mb_in_proj[i], nullptr, nullptr, + (size_t)MB_IN_PROJ, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->mb_out_proj[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)MB_INNER}, + (const bf16 *)bw->mb_conv_w[i], (const bf16 *)bw->mb_conv_b[i], + (const float *)bw->mb_A[i], (const float *)bw->mb_D[i], + (const float *)bw->mb_dt_bias[i], (const bf16 *)bw->mb_gnorm[i], + nullptr, 0, nullptr, 0, + }); + } + std::vector at; + for (int i = 0; i < 3; ++i) { + at.push_back(AttnLayer{ + weight_gl{(bf16 *)bw->at_qkv[i], nullptr, nullptr, (size_t)BB_QKV, + (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->at_o[i], nullptr, nullptr, (size_t)D_MODEL, + (size_t)D_MODEL}, + nullptr, nullptr, 0, 0, + }); + } + std::vector mo; + bf16 *zero_gate_bias = nullptr; + for (int i = 0; i < 14; ++i) + if (!bw->moe_gate_bias[i] && !zero_gate_bias) { + mk2_check(owned.alloc(&zero_gate_bias, BB_EXPERTS * sizeof(bf16))); + mk2_check(cudaMemset(zero_gate_bias, 0, BB_EXPERTS * sizeof(bf16))); + } + for (int i = 0; i < 14; ++i) { + mo.push_back(MoeLayer{ + weight_gl{(bf16 *)bw->moe_exp_up[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * BB_MOE_INTER), (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_exp_down[i], nullptr, nullptr, + (size_t)(BB_EXPERTS * D_MODEL), (size_t)BB_MOE_INTER}, + weight_gl{(bf16 *)bw->moe_sh_up[i], nullptr, nullptr, + (size_t)BB_SHARED_INTER, (size_t)D_MODEL}, + weight_gl{(bf16 *)bw->moe_sh_down[i], nullptr, nullptr, + (size_t)D_MODEL, (size_t)BB_SHARED_INTER}, + (const bf16 *)bw->moe_gate[i], + bw->moe_gate_bias[i] ? (const bf16 *)bw->moe_gate_bias[i] : zero_gate_bias, + }); + } + bf16 *h, *r, *normed, *zxbcdt, *gated, *qkv, *moe_h, *moe_sh, *xbc, *ph_logits; + float *moe_out, *moe_w; + int *ph_out; + mk2_check(owned.alloc(&h, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&r, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&normed, (size_t)BATCH * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&zxbcdt, (size_t)BATCH * MB_IN_PROJ * sizeof(bf16))); + mk2_check(owned.alloc(&gated, (size_t)BATCH * MB_INNER * sizeof(bf16))); + mk2_check(owned.alloc(&qkv, (size_t)BATCH * BB_QKV * sizeof(bf16))); + mk2_check(owned.alloc(&moe_h, + (size_t)BATCH_PAD * BB_EXPERTS * BB_MOE_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_sh, (size_t)BATCH * BB_SHARED_INTER * sizeof(bf16))); + mk2_check(owned.alloc(&moe_out, (size_t)BATCH * (BB_EXPERTS + 1) * D_MODEL * sizeof(float))); + mk2_check(owned.alloc(&moe_w, (size_t)BATCH * BB_EXPERTS * sizeof(float))); + mk2_check(owned.alloc(&xbc, (size_t)BATCH * MB_CONV_DIM * sizeof(bf16))); + mk2_check(owned.alloc(&ph_logits, (size_t)BATCH * BB_PHON_PADDED * sizeof(bf16))); + mk2_check(owned.alloc(&ph_out, (size_t)BATCH * sizeof(int))); + unsigned long long *launch_counter = nullptr; + mk2_check(owned.alloc(&launch_counter, 6 * sizeof(unsigned long long))); + mk2_check(cudaMemset(launch_counter, 0, 6 * sizeof(unsigned long long))); + bf16 *ph_w = nullptr, *ph_b = nullptr; + mk2_check(owned.alloc(&ph_w, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemset(ph_w, 0, (size_t)BB_PHON_PADDED * D_MODEL * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_w, bw->phoneme_w, + (size_t)BB_PHONEME * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + mk2_check(owned.alloc(&ph_b, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemset(ph_b, 0, (size_t)BB_PHON_PADDED * sizeof(bf16))); + mk2_check(cudaMemcpy(ph_b, bw->phoneme_b, (size_t)BB_PHONEME * sizeof(bf16), + cudaMemcpyDeviceToDevice)); + Backbone hb{ + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + upload_layers(mb, owned), + (const bf16 *)bw->norm_f, + upload_layers(at, owned), + upload_layers(mo, owned), + {}, {}, {}, {}, + 0, + h, r, normed, zxbcdt, gated, + nullptr, + qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + nullptr, + xbc, + nullptr, + 0, 0, 1, 0, nullptr, nullptr, + h, r, normed, zxbcdt, xbc, gated, qkv, moe_h, moe_sh, moe_out, moe_w, + nullptr, + nullptr, + act_gl{normed, nullptr, nullptr, (size_t)BATCH, (size_t)D_MODEL}, + act_gl{gated, nullptr, nullptr, (size_t)BATCH, (size_t)MB_INNER}, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + weight_gl{ph_w, nullptr, nullptr, (size_t)BB_PHON_PADDED, (size_t)D_MODEL}, + ph_b, ph_logits, ph_out, + launch_counter, + nullptr, + act_gl{qkv, nullptr, nullptr, (size_t)BATCH, (size_t)BB_QKV}, + act_gl{moe_h, nullptr, nullptr, (size_t)BATCH, + (size_t)(BB_EXPERTS * BB_MOE_INTER)}, + act_gl{moe_sh, nullptr, nullptr, (size_t)BATCH, (size_t)BB_SHARED_INTER}, + }; + for (int l = 0; l < BB_LAYERS; ++l) { + hb.norm_w[l] = (const bf16 *)bw->norm[l]; + hb.mamba_of_layer[l] = mamba_of[l]; + hb.attn_of_layer[l] = attn_of[l]; + hb.moe_of_layer[l] = moe_of[l]; + } + if constexpr (MK2_SPARSE_MOE && BATCH >= 2) { + bf16 *gathered = nullptr; + int *rows = nullptr, *counts = nullptr; + SparseMoe *device_sparse = nullptr; + mk2_check(owned.alloc(&gathered, + (size_t)BB_EXPERTS * BATCH_PAD * D_MODEL * sizeof(bf16))); + mk2_check(owned.alloc(&rows, (size_t)BB_EXPERTS * BATCH_PAD * sizeof(int))); + mk2_check(owned.alloc(&counts, BB_EXPERTS * sizeof(int))); + SparseMoe sparse{gathered, rows, counts, + sparse_act_gl{gathered, nullptr, nullptr, (size_t)BB_EXPERTS * BATCH_PAD, + (size_t)D_MODEL}, + sparse_act_gl{moe_h, nullptr, nullptr, (size_t)BB_EXPERTS * BATCH_PAD, + (size_t)BB_MOE_INTER}}; + mk2_check(owned.alloc(&device_sparse, sizeof(SparseMoe))); + mk2_check(cudaMemcpy(device_sparse, &sparse, sizeof(sparse), cudaMemcpyHostToDevice)); + hb.sparse_moe = device_sparse; + } + Backbone *d_bb = nullptr, *bb_stage = nullptr; + MambaLayer *mb_stage = nullptr; + AttnLayer *at_stage = nullptr; + mk2_check(owned.alloc(&d_bb, sizeof(Backbone))); + mk2_check(pinned.alloc(&bb_stage, sizeof(Backbone))); + mk2_check(pinned.alloc(&mb_stage, 14 * sizeof(MambaLayer))); + mk2_check(pinned.alloc(&at_stage, 3 * sizeof(AttnLayer))); + auto *st = new MK2Stage0State{ + MK2State{pred->g, pred->batch, pred->grid, pred->smem, + pred->hidden_bf16, 1, pred->device, {}, pred->run_predicate}, + d_bb, hb, std::move(mb), std::move(at), bb_stage, mb_stage, at_stage, + {}, {}}; + st->pred.maximum_grid = pred->maximum_grid; + st->pred.owned.swap(pred->owned); + st->owned.swap(owned.ptrs); + st->pinned.swap(pinned.ptrs); + g_last_error.clear(); + return reinterpret_cast(st); + } catch (const std::exception &e) { g_last_error = e.what(); return nullptr; } + catch (...) { g_last_error = "unknown full Stage 0 construction failure"; return nullptr; } +} +static int mk2_enqueue_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream, bool native) { + try { + const auto *pred = reinterpret_cast(p); + if (!pred || pred->kind != 1 || pred->batch != BATCH || !b) { + g_last_error = "full Stage 0 requires a matching Stage 0 handle and bindings"; + return 1; + } + int device = -1; + mk2_check(cudaGetDevice(&device)); + if (device != pred->device) + throw std::runtime_error("current CUDA device differs from full Stage 0 handle device"); + if (!(b->temperature > 0.f) || !std::isfinite(b->temperature)) + throw std::runtime_error("full Stage 0 temperature must be finite and positive"); + if (b->kv_max_blocks_per_seq <= 0 || b->kv_tokens_per_block <= 0 || + b->kv_tokens_per_block > INT_MAX / (BB_KV_HEADS * BB_HEAD_DIM)) + throw std::runtime_error("invalid full Stage 0 KV geometry"); + for (int i = 0; i < 14; ++i) { + if (b->conv_slot_stride[i] < MB_CONV_DIM * (MB_CONV_K - 1) || + b->ssm_slot_stride[i] < MB_HEADS * MB_HEAD_DIM * MB_STATE) + throw std::runtime_error("full Stage 0 state slot stride is smaller than its state"); + mk2_stage0_device_pointer(b->conv_state[i], device, "conv_state"); + mk2_stage0_device_pointer(b->ssm_state[i], device, "ssm_state"); + } + for (int i = 0; i < 3; ++i) { + mk2_stage0_device_pointer(b->kv_pool_base[i], device, "kv_pool_base"); + mk2_stage0_device_pointer(b->kv_block_table[i], device, "kv_block_table"); + } + mk2_stage0_device_pointer(b->slot_map, device, "slot_map"); + mk2_stage0_device_pointer(b->seq_len, device, "seq_len"); + mk2_stage0_device_pointer(b->embedding, device, "embedding"); + mk2_stage0_device_pointer(b->gumbel, device, "gumbel"); + mk2_stage0_device_pointer(b->codes_out, device, "codes_out"); + mk2_stage0_device_pointer(b->phoneme_out, device, "phoneme_out"); + if (b->moe_w_inject_debug) + mk2_stage0_device_pointer(b->moe_w_inject_debug, device, "moe_w_inject_debug"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && !mk2_stage0_same_pinned_bindings(st->pinned_bindings, *b)) + throw std::runtime_error("full Stage 0 captured state bindings changed; create a new handle"); + cudaStream_t s = static_cast(stream); + Backbone &hb = st->host_bb; + const bool first_bind = !st->bindings_initialized; + if (first_bind) { + for (int i = 0; i < 14; ++i) { + st->mb_host[i].conv_state = (bf16 *)b->conv_state[i]; + st->mb_host[i].conv_slot_stride = b->conv_slot_stride[i]; + st->mb_host[i].ssm_state = (float *)b->ssm_state[i]; + st->mb_host[i].ssm_slot_stride = b->ssm_slot_stride[i]; + } + for (int i = 0; i < 3; ++i) { + st->at_host[i].pool_base = (bf16 *)b->kv_pool_base[i]; + st->at_host[i].block_table = (const int *)b->kv_block_table[i]; + st->at_host[i].max_blocks_per_seq = b->kv_max_blocks_per_seq; + st->at_host[i].tokens_per_block = b->kv_tokens_per_block; + } + hb.moe_w_inject = (const float *)b->moe_w_inject_debug; + hb.slot_map = (const int *)b->slot_map; + hb.seq_len = (const int *)b->seq_len; + hb.phoneme_out = (int *)b->phoneme_out; + memcpy(st->mb_stage, st->mb_host.data(), 14 * sizeof(MambaLayer)); + memcpy(st->at_stage, st->at_host.data(), 3 * sizeof(AttnLayer)); + memcpy(st->bb_stage, &hb, sizeof(Backbone)); + st->pinned_bindings = *b; + st->bindings_initialized = true; + } + if (!native || first_bind) { + mk2_check(cudaMemcpyAsync(const_cast(hb.mamba), st->mb_stage, + 14 * sizeof(MambaLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(const_cast(hb.attn), st->at_stage, + 3 * sizeof(AttnLayer), cudaMemcpyHostToDevice, s)); + mk2_check(cudaMemcpyAsync(st->d_bb, st->bb_stage, sizeof(Backbone), + cudaMemcpyHostToDevice, s)); + } + if (!native) { + mk2_check(cudaMemcpyAsync(st->host_bb.h, b->embedding, + (size_t)BATCH * D_MODEL * sizeof(bf16), + cudaMemcpyDeviceToDevice, s)); + mk2_check(cudaMemsetAsync(st->host_bb.r, 0, + (size_t)BATCH * D_MODEL * sizeof(bf16), s)); + } + st->pred.g.native_embedding = native ? (const bf16 *)b->embedding : nullptr; + st->pred.g.native_codes = native ? (long long *)b->codes_out : nullptr; + st->pred.g.native_seq_len = native ? (int *)b->seq_len : nullptr; + st->pred.g.gumbel = static_cast(b->gumbel); + st->pred.g.temperature = b->temperature; + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + mk2_check(cudaStreamIsCapturing(s, &capture_status)); + if (st->pred.grid_reset_pending || capture_status != cudaStreamCaptureStatusNone) { + mk2_check(cudaMemsetAsync(st->pred.g.barrier, 0, 64 * sizeof(unsigned int), s)); + st->pred.grid_reset_pending = false; + } + predictor_mk2<<pred.grid, THREADS, st->pred.smem, s>>>( + st->pred.g, st->d_bb, N_LAYERS, N_POS, BB_LAYERS, nullptr, + native ? st->pred.native_predictor_active_mask : nullptr + MK2_PROGRESS_ARGUMENT(native ? st->pred.native_progress_mask : nullptr)); + mk2_check(cudaPeekAtLastError()); + if (!native) + codes_to_int64<<<(BATCH * N_POS + 255) / 256, 256, 0, s>>>( + st->pred.g.sampled, static_cast(b->codes_out), BATCH * N_POS); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { g_last_error = cudaGetErrorString(e); return 1; } + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown full Stage 0 enqueue failure"; return 1; } +} +extern "C" int mk2_run_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, false); +} +extern "C" int mk2_run_native_stage0(MK2Predictor *p, const MK2Stage0Bindings *b, + void *stream) { + return mk2_enqueue_stage0(p, b, stream, true); +} +extern "C" int mk2_set_native_predictor_active_mask(MK2Predictor *p, + const unsigned char *mask) { + try { + if (!p || reinterpret_cast(p)->kind != 1) + throw std::runtime_error("startup mask requires a full Stage0 handle"); + auto *st = reinterpret_cast(p); + if (st->bindings_initialized && st->pred.native_predictor_active_mask != mask) + throw std::runtime_error("startup mask pointer cannot change after first Stage0 enqueue"); + if (mask) + mk2_stage0_device_pointer(mask, st->pred.device, "native_predictor_active_mask"); + st->pred.native_predictor_active_mask = mask; + g_last_error.clear(); + return 0; + } catch (const std::exception &e) { g_last_error = e.what(); return 1; } + catch (...) { g_last_error = "unknown startup mask binding failure"; return 1; } +} +extern "C" int mk2_native_startup_skip_abi_version() { return 1; } +#include "../../inflight/runtime.cuh" +extern "C" unsigned long long mk2_stage0_launches(MK2Predictor *p) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches) return 0ull; + unsigned long long n[2] = {0ull, 0ull}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return 0ull; + return (n[1] << 32) | (n[0] & 0xffffffffull); +} +extern "C" int mk2_stage0_timing(MK2Predictor *p, unsigned long long *out) { + auto *st = reinterpret_cast(p); + if (!st || st->pred.kind != 1 || !st->host_bb.launches || !out) return -1; + unsigned long long n[6] = {0, 0, 0, 0, 0, 0}; + if (cudaMemcpy(n, st->host_bb.launches, sizeof(n), + cudaMemcpyDeviceToHost) != cudaSuccess) + return -2; + out[0] = n[3]; + out[1] = n[4]; + out[2] = n[5]; + out[3] = n[0]; + return 0; +} +extern "C" void mk2_destroy(MK2Predictor *p) { + if (!p) return; + if (reinterpret_cast(p)->kind == 1) { + auto *st = reinterpret_cast(p); + cudaError_t first_error = cudaSuccess; + auto record = [&](cudaError_t e) { + if (first_error == cudaSuccess && e != cudaSuccess) first_error = e; + }; + int previous = -1; + const cudaError_t current = cudaGetDevice(&previous); + record(current); + record(cudaSetDevice(st->pred.device)); + for (void *owned : st->owned) record(cudaFree(owned)); + for (void *owned : st->pred.owned) record(cudaFree(owned)); + for (void *pinned : st->pinned) record(cudaFreeHost(pinned)); + if (current == cudaSuccess) record(cudaSetDevice(previous)); + delete st; + if (first_error != cudaSuccess) g_last_error = cudaGetErrorString(first_error); + } else { + auto *st = reinterpret_cast(p); + int previous = st->device; + cudaGetDevice(&previous); + cudaSetDevice(st->device); + for (void *owned : st->owned) { + const auto e = cudaFree(owned); + if (e != cudaSuccess) g_last_error = cudaGetErrorString(e); + } + cudaSetDevice(previous); + delete st; + } +} +extern "C" int mk2_batch(const MK2Predictor *p) { + return p ? reinterpret_cast(p)->batch : 0; +} +extern "C" int mk2_positions(const MK2Predictor *) { return N_POS; } +extern "C" const char *mk2_last_error(void) { return g_last_error.c_str(); } +#include "native_prefill_impl.inc" +extern "C" int mk2_set_native_stage0_grid(MK2Predictor *p, int grid) { + if (!p) { g_last_error = "grid setter requires full Stage0 handle"; return 1; } + auto *st = reinterpret_cast(p); + if (st->kind != 1 || grid < 1 || grid > st->maximum_grid) { + g_last_error = "requested Stage0 grid outside independently resident capacity"; + return 1; + } + if (st->grid != grid) st->grid_reset_pending = true; + st->grid = grid; + g_last_error.clear(); + return 0; +} diff --git a/deploy/nemotron_native/kernels/stage0/small/sparse_moe_impl.cuh b/deploy/nemotron_native/kernels/stage0/small/sparse_moe_impl.cuh new file mode 100644 index 0000000..056890d --- /dev/null +++ b/deploy/nemotron_native/kernels/stage0/small/sparse_moe_impl.cuh @@ -0,0 +1,139 @@ +__device__ void sparse_moe_pack(const Backbone *bb) { + const SparseMoe &s = *bb->sparse_moe; + __shared__ int count; + __shared__ int local_rows[BATCH]; + const int experts_per_wave = min(int(gridDim.x), BB_EXPERTS); + const int first_expert = int(blockIdx.x) % experts_per_wave; + const int rank = int(blockIdx.x) / experts_per_wave; + const int team_size = (int(gridDim.x) + experts_per_wave - 1 - first_expert) / + experts_per_wave; + for (int e = first_expert; e < BB_EXPERTS; e += experts_per_wave) { + if (threadIdx.x == 0) { + int n = 0; + for (int row = 0; row < BATCH; ++row) + if (bb->moe_w[row * BB_EXPERTS + e] != 0.f) + local_rows[n++] = row; + count = n; + if (rank == 0) { + s.counts[e] = n; + for (int i = 0; i < n; ++i) + s.rows[e * BATCH_PAD + i] = local_rows[i]; + } + } + __syncthreads(); + constexpr int chunks = D_MODEL / 8; + const int padded = ((count + SPARSE_N - 1) / SPARSE_N) * SPARSE_N; + for (int i = rank * THREADS + threadIdx.x; i < padded * chunks; + i += team_size * THREADS) { + const int row = i / chunks, col = i % chunks; + const int packed = e * BATCH_PAD + row; + auto *dst = reinterpret_cast(s.gathered + + (size_t)packed * D_MODEL + col * 8); + if (row < count) { + const int original = local_rows[row]; + *dst = *reinterpret_cast(bb->normed + + (size_t)original * D_MODEL + col * 8); + } else { + *dst = make_float4(0.f, 0.f, 0.f, 0.f); + } + } + __syncthreads(); + } +} +template +__device__ void sparse_moe_gemm( + const weight_gl &w, const sparse_act_gl &a, const SparseMoe &s, + bf16 *hidden, const float *routes, float *contributions, + weight_tile (&w_s)[STAGES], act_tile (&a_s)[STAGES], out_tile &o_s, + semaphore *filled, semaphore *drained, int &iter_io, int &tile_io) { + constexpr int width = DOWN ? D_MODEL : BB_MOE_INTER; + constexpr int k_total = DOWN ? BB_MOE_INTER : D_MODEL; + constexpr int out_tiles = width / OUT_ATOM; + constexpr int iters = k_total / KCHUNK; + static_assert(BATCH_PAD >= SPARSE_N && BATCH_PAD % SPARSE_N == 0); + static_assert(sizeof(sparse_act_tile) <= sizeof(act_tile)); + static_assert(sizeof(sparse_out_tile) <= sizeof(out_tile)); + __shared__ int prefix[BB_EXPERTS + 1]; + __shared__ int row_tiles[BB_EXPERTS]; + if (threadIdx.x == 0) { + int total = 0; + for (int e = 0; e < BB_EXPERTS; ++e) { + prefix[e] = total; + row_tiles[e] = (s.counts[e] + SPARSE_N - 1) / SPARSE_N; + total += row_tiles[e] * out_tiles; + } + prefix[BB_EXPERTS] = total; + } + __syncthreads(); + const int n_tiles = prefix[BB_EXPERTS]; + int start = int(blockIdx.x) - tile_io % int(gridDim.x); + if (start < 0) start += gridDim.x; + tile_io += n_tiles; + const int mine = start < n_tiles ? 1 + (n_tiles - 1 - start) / gridDim.x : 0; + const int iter_base = iter_io; + iter_io += mine * iters; + if (warpgroup::groupid() == 1) { + if (warpgroup::laneid() == 0) { + int iter = iter_base; + for (int f = start; f < n_tiles; f += gridDim.x) { + int e = 0; + while (prefix[e + 1] <= f) ++e; + const int rem = f - prefix[e], nt = row_tiles[e]; + const int feature_tile = rem / nt; + const int token_tile = rem - feature_tile * nt; + const int weight_row = e * out_tiles + feature_tile; + const int act_row = e * (BATCH_PAD / SPARSE_N) + token_tile; + for (int k = 0; k < iters; ++k, ++iter) { + const int stage = iter % STAGES, phase = (iter / STAGES) & 1; + wait(drained[stage], phase); + tma::expect_bytes(filled[stage], sizeof(weight_tile) + sizeof(sparse_act_tile)); + tma::load_async<2, cache_policy::EVICT_FIRST>( + w_s[stage], w, {weight_row, k}, filled[stage]); + tma::load_async<2, cache_policy::EVICT_LAST>( + *reinterpret_cast(&a_s[stage]), + a, {act_row, k}, filled[stage]); + } + } + } + } else if (warpgroup::groupid() == 0) { + int iter = iter_base; + for (int f = start; f < n_tiles; f += gridDim.x) { + int e = 0; + while (prefix[e + 1] <= f) ++e; + const int rem = f - prefix[e], nt = row_tiles[e]; + const int feature_tile = rem / nt; + const int token_tile = rem - feature_tile * nt; + rt_fl<16, SPARSE_N> accum; + warp::zero(accum); + for (int k = 0; k < iters; ++k, ++iter) { + const int stage = iter % STAGES, phase = (iter / STAGES) & 1; + wait(filled[stage], phase); + warpgroup::mma_ABt(accum, w_s[stage], + *reinterpret_cast(&a_s[stage])); + warpgroup::mma_async_wait(); + if (warp::laneid() == 0) arrive(drained[stage]); + } + sparse_out_tile &out = *reinterpret_cast(&o_s); + warpgroup::store(out, accum); + warpgroup::sync(1); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + for (int i = warpgroup::laneid(); i < OUT_ATOM * SPARSE_N; i += 128) { + const int row = i / OUT_ATOM, col = i % OUT_ATOM; + const int local = token_tile * SPARSE_N + row; + const int packed = e * BATCH_PAD + local; + const int output = feature_tile * OUT_ATOM + col; + const float value = float(out[{col, row}]); + if constexpr (!DOWN) { + hidden[(size_t)packed * BB_MOE_INTER + output] = + bf16(value / (1.f + __expf(-value))); + } else if (local < s.counts[e]) { + const int original = s.rows[packed]; + const float weight = routes[original * BB_EXPERTS + e]; + contributions[((size_t)original * (BB_EXPERTS + 1) + e) * + D_MODEL + output] = __fmul_rn(weight, value); + } + } + warpgroup::sync(1); + } + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/LICENSE b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/LICENSE new file mode 100644 index 0000000..79e0fdb --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024-2026 HazyResearch + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/common/base_ops.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/common/base_ops.cuh new file mode 100644 index 0000000..a5a36da --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/common/base_ops.cuh @@ -0,0 +1,446 @@ + + + + + +#pragma once + +#include "base_types.cuh" + +namespace kittens { + + + + + + +namespace base_ops { + + + + + + + + + + + +struct zero { + template __device__ static inline constexpr T op(args... _) { return base_types::constants::zero(); } +}; + + + + + + + + +struct one { + template __device__ static inline constexpr T op(args... _) { return base_types::constants::one(); } +}; + + + + + + + + +struct pos_infty { + template __device__ static inline constexpr T op(args... _) { return base_types::constants::pos_infty(); } +}; + + + + + + + + +struct neg_infty { + template __device__ static inline constexpr T op(args... _) { return base_types::constants::neg_infty(); } +}; + + + + + + + + + + + + + +struct exp { + template static __device__ inline T op(const T &x) { return exp(x); } +}; +template<> __device__ inline float exp::op (const float &x ) { return __expf(x); } +template<> __device__ inline float2 exp::op(const float2 &x) { return float2{__expf(x.x), __expf(x.y)}; } +template<> __device__ inline bf16 exp::op (const bf16 &x ) { return hexp(x); } +template<> __device__ inline bf16_2 exp::op(const bf16_2 &x) { return h2exp(x); } +template<> __device__ inline half exp::op (const half &x ) { return hexp(x); } +template<> __device__ inline half_2 exp::op(const half_2 &x) { return h2exp(x); } + + + + + + + + + + +struct exp2 { + template static __device__ inline T op(const T &x) { return exp2f(x); } +}; +template<> __device__ inline float exp2::op (const float &x ) { return exp2f(x); } +template<> __device__ inline float2 exp2::op(const float2 &x) { return float2{exp2f(x.x), exp2f(x.y)}; } +template<> __device__ inline bf16 exp2::op (const bf16 &x ) { return hexp2(x); } +template<> __device__ inline bf16_2 exp2::op(const bf16_2 &x) { return h2exp2(x); } +template<> __device__ inline half exp2::op (const half &x ) { return hexp2(x); } +template<> __device__ inline half_2 exp2::op(const half_2 &x) { return h2exp2(x); } + + + + + + + + + +struct log { + template static __device__ inline T op(const T &x) { return log(x); } +}; +template<> __device__ inline float log::op (const float &x ) { return __logf(x); } +template<> __device__ inline float2 log::op(const float2 &x) { return float2{__logf(x.x), __logf(x.y)}; } +template<> __device__ inline bf16 log::op (const bf16 &x ) { return hlog(x); } +template<> __device__ inline bf16_2 log::op(const bf16_2 &x) { return h2log(x); } +template<> __device__ inline half log::op (const half &x ) { return hlog(x); } +template<> __device__ inline half_2 log::op(const half_2 &x) { return h2log(x); } + + + + + + + + + +struct log2 { + template static __device__ inline T op(const T &x) { return log2(x); } +}; +template<> __device__ inline float log2::op (const float &x ) { return __log2f(x); } +template<> __device__ inline float2 log2::op(const float2 &x) { return float2{__log2f(x.x), __log2f(x.y)}; } +template<> __device__ inline bf16 log2::op (const bf16 &x ) { return hlog2(x); } +template<> __device__ inline bf16_2 log2::op(const bf16_2 &x) { return h2log2(x); } +template<> __device__ inline half log2::op (const half &x ) { return hlog2(x); } +template<> __device__ inline half_2 log2::op(const half_2 &x) { return h2log2(x); } + + + + + + + + + +struct abs { + template static __device__ inline T op(const T &x) { return abs(x); } +}; +template<> __device__ inline float abs::op (const float &x ) { return fabsf(x); } +template<> __device__ inline float2 abs::op(const float2 &x) { return float2{fabsf(x.x), fabsf(x.y)}; } +template<> __device__ inline bf16 abs::op (const bf16 &x ) { return __habs(x); } +template<> __device__ inline bf16_2 abs::op(const bf16_2 &x) { return __habs2(x); } +template<> __device__ inline half abs::op (const half &x ) { return __habs(x); } +template<> __device__ inline half_2 abs::op(const half_2 &x) { return __habs2(x); } + + + + + + + + + + +struct relu { + template static __device__ inline T op(const T &x) { return max(x, base_types::constants::zero()); } +}; +template<> __device__ inline float relu::op (const float &x ) { return max(x, 0.f); } +template<> __device__ inline float2 relu::op(const float2 &x) { return float2{max(x.x, 0.f), max(x.y, 0.f)}; } +template<> __device__ inline bf16 relu::op (const bf16 &x ) { return __hmax(x, base_types::constants::zero()); } +template<> __device__ inline bf16_2 relu::op(const bf16_2 &x) { return __hmax2(x, base_types::constants::zero()); } +template<> __device__ inline half relu::op (const half &x ) { return __hmax(x, base_types::constants::zero()); } +template<> __device__ inline half_2 relu::op(const half_2 &x) { return __hmax2(x, base_types::constants::zero()); } + + + + + + + +struct neg { + template static __device__ inline T op(const T &x) { return -x; } +}; +template<> __device__ inline float neg::op (const float &x ) { return -x; } +template<> __device__ inline float2 neg::op (const float2 &x ) { return float2{-x.x, -x.y}; } +template<> __device__ inline bf16 neg::op (const bf16 &x ) { return __hneg(x); } +template<> __device__ inline bf16_2 neg::op (const bf16_2 &x ) { return __hneg2(x); } +template<> __device__ inline half neg::op (const half &x ) { return __hneg(x); } +template<> __device__ inline half_2 neg::op (const half_2 &x ) { return __hneg2(x); } + + + + + + + +struct sqrt { + template static __device__ inline T op(const T &x) { return sqrt(x); } +}; +template<> __device__ inline float sqrt::op (const float &x ) { return sqrtf(x); } +template<> __device__ inline float2 sqrt::op (const float2 &x ) { return float2{sqrtf(x.x), sqrtf(x.y)}; } +template<> __device__ inline bf16 sqrt::op (const bf16 &x ) { return hsqrt(x); } +template<> __device__ inline bf16_2 sqrt::op (const bf16_2 &x ) { return h2sqrt(x); } +template<> __device__ inline half sqrt::op (const half &x ) { return hsqrt(x); } +template<> __device__ inline half_2 sqrt::op (const half_2 &x ) { return h2sqrt(x); } + + + + + + + +struct rsqrt { + template static __device__ inline T op(const T &x) { return rsqrt(x); } +}; +template<> __device__ inline float rsqrt::op (const float &x ) { return rsqrtf(x); } +template<> __device__ inline float2 rsqrt::op (const float2 &x ) { return float2{rsqrtf(x.x), rsqrtf(x.y)}; } +template<> __device__ inline bf16 rsqrt::op (const bf16 &x ) { return hrsqrt(x); } +template<> __device__ inline bf16_2 rsqrt::op (const bf16_2 &x ) { return h2rsqrt(x); } +template<> __device__ inline half rsqrt::op (const half &x ) { return hrsqrt(x); } +template<> __device__ inline half_2 rsqrt::op (const half_2 &x ) { return h2rsqrt(x); } + + + + + + + + + +struct copy { + template static __device__ inline T op(const T &a) { return a; } +}; + + + + + + + + + + + + + + +struct copy2 { + template static __device__ inline T op(const T &a, const T &b) { return b; } +}; + + + + + + + + + + +struct sum { + template static __device__ inline T op(const T &a, const T &b) { return a+b; } +}; +template<> __device__ inline float2 sum::op(const float2 &a, const float2 &b) { +#ifdef KITTENS_SM10X + float2 c; + asm volatile("add.f32x2 %0, %1, %2;" : "=l"(*(uint64_t*)&c) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b)); + return c; +#else + return float2{a.x+b.x, a.y+b.y}; +#endif +} +template<> __device__ inline bf16 sum::op (const bf16 &a, const bf16 &b) { return __hadd(a, b); } +template<> __device__ inline bf16_2 sum::op(const bf16_2 &a, const bf16_2 &b) { return __hadd2(a, b); } +template<> __device__ inline half sum::op (const half &a, const half &b) { return __hadd(a, b); } +template<> __device__ inline half_2 sum::op(const half_2 &a, const half_2 &b) { return __hadd2(a, b); } + + + + + + + + + + +struct sub { + template static __device__ inline T op(const T &a, const T &b) { return a-b; } +}; +template<> __device__ inline float2 sub::op(const float2 &a, const float2 &b) { +#ifdef KITTENS_SM10X + float2 c; + asm volatile("sub.f32x2 %0, %1, %2;" : "=l"(*(uint64_t*)&c) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b)); + return c; +#else + return float2{a.x-b.x, a.y-b.y}; +#endif +} +template<> __device__ inline bf16 sub::op (const bf16 &a, const bf16 &b) { return __hsub(a, b); } +template<> __device__ inline bf16_2 sub::op(const bf16_2 &a, const bf16_2 &b) { return __hsub2(a, b); } +template<> __device__ inline half sub::op (const half &a, const half &b) { return __hsub(a, b); } +template<> __device__ inline half_2 sub::op(const half_2 &a, const half_2 &b) { return __hsub2(a, b); } + + + + + + + + + + +struct mul { + template static __device__ inline T op(const T &a, const T &b) { return a*b; } +}; +template<> __device__ inline float2 mul::op(const float2 &a, const float2 &b) { +#ifdef KITTENS_SM10X + float2 c; + asm volatile("mul.f32x2 %0, %1, %2;" : "=l"(*(uint64_t*)&c) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b)); + return c; +#else + return float2{a.x*b.x, a.y*b.y}; +#endif +} +template<> __device__ inline bf16 mul::op (const bf16 &a, const bf16 &b) { return __hmul(a, b); } +template<> __device__ inline bf16_2 mul::op(const bf16_2 &a, const bf16_2 &b) { return __hmul2(a, b); } +template<> __device__ inline half mul::op (const half &a, const half &b) { return __hmul(a, b); } +template<> __device__ inline half_2 mul::op(const half_2 &a, const half_2 &b) { return __hmul2(a, b); } + + + + + + + + + + +struct div { + template static __device__ inline T op(const T &a, const T &b) { return a/b; } +}; +template<> __device__ inline float2 div::op(const float2 &a, const float2 &b) { return float2{a.x/b.x, a.y/b.y}; } +template<> __device__ inline bf16 div::op (const bf16 &a, const bf16 &b) { return __hdiv(a, b); } +template<> __device__ inline bf16_2 div::op(const bf16_2 &a, const bf16_2 &b) { return __h2div(a, b); } +template<> __device__ inline half div::op (const half &a, const half &b) { return __hdiv(a, b); } +template<> __device__ inline half_2 div::op(const half_2 &a, const half_2 &b) { return __h2div(a, b); } + + + + + + + + + + + struct max { + template static __device__ inline T op(const T &a, const T &b) { return ::max(a, b); } +}; +template<> __device__ inline float2 max::op(const float2 &a, const float2 &b) { return float2{::max(a.x, b.x), ::max(a.y, b.y)}; } +template<> __device__ inline bf16 max::op (const bf16 &a, const bf16 &b) { return __hmax(a, b); } +template<> __device__ inline bf16_2 max::op(const bf16_2 &a, const bf16_2 &b) { return __hmax2(a, b); } +template<> __device__ inline half max::op (const half &a, const half &b) { return __hmax(a, b); } +template<> __device__ inline half_2 max::op(const half_2 &a, const half_2 &b) { return __hmax2(a, b); } + + + + + + + + + + +struct min { + template static __device__ inline T op(const T &a, const T &b) { return ::min(a, b); } +}; +template<> __device__ inline float2 min::op(const float2 &a, const float2 &b) { return float2{::min(a.x, b.x), ::min(a.y, b.y)}; } +template<> __device__ inline bf16 min::op (const bf16 &a, const bf16 &b) { return __hmin(a, b); } +template<> __device__ inline bf16_2 min::op(const bf16_2 &a, const bf16_2 &b) { return __hmin2(a, b); } +template<> __device__ inline half min::op (const half &a, const half &b) { return __hmin(a, b); } +template<> __device__ inline half_2 min::op(const half_2 &a, const half_2 &b) { return __hmin2(a, b); } + + + + + + + + + + + + + + + +struct fma_AxBtC { + template static __device__ inline T op(const T &a, const T &b, const T &c) { + return sum::op(mul::op(a, b), c); + } +}; +template<> __device__ inline float2 fma_AxBtC::op(const float2 &a, const float2 &b, const float2 &c) { +#ifdef KITTENS_SM10X + float2 d; + asm volatile("fma.rn.f32x2 %0, %1, %2, %3;" : "=l"(*(uint64_t*)&d) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b), "l"(*(uint64_t*)&c)); + return d; +#else + return float2{a.x*b.x+c.x, a.y*b.y+c.y}; +#endif +} + + + + + + + + + + + + +struct fma_AxCtB { + template static __device__ inline T op(const T &a, const T &b, const T &c) { + return sum::op(mul::op(a, c), b); + } +}; +template<> __device__ inline float2 fma_AxCtB::op(const float2 &a, const float2 &b, const float2 &c) { +#ifdef KITTENS_SM10X + float2 d; + asm volatile("fma.rn.f32x2 %0, %1, %2, %3;" : "=l"(*(uint64_t*)&d) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&c), "l"(*(uint64_t*)&b)); + return d; +#else + return float2{a.x*c.x+b.x, a.y*c.y+b.y}; +#endif +} + +} + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/common/base_types.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/common/base_types.cuh new file mode 100644 index 0000000..3a8efbb --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/common/base_types.cuh @@ -0,0 +1,777 @@ + + + + + + + + + + +#pragma once + +namespace kittens { + + + + +using bf16 = __nv_bfloat16; + + + +using half = __half; + + + +using int8 = signed char; + + + +using uint8 = unsigned char; + + + + +using bf16_2 = __nv_bfloat162; + + + +using half_2 = __half2; + + + + +using int8_2 = char2; +using int8_4 = char4; + + + +using uint8_2 = uchar2; +using uint8_4 = uchar4; + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + +using fp8e4m3 = __nv_fp8_e4m3; +using fp8e5m2 = __nv_fp8_e5m2; +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) +using fp8e8m0 = __nv_fp8_e8m0; +#endif + + + +using fp8e4m3_2 = __nv_fp8x2_e4m3; +using fp8e5m2_2 = __nv_fp8x2_e5m2; +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) +using fp8e8m0_2 = __nv_fp8x2_e8m0; +#endif + + + +using fp8e4m3_4 = __nv_fp8x4_e4m3; +using fp8e5m2_4 = __nv_fp8x4_e5m2; +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) +using fp8e8m0_4 = __nv_fp8x4_e8m0; +#endif +#endif + +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + +using fp4e2m1 = __nv_fp4_e2m1; + + + +using fp4e2m1_2 = __nv_fp4x2_e2m1; + + + +using fp4e2m1_4 = __nv_fp4x4_e2m1; +#endif + +namespace ducks { + + + + + +namespace base_types { + +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template +concept T2 = std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v; +template +concept T1 = std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v; +#elif defined(KITTENS_SM90) +template +concept T2 = std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v; +template +concept T1 = std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v; +#else +template +concept T2 = std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v; +template +concept T1 = std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v; +#endif + +} +} + + + + + + +namespace base_types { + + + + + + +template struct constants { + + + + + static __device__ inline constexpr T zero() { return T{0}; } + + + + + static __device__ inline constexpr T one() { return T{1}; } + + + + + static __device__ inline constexpr T pos_infty() { static_assert(sizeof(T) == 9999, "pos_infty not defined for this type"); return T{}; } + + + + + static __device__ inline constexpr T neg_infty() { static_assert(sizeof(T) == 9999, "neg_infty not defined for this type"); return T{}; } +}; +template<> struct constants { + static __device__ inline constexpr float zero() { return 0.f; } + static __device__ inline constexpr float one() { return 1.f; } + static __device__ inline constexpr float pos_infty() { return std::bit_cast(uint32_t(0x7f800000)); } + static __device__ inline constexpr float neg_infty() { return std::bit_cast(uint32_t(0xff800000)); } +}; +template<> struct constants { + static __device__ inline constexpr float2 zero() { return float2{0.f, 0.f}; } + static __device__ inline constexpr float2 one() { return float2{1.f, 1.f}; } + static __device__ inline constexpr float2 pos_infty() { return float2{constants::pos_infty(), constants::pos_infty()}; } + static __device__ inline constexpr float2 neg_infty() { return float2{constants::neg_infty(), constants::neg_infty()}; } +}; +template<> struct constants { + static __device__ inline constexpr bf16 zero() { return std::bit_cast<__nv_bfloat16>(uint16_t(0x0000)); } + static __device__ inline constexpr bf16 one() { return std::bit_cast<__nv_bfloat16>(uint16_t(0x3F80)); } + static __device__ inline constexpr bf16 pos_infty() { return std::bit_cast<__nv_bfloat16>(uint16_t(0x7F80)); } + static __device__ inline constexpr bf16 neg_infty() { return std::bit_cast<__nv_bfloat16>(uint16_t(0xFF80)); } +}; +template<> struct constants { + static __device__ inline constexpr bf16_2 zero() { return bf16_2{constants::zero(), constants::zero()}; } + static __device__ inline constexpr bf16_2 one() { return bf16_2{constants::one(), constants::one()}; } + static __device__ inline constexpr bf16_2 pos_infty() { return bf16_2{constants::pos_infty(), constants::pos_infty()}; } + static __device__ inline constexpr bf16_2 neg_infty() { return bf16_2{constants::neg_infty(), constants::neg_infty()}; } +}; +template<> struct constants { + static __device__ inline constexpr half zero() { return std::bit_cast<__half>(uint16_t(0x0000)); } + static __device__ inline constexpr half one() { return std::bit_cast<__half>(uint16_t(0x3C00)); } + static __device__ inline constexpr half pos_infty() { return std::bit_cast<__half>(uint16_t(0x7C00)); } + static __device__ inline constexpr half neg_infty() { return std::bit_cast<__half>(uint16_t(0xFC00)); } +}; +template<> struct constants { + static __device__ inline constexpr half_2 zero() { return half_2{constants::zero(), constants::zero()}; } + static __device__ inline constexpr half_2 one() { return half_2{constants::one(), constants::one()}; } + static __device__ inline constexpr half_2 pos_infty() { return half_2{constants::pos_infty(), constants::pos_infty()}; } + static __device__ inline constexpr half_2 neg_infty() { return half_2{constants::neg_infty(), constants::neg_infty()}; } +}; +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template<> struct constants { + static __device__ inline constexpr fp8e4m3 zero() { return std::bit_cast<__nv_fp8_e4m3>(uint8_t(0x00)); } + static __device__ inline constexpr fp8e4m3 one() { return std::bit_cast<__nv_fp8_e4m3>(uint8_t(0x38)); } +}; +template<> struct constants { + static __device__ inline constexpr fp8e4m3_2 zero() { return std::bit_cast(uint16_t(0x0000)); } + static __device__ inline constexpr fp8e4m3_2 one() { return std::bit_cast(uint16_t(0x3838)); } +}; +template<> struct constants { + static __device__ inline constexpr fp8e4m3_4 zero() { return std::bit_cast(uint32_t(0x00000000)); } + static __device__ inline constexpr fp8e4m3_4 one() { return std::bit_cast(uint32_t(0x38383838)); } +}; +template<> struct constants { + static __device__ inline constexpr fp8e5m2 zero() { return std::bit_cast<__nv_fp8_e5m2>(uint8_t(0x00)); } + static __device__ inline constexpr fp8e5m2 one() { return std::bit_cast<__nv_fp8_e5m2>(uint8_t(0x3C)); } +}; +template<> struct constants { + static __device__ inline constexpr fp8e5m2_2 zero() { return std::bit_cast(uint16_t(0x0000)); } + static __device__ inline constexpr fp8e5m2_2 one() { return std::bit_cast(uint16_t(0x3C3C)); } +}; +template<> struct constants { + static __device__ inline constexpr fp8e5m2_4 zero() { return std::bit_cast(uint32_t(0x00000000)); } + static __device__ inline constexpr fp8e5m2_4 one() { return std::bit_cast(uint32_t(0x3C3C3C3C)); } +}; +#endif +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template<> struct constants { + static __device__ inline constexpr fp4e2m1 zero() { return std::bit_cast(uint8_t(0x00)); } + static __device__ inline constexpr fp4e2m1 one() { return std::bit_cast(uint8_t(0x02)); } +}; +template<> struct constants { + static __device__ inline constexpr fp4e2m1_2 zero() { return std::bit_cast(uint8_t(0x00)); } + static __device__ inline constexpr fp4e2m1_2 one() { return std::bit_cast(uint8_t(0x22)); } +}; +template<> struct constants { + static __device__ inline constexpr fp4e2m1_4 zero() { return std::bit_cast(uint16_t(0x0000)); } + static __device__ inline constexpr fp4e2m1_4 one() { return std::bit_cast(uint16_t(0x2222)); } +}; +template<> struct constants { + static __device__ inline constexpr fp8e8m0 zero() { return std::bit_cast(uint8_t(0x00)); } + static __device__ inline constexpr fp8e8m0 one() { return std::bit_cast(uint8_t(0x7F)); } +}; +template<> struct constants { + static __device__ inline constexpr fp8e8m0_2 zero() { return std::bit_cast(uint16_t(0x0000)); } + static __device__ inline constexpr fp8e8m0_2 one() { return std::bit_cast(uint16_t(0x7F7F)); } +}; +template<> struct constants { + static __device__ inline constexpr fp8e8m0_4 zero() { return std::bit_cast(uint32_t(0x00000000)); } + static __device__ inline constexpr fp8e8m0_4 one() { return std::bit_cast(uint32_t(0x7F7F7F7F)); } +}; +#endif + +template<> struct constants { + static __device__ inline constexpr int zero() { return 0; } + static __device__ inline constexpr int one() { return 1; } +}; +template<> struct constants { + static __device__ inline constexpr int2 zero() { return int2{0, 0}; } + static __device__ inline constexpr int2 one() { return int2{1, 1}; } +}; + + + + + + +template struct packing { + + + + + + static __device__ inline constexpr int num() { return 1; } + + + + + + + static __device__ inline constexpr T pack(const bf16 &i); +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = bf16; + using packed_type = bf16_2; + static __device__ inline constexpr bf16_2 pack(const bf16 &i) { return bf16_2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = bf16; + using packed_type = bf16_2; + static __device__ inline constexpr bf16_2 pack(const bf16 &i) { return bf16_2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = half; + using packed_type = half_2; + static __device__ inline constexpr half_2 pack(const half &i) { return half_2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = half; + using packed_type = half_2; + static __device__ inline constexpr half_2 pack(const half &i) { return half_2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = float; + using packed_type = float2; + static __device__ inline constexpr float2 pack(const float &i) { return float2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = float; + using packed_type = float2; + static __device__ inline constexpr float2 pack(const float &i) { return float2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = double; + using packed_type = double2; + static __device__ inline constexpr double2 pack(const double &i) { return double2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = double; + using packed_type = double2; + static __device__ inline constexpr double2 pack(const double &i) { return double2{i, i}; } +}; +struct bool_2 { bool x, y; }; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = bool; + using packed_type = bool_2; + static __device__ inline constexpr bool_2 pack(const bool &i) { return bool_2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = bool; + using packed_type = bool_2; + static __device__ inline constexpr bool_2 pack(const bool &i) { return bool_2{i, i}; } +}; + +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = int8; + using packed_type = int8_4; + static __device__ inline constexpr int8_4 pack(const int8 &i) { return int8_4{i, i, i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 4; } + using unpacked_type = int8; + using packed_type = int8_4; + static __device__ inline constexpr int8_4 pack(const int8 &i) { return int8_4{i, i, i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = uint8; + using packed_type = uint8_4; + static __device__ inline constexpr uint8_4 pack(const uint8 &i) { return uint8_4{i, i, i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 4; } + using unpacked_type = uint8; + using packed_type = uint8_4; + static __device__ inline constexpr uint8_4 pack(const uint8 &i) { return uint8_4{i, i, i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = short; + using packed_type = short2; + static __device__ inline constexpr short2 pack(const short &i) { return short2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = short; + using packed_type = short2; + static __device__ inline constexpr short2 pack(const short &i) { return short2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = unsigned short; + using packed_type = ushort2; + static __device__ inline constexpr ushort2 pack(const unsigned short &i) { return ushort2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = unsigned short; + using packed_type = ushort2; + static __device__ inline constexpr ushort2 pack(const unsigned short &i) { return ushort2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = int; + using packed_type = int2; + static __device__ inline constexpr int2 pack(const int &i) { return int2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = int; + using packed_type = int2; + static __device__ inline constexpr int2 pack(const int &i) { return int2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = uint; + using packed_type = uint2; + static __device__ inline constexpr uint2 pack(const uint &i) { return uint2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = uint; + using packed_type = uint2; + static __device__ inline constexpr uint2 pack(const uint &i) { return uint2{i, i}; } +}; +struct int64_2 { int64_t x, y; }; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = int64_t; + using packed_type = int64_2; + static __device__ inline constexpr int64_2 pack(const int64_t &i) { return int64_2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = int64_t; + using packed_type = int64_2; + static __device__ inline constexpr int64_2 pack(const int64_t &i) { return int64_2{i, i}; } +}; +struct uint64_2 { uint64_t x, y; }; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = uint64_t; + using packed_type = uint64_2; + static __device__ inline constexpr uint64_2 pack(const uint64_t &i) { return uint64_2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = uint64_t; + using packed_type = uint64_2; + static __device__ inline constexpr uint64_2 pack(const uint64_t &i) { return uint64_2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 4; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 4; } +}; +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = fp8e4m3; + using packed_type = fp8e4m3_4; +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 4; } + using unpacked_type = fp8e4m3; + using packed_type = fp8e4m3_4; +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = fp8e5m2; + using packed_type = fp8e5m2_4; +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 4; } + using unpacked_type = fp8e5m2; + using packed_type = fp8e5m2_4; +}; +#endif +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = fp8e8m0; + using packed_type = fp8e8m0_4; +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 4; } + using unpacked_type = fp8e8m0; + using packed_type = fp8e8m0_4; +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = fp4e2m1_2; + using packed_type = fp4e2m1_4; +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 4; } + using unpacked_type = fp4e2m1_2; + using packed_type = fp4e2m1_4; +}; +#endif + + + + + + + +template struct convertor { + + + + + + + static __host__ __device__ inline T convert(const U & u) { + return (T)u; + } +}; +template<> struct convertor { + static __host__ __device__ inline float convert(const bf16 & u) { + return __bfloat162float(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline bf16 convert(const float & u) { + return __float2bfloat16_rn(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float2 convert(const bf16_2 & u) { + return __bfloat1622float2(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline bf16_2 convert(const float2 & u) { + return __float22bfloat162_rn(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float convert(const half & u) { + return __half2float(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline half convert(const float & u) { + return __float2half(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float2 convert(const half_2 & u) { + return __half22float2(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline half_2 convert(const float2 & u) { + return __float22half2_rn(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline bf16 convert(const half & u) { + return __float2bfloat16_rn(__half2float(u)); + } +}; +template<> struct convertor { + static __host__ __device__ inline half convert(const bf16 & u) { + return __float2half(__bfloat162float(u)); + } +}; +template<> struct convertor { + static __host__ __device__ inline bf16_2 convert(const half_2 & u) { + return __float22bfloat162_rn(__half22float2(u)); + } +}; +template<> struct convertor { + static __host__ __device__ inline half_2 convert(const bf16_2 & u) { + return __float22half2_rn(__bfloat1622float2(u)); + } +}; +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + +template<> struct convertor { + static __host__ __device__ inline fp8e8m0_4 convert(const float4& u) { + return __nv_fp8x4_e8m0(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float4 convert(const fp8e8m0_4& u) { + __nv_fp8_e8m0 *vals = reinterpret_cast<__nv_fp8_e8m0*>(const_cast<__nv_fp8x4_e8m0*>(&u)); + return make_float4(float(vals[0]), float(vals[1]), float(vals[2]), float(vals[3])); + } +}; +template<> struct convertor { + static __host__ __device__ inline fp8e8m0_2 convert(const float2& u) { + return __nv_fp8x2_e8m0(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float2 convert(const fp8e8m0_2& u) { + __nv_fp8_e8m0 *vals = reinterpret_cast<__nv_fp8_e8m0*>(const_cast<__nv_fp8x2_e8m0*>(&u)); + return make_float2(float(vals[0]), float(vals[1])); + } +}; +template<> struct convertor { + static __host__ __device__ inline fp8e8m0 convert(const float & u) { + return __nv_fp8_e8m0(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float convert(const fp8e8m0 & u) { + return float(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline bf16_2 convert(const fp8e8m0_4 & u) { + float4 f4 = convertor::convert(u); + float2 f2 = make_float2(f4.x, f4.y); + return __float22bfloat162_rn(f2); + } +}; +template<> struct convertor { + static __host__ __device__ inline fp8e8m0_4 convert(const bf16_2 & u) { + float2 f2 = __bfloat1622float2(u); + float4 f4 = make_float4(f2.x, f2.y, 0.0f, 0.0f); + return __nv_fp8x4_e8m0(f4); + } +}; + +template<> struct convertor { + static __device__ inline fp4e2m1 convert(const float & u) { + __nv_fp4_storage_t storage = __nv_cvt_float_to_fp4(u, __NV_E2M1, cudaRoundNearest); + return std::bit_cast(storage); + } +}; +template<> struct convertor { + static __device__ inline float convert(const fp4e2m1 & u) { + __nv_fp4_storage_t storage = std::bit_cast<__nv_fp4_storage_t>(u); + __half_raw hr = __nv_cvt_fp4_to_halfraw(storage, __NV_E2M1); + return __half2float(*reinterpret_cast<__half*>(&hr)); + } +}; +template<> struct convertor { + static __device__ inline fp4e2m1_2 convert(const float2 & u) { + __nv_fp4x2_storage_t storage = __nv_cvt_float2_to_fp4x2(u, __NV_E2M1, cudaRoundNearest); + return std::bit_cast(storage); + } +}; +template<> struct convertor { + static __device__ inline float2 convert(const fp4e2m1_2 & u) { + __nv_fp4x2_storage_t storage = std::bit_cast<__nv_fp4x2_storage_t>(u); + __half2_raw hr2 = __nv_cvt_fp4x2_to_halfraw2(storage, __NV_E2M1); + return __half22float2(*reinterpret_cast<__half2*>(&hr2)); + } +}; +#endif +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + +template<> struct convertor { + static __host__ __device__ inline fp8e4m3_4 convert(const float4& u) { + return __nv_fp8x4_e4m3(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float4 convert(const fp8e4m3_4& u) { + __nv_fp8_e4m3 *vals = reinterpret_cast<__nv_fp8_e4m3*>(const_cast<__nv_fp8x4_e4m3*>(&u)); + return make_float4(float(vals[0]), float(vals[1]), float(vals[2]), float(vals[3])); + } +}; +template<> struct convertor { + static __host__ __device__ inline fp8e4m3_2 convert(const float2& u) { + return __nv_fp8x2_e4m3(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float2 convert(const fp8e4m3_2& u) { + __nv_fp8_e4m3 *vals = reinterpret_cast<__nv_fp8_e4m3*>(const_cast<__nv_fp8x2_e4m3*>(&u)); + return make_float2(float(vals[0]), float(vals[1])); + } +}; +template<> struct convertor { + static __host__ __device__ inline fp8e4m3 convert(const float & u) { + return __nv_fp8_e4m3(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float convert(const fp8e4m3 & u) { + return float(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline bf16_2 convert(const fp8e4m3_4 & u) { + float4 f4 = convertor::convert(u); + float2 f2 = make_float2(f4.x, f4.y); + return __float22bfloat162_rn(f2); + } +}; +template<> struct convertor { + static __host__ __device__ inline fp8e4m3_4 convert(const bf16_2 & u) { + float2 f2 = __bfloat1622float2(u); + float4 f4 = make_float4(f2.x, f2.y, 0.0f, 0.0f); + return __nv_fp8x4_e4m3(f4); + } +}; + +template<> struct convertor { + static __host__ __device__ inline fp8e5m2_4 convert(const float4& u) { + return __nv_fp8x4_e5m2(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float4 convert(const fp8e5m2_4& u) { + __nv_fp8_e5m2 *vals = reinterpret_cast<__nv_fp8_e5m2*>(const_cast<__nv_fp8x4_e5m2*>(&u)); + return make_float4(float(vals[0]), float(vals[1]), float(vals[2]), float(vals[3])); + } +}; +template<> struct convertor { + static __host__ __device__ inline fp8e5m2_2 convert(const float2& u) { + return __nv_fp8x2_e5m2(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float2 convert(const fp8e5m2_2& u) { + __nv_fp8_e5m2 *vals = reinterpret_cast<__nv_fp8_e5m2*>(const_cast<__nv_fp8x2_e5m2*>(&u)); + return make_float2(float(vals[0]), float(vals[1])); + } +}; +template<> struct convertor { + static __host__ __device__ inline fp8e5m2 convert(const float & u) { + return __nv_fp8_e5m2(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float convert(const fp8e5m2 & u) { + return float(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline bf16_2 convert(const fp8e5m2_4 & u) { + float4 f4 = convertor::convert(u); + float2 f2 = make_float2(f4.x, f4.y); + return __float22bfloat162_rn(f2); + } +}; +template<> struct convertor { + static __host__ __device__ inline fp8e5m2_4 convert(const bf16_2 & u) { + float2 f2 = __bfloat1622float2(u); + float4 f4 = make_float4(f2.x, f2.y, 0.0f, 0.0f); + return __nv_fp8x4_e5m2(f4); + } +}; +#endif +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/common/common.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/common/common.cuh new file mode 100644 index 0000000..855bb69 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/common/common.cuh @@ -0,0 +1,15 @@ + + + + + + +#pragma once + +#include "base_types.cuh" +#include "base_ops.cuh" +#include "util.cuh" + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +#include "multimem.cuh" +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/common/multimem.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/common/multimem.cuh new file mode 100644 index 0000000..ac5f226 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/common/multimem.cuh @@ -0,0 +1,407 @@ + + + + + +#pragma once + +#include "base_types.cuh" + +namespace kittens { + +enum class reduce_op { + ADD = 0, + MIN = 1, + MAX = 2 +}; + +enum class memory_model { + WEAK = 0, + STRONG = 1 +}; + +template +struct multimem; + +template <> +struct multimem { + template + __device__ static inline void ld_reduce(int &dst, const int *src) { + if constexpr (Op == reduce_op::ADD) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.add.s32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.add.s32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MIN) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.min.s32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.min.s32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MAX) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.max.s32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.max.s32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } + } + } + template + __device__ static inline void st(int *dst, const int &src) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.st.weak.global.s32 [%0], %1;" + :: "l"(dst), "r"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.st.release.sys.global.s32 [%0], %1;" + :: "l"(dst), "r"(src) : "memory"); + } + } + template + __device__ static inline void red(int *dst, const int &src) { + if constexpr (Op == reduce_op::ADD) { + asm volatile("multimem.red.release.sys.global.add.s32 [%0], %1;" + : : "l"(dst), "r"(src) : "memory"); + } else if constexpr (Op == reduce_op::MIN) { + asm volatile("multimem.red.release.sys.global.min.s32 [%0], %1;" + : : "l"(dst), "r"(src) : "memory"); + } else if constexpr (Op == reduce_op::MAX) { + asm volatile("multimem.red.release.sys.global.max.s32 [%0], %1;" + : : "l"(dst), "r"(src) : "memory"); + } + } +}; + +template <> +struct multimem { + template + __device__ static inline void ld_reduce(uint &dst, const uint *src) { + if constexpr (Op == reduce_op::ADD) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.add.u32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.add.u32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MIN) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.min.u32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.min.u32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MAX) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.max.u32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.max.u32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } + } + } + template + __device__ static inline void st(uint *dst, const uint &src) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.st.weak.global.u32 [%0], %1;" + :: "l"(dst), "r"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.st.release.sys.global.u32 [%0], %1;" + :: "l"(dst), "r"(src) : "memory"); + } + } + template + __device__ static inline void red(uint *dst, const uint &src) { + if constexpr (Op == reduce_op::ADD) { + asm volatile("multimem.red.release.sys.global.add.u32 [%0], %1;" + : : "l"(dst), "r"(src) : "memory"); + } else if constexpr (Op == reduce_op::MIN) { + asm volatile("multimem.red.release.sys.global.min.u32 [%0], %1;" + : : "l"(dst), "r"(src) : "memory"); + } else if constexpr (Op == reduce_op::MAX) { + asm volatile("multimem.red.release.sys.global.max.u32 [%0], %1;" + : : "l"(dst), "r"(src) : "memory"); + } + } +}; + +template <> +struct multimem { + template + __device__ static inline void ld_reduce(float &dst, const float *src) { + static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for f32 ld_reduce operations"); + if constexpr (Op == reduce_op::ADD) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.add.f32 %0, [%1];" + : "=f"(dst) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.add.f32 %0, [%1];" + : "=f"(dst) : "l"(src) : "memory"); + } + } + } + template + __device__ static inline void st(float *dst, const float &src) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.st.weak.global.f32 [%0], %1;" + :: "l"(dst), "f"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.st.release.sys.global.f32 [%0], %1;" + :: "l"(dst), "f"(src) : "memory"); + } + } + template + __device__ static inline void red(float *dst, const float &src) { + static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for f32 red operations"); + if constexpr (Op == reduce_op::ADD) { + asm volatile("multimem.red.release.sys.global.add.f32 [%0], %1;" + : : "l"(dst), "f"(src) : "memory"); + } + } +}; + + +template <> +struct multimem { + template + __device__ static inline void ld_reduce(float2 &dst, const float2 *src) { + static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for f32 ld_reduce operations"); + if constexpr (Op == reduce_op::ADD) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.add.v2.f32 {%0, %1}, [%2];" + : "=f"(dst.x), "=f"(dst.y) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.add.v2.f32 {%0, %1}, [%2];" + : "=f"(dst.x), "=f"(dst.y) : "l"(src) : "memory"); + } + } + } + template + __device__ static inline void st(float2 *dst, const float2 &src) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.st.weak.global.v2.f32 [%0], {%1, %2};" + :: "l"(dst), "f"(src.x), "f"(src.y) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.st.release.sys.global.v2.f32 [%0], {%1, %2};" + :: "l"(dst), "f"(src.x), "f"(src.y) : "memory"); + } + } + template + __device__ static inline void red(float2 *dst, const float2 &src) { + static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for f32 red operations"); + if constexpr (Op == reduce_op::ADD) { + asm volatile("multimem.red.release.sys.global.add.v2.f32 [%0], {%1, %2};" + : : "l"(dst), "f"(src.x), "f"(src.y) : "memory"); + } + } +}; + +template <> +struct multimem { + template + __device__ static inline void ld_reduce(bf16 &dst, const bf16 *src) { + if constexpr (Op == reduce_op::ADD) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.add.acc::f32.bf16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.add.acc::f32.bf16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MIN) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.min.bf16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.min.bf16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MAX) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.max.bf16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.max.bf16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } + } + template + __device__ static inline void st(bf16 *dst, const bf16 &src) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.st.weak.global.bf16 [%0], %1;" + :: "l"(dst), "h"(*reinterpret_cast(&src)) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.st.release.sys.global.bf16 [%0], %1;" + :: "l"(dst), "h"(*reinterpret_cast(&src)) : "memory"); + } + } + template + __device__ static inline void red(bf16 *dst, const bf16 &src) { + static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for bf16 red operations"); + if constexpr (Op == reduce_op::ADD) { + asm volatile("multimem.red.release.sys.global.add.bf16 [%0], %1;" + : : "l"(dst), "h"(*reinterpret_cast(&src)) : "memory"); + } + } +}; + +template <> +struct multimem { + template + __device__ static inline void ld_reduce(bf16_2 &dst, const bf16_2 *src) { + if constexpr (Op == reduce_op::ADD) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.add.acc::f32.bf16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.add.acc::f32.bf16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MIN) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.min.bf16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.min.bf16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MAX) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.max.bf16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.max.bf16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } + } + template + __device__ static inline void st(bf16_2 *dst, const bf16_2 &src) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.st.weak.global.bf16x2 [%0], %1;" + :: "l"(dst), "r"(*reinterpret_cast(&src)) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.st.release.sys.global.bf16x2 [%0], %1;" + :: "l"(dst), "r"(*reinterpret_cast(&src)) : "memory"); + } + } + template + __device__ static inline void red(bf16_2 *dst, const bf16_2 &src) { + static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for bf16_2 red operations"); + if constexpr (Op == reduce_op::ADD) { + asm volatile("multimem.red.release.sys.global.add.bf16x2 [%0], %1;" + : : "l"(dst), "r"(*reinterpret_cast(&src)) : "memory"); + } + } +}; + +template <> +struct multimem { + template + __device__ static inline void ld_reduce(half &dst, const half *src) { + if constexpr (Op == reduce_op::ADD) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.add.acc::f32.f16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.add.acc::f32.f16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MIN) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.min.f16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.min.f16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MAX) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.max.f16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.max.f16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } + } + template + __device__ static inline void st(half *dst, const half &src) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.st.weak.global.f16 [%0], %1;" + :: "l"(dst), "h"(*reinterpret_cast(&src)) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.st.release.sys.global.f16 [%0], %1;" + :: "l"(dst), "h"(*reinterpret_cast(&src)) : "memory"); + } + } + template + __device__ static inline void red(half *dst, const half &src) { + static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for f16 red operations"); + if constexpr (Op == reduce_op::ADD) { + asm volatile("multimem.red.release.sys.global.add.f16 [%0], %1;" + : : "l"(dst), "h"(*reinterpret_cast(&src)) : "memory"); + } + } +}; + +template <> +struct multimem { + template + __device__ static inline void ld_reduce(half_2 &dst, const half_2 *src) { + if constexpr (Op == reduce_op::ADD) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.add.acc::f32.f16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.add.acc::f32.f16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MIN) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.min.f16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.min.f16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MAX) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.max.f16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.max.f16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } + } + template + __device__ static inline void st(half_2 *dst, const half_2 &src) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.st.weak.global.f16x2 [%0], %1;" + :: "l"(dst), "r"(*reinterpret_cast(&src)) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.st.release.sys.global.f16x2 [%0], %1;" + :: "l"(dst), "r"(*reinterpret_cast(&src)) : "memory"); + } + } + template + __device__ static inline void red(half_2 *dst, const half_2 &src) { + static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for f16_2 red operations"); + if constexpr (Op == reduce_op::ADD) { + asm volatile("multimem.red.release.sys.global.add.f16x2 [%0], %1;" + : : "l"(dst), "r"(*reinterpret_cast(&src)) : "memory"); + } + } +}; + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/common/util.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/common/util.cuh new file mode 100644 index 0000000..ac666b3 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/common/util.cuh @@ -0,0 +1,551 @@ + + + + + +#pragma once + + + + + + +namespace kittens { + + + + + + +constexpr int BASE_TILE_DIM = 16; +template constexpr int TILE_COL_DIM = sizeof(T) == 1 ? BASE_TILE_DIM * 2 : BASE_TILE_DIM; +template constexpr int TILE_ROW_DIM = BASE_TILE_DIM; + + + +template constexpr int TILE_ELEMENTS{TILE_COL_DIM*TILE_ROW_DIM}; + + + +constexpr int WARP_THREADS{32}; + + + +constexpr int WARPGROUP_THREADS{128}; + + + + +constexpr int WARPGROUP_WARPS{4}; + + + + + +__device__ __forceinline__ int warpid() { return threadIdx.x >> 5; } + + + + +__device__ __forceinline__ int warpgroupid() { return threadIdx.x >> 7; } + + + + +__device__ __forceinline__ int laneid() { return threadIdx.x & 0x1f; } + + + + +__device__ __forceinline__ int smid() { + uint32_t r; + asm volatile("mov.u32 %0, %%smid;" : "=r"(r)); + return static_cast(r); +} + +#if defined(KITTENS_SM80) +constexpr int MAX_SHARED_MEMORY = 163 * 1024; +#elif defined(KITTENS_SM90) || defined(KITTENS_SM10X) +constexpr int MAX_SHARED_MEMORY = 227 * 1024; +#elif defined(KITTENS_SM120) +constexpr int MAX_SHARED_MEMORY = 99 * 1024; +#endif + +struct transpose { + static constexpr int N = 0; + static constexpr int T = 1; +}; +struct axis { + static constexpr int ROW = 0; + static constexpr int COL = 1; +}; + + + + + + + + + + + +namespace ducks { + + + + +struct default_type {}; + + +#define typeof(A) typename std::remove_const::type>::type + +} + + + + + + +static constexpr uint32_t MASK_ALL = 0xFFFFFFFF; + + + + + + + + + +template +__device__ static inline T packed_shfl_down_sync(uint32_t mask, const T &f, int delta) { + return __shfl_down_sync(mask, f, delta); +} +template<> +__device__ inline float2 packed_shfl_down_sync(uint32_t mask, const float2 &f, int delta) { + float2 r; + r.x = __shfl_down_sync(mask, f.x, delta); + r.y = __shfl_down_sync(mask, f.y, delta); + return r; +} + + + + + + + + +template +__device__ static inline T packed_shfl_xor_sync(uint32_t mask, const T &f, int lane_mask) { + return __shfl_xor_sync(mask, f, lane_mask); +} +template<> +__device__ inline float2 packed_shfl_xor_sync(uint32_t mask, const float2 &f, int lane_mask) { + float2 r; + r.x = __shfl_xor_sync(mask, f.x, lane_mask); + r.y = __shfl_xor_sync(mask, f.y, lane_mask); + return r; +} + + + + + + + + +template +__device__ static inline T packed_shfl_sync(uint32_t mask, const T &f, int src) { + return __shfl_sync(mask, f, src); +} +template<> +__device__ inline float2 packed_shfl_sync(uint32_t mask, const float2 &f, int src) { + float2 r; + r.x = __shfl_sync(mask, f.x, src); + r.y = __shfl_sync(mask, f.y, src); + return r; +} + + + + + + + + + + + + + + + + + + + + + + + + +#if defined(__CUDACC__) +#define KITTENS_ALIGN_AS(n) __align__(n) +#else +#define KITTENS_ALIGN_AS(n) alignas(n) +#endif + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +#define KITTENS_DEFAULT_ALIGN KITTENS_ALIGN_AS(128) +#else +#define KITTENS_DEFAULT_ALIGN KITTENS_ALIGN_AS(16) +#endif + + + + + + + + +template +__host__ __device__ constexpr T cdiv(T a, T b) { + return (a + b - 1) / b; +} + + + + +struct KITTENS_DEFAULT_ALIGN alignment_dummy { int dummy; }; + + + + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template +#else +template +#endif +struct shared_allocator { + int *ptr; + +private: + + template + struct variadic_array; + template + struct variadic_array { + using type = typename variadic_array::type[first_dim]; + }; + template + struct variadic_array { + using type = A; + }; + template + using variadic_array_t = typename variadic_array::type; + + template + __device__ inline void align_ptr() { + if constexpr (alignment > 0) { + uint64_t p = reinterpret_cast(ptr); + if(p % alignment != 0) { + ptr = (int*)(p + (alignment-(p%alignment))); + } + } + } + +public: + + + + + __device__ shared_allocator(int *_ptr): ptr(_ptr) {} + + + + + + + template + __device__ inline variadic_array_t& allocate() { + + align_ptr(); + using at = variadic_array_t; + at*p = reinterpret_cast(ptr); + ptr += sizeof(at)/sizeof(int); + return *p; + } + + + + + + + + template + __device__ inline variadic_array_t& allocate() { + + align_ptr(); + using at = variadic_array_t; + at*p = reinterpret_cast(ptr); + ptr += sizeof(at)/sizeof(int); + return *p; + } +}; + +#if (defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120)) + + + +using tma_allocator = shared_allocator<1024>; +using tma_swizzle_allocator = tma_allocator; + + + + +__device__ static inline int3 clusterIdx() { + int3 cluster_idx; + asm volatile("mov.u32 %0, %clusterid.x;\n" : "=r"(cluster_idx.x)); + asm volatile("mov.u32 %0, %clusterid.y;\n" : "=r"(cluster_idx.y)); + asm volatile("mov.u32 %0, %clusterid.z;\n" : "=r"(cluster_idx.z)); + return cluster_idx; +} + + +__device__ static inline int3 clusterid() { + int3 clusterid; + asm volatile("mov.u32 %0, %clusterid.x;\n" : "=r"(clusterid.x)); + asm volatile("mov.u32 %0, %clusterid.y;\n" : "=r"(clusterid.y)); + asm volatile("mov.u32 %0, %clusterid.z;\n" : "=r"(clusterid.z)); + return clusterid; +} + + +__device__ static inline int3 nclusterid() { + int3 nclusterid; + asm volatile("mov.u32 %0, %nclusterid.x;\n" : "=r"(nclusterid.x)); + asm volatile("mov.u32 %0, %nclusterid.y;\n" : "=r"(nclusterid.y)); + asm volatile("mov.u32 %0, %nclusterid.z;\n" : "=r"(nclusterid.z)); + return nclusterid; +} + + +__device__ static inline int cluster_ctarank() { + uint32_t ctarank; + asm volatile("mov.u32 %0, %cluster_ctarank;\n" : "=r"(ctarank)); + return ctarank; +} + + +__device__ static inline int cluster_nctarank() { + uint32_t nctarank; + asm volatile("mov.u32 %0, %cluster_nctarank;\n" : "=r"(nctarank)); + return nctarank; +} + +#endif + + + + + + + + + + + + + + + + + + +template +__device__ static inline int2 get_swizzled_2d_idx(const int num_rows, const int num_cols, const int linear_idx) { + static_assert(SUPERGROUP_SIZE > 0, "SUPERGROUP_SIZE must be greater than 0"); + if constexpr (ROW_MAJOR) { + const int supergroup_numel = num_rows*SUPERGROUP_SIZE; + const int supergroup_idx = linear_idx/supergroup_numel; + const int supersection_cols = (num_cols/SUPERGROUP_SIZE)*SUPERGROUP_SIZE; + const int supersection_numel = num_rows*supersection_cols; + const int finalsection_cols = num_cols-supersection_cols; + int row_idx, col_idx; + if (linear_idx < supersection_numel) { + row_idx = (linear_idx%supergroup_numel)/SUPERGROUP_SIZE; + col_idx = supergroup_idx*SUPERGROUP_SIZE + linear_idx%SUPERGROUP_SIZE; + } else { + const int remainder_task_id = linear_idx - supersection_numel; + row_idx = remainder_task_id/finalsection_cols; + col_idx = supersection_cols + remainder_task_id%finalsection_cols; + } + return { (supergroup_idx&1) ? num_rows-row_idx-1 : row_idx, col_idx }; + } else { + const int supergroup_numel = num_cols*SUPERGROUP_SIZE; + const int supergroup_idx = linear_idx/supergroup_numel; + const int supersection_rows = (num_rows/SUPERGROUP_SIZE)*SUPERGROUP_SIZE; + const int supersection_numel = num_cols*supersection_rows; + const int finalsection_rows = num_rows-supersection_rows; + int row_idx, col_idx; + if (linear_idx < supersection_numel) { + row_idx = supergroup_idx*SUPERGROUP_SIZE + linear_idx%SUPERGROUP_SIZE; + col_idx = (linear_idx%supergroup_numel)/SUPERGROUP_SIZE; + } else { + const int remainder_task_id = linear_idx - supersection_numel; + row_idx = supersection_rows + remainder_task_id%finalsection_rows; + col_idx = remainder_task_id/finalsection_rows; + } + return { row_idx, (supergroup_idx&1) ? num_cols-col_idx-1 : col_idx }; + } +} + +template +__device__ static inline int get_phasebit(uint32_t bitfield, int ring_id) { + if constexpr (half == 0) + return (bitfield >> (ring_id)) & 0b1; + else if constexpr (half == 1) + return (bitfield >> (ring_id + 16)) & 0b1; + else + asm volatile ("brkpt;\n"); + return -1; +} + +template +__device__ static inline void update_phasebit(uint32_t &bitfield, int ring_id) { + if constexpr (half == 0) + bitfield ^= (1 << (ring_id)); + else if constexpr (half == 1) + bitfield ^= (1 << (ring_id + 16)); + else + asm volatile ("brkpt;\n"); +} + +template __device__ static inline int ring_advance(int ring, int distance=1) { return (ring + distance) % N; } +template __device__ static inline int ring_retreat(int ring, int distance=1) { return (ring + 16*N - distance) % N; } + + + +#ifndef KITTENS_NO_HOST + + +#define CUCHECK(cmd) do { \ + CUresult err = cmd; \ + if (err != CUDA_SUCCESS) { \ + const char *errStr; \ + cuGetErrorString(err, &errStr); \ + fprintf(stderr, "Failed: CUDA error %s:%d '%s'\n", \ + __FILE__, __LINE__, errStr); \ + exit(EXIT_FAILURE); \ + } \ +} while(0) + + +#define CUDACHECK(cmd) do { \ + cudaError_t err = cmd; \ + if (err != cudaSuccess) { \ + fprintf(stderr, "Failed: CUDA error %s:%d '%s'\n", \ + __FILE__, __LINE__, cudaGetErrorString(err)); \ + exit(EXIT_FAILURE); \ + } \ +} while(0) + + +#define CHECK_CUDA_ERROR(val) check((val), #val, __FILE__, __LINE__) +template __host__ void check( + T err, char const* const func, char const* const file, int const line +) { + if (err != cudaSuccess) { + std::cerr << "CUDA Runtime Error at: " << file << ":" << line + << std::endl; + std::cerr << cudaGetErrorString(err) << " " << func << std::endl; + std::exit(EXIT_FAILURE); + } +} + + + + + + + __host__ inline int num_sms(int device_id = -1) { + if (device_id < 0) CUDACHECK(cudaGetDevice(&device_id)); + int sm_count; + CUDACHECK(cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device_id)); + return sm_count; +} + +template +struct LaunchConfig { + static constexpr unsigned int num_attributes = CLUSTER ? (PDL ? 3 : 2) : 1; + + cudaLaunchAttribute attributes[num_attributes]; + cudaLaunchConfig_t config = {0}; + + __host__ inline LaunchConfig(dim3 grid, dim3 block, size_t dynamic_shared_memory, + cudaStream_t stream) noexcept requires(!CLUSTER) { + if constexpr (PDL) { + attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attributes[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attributes; + config.numAttrs = num_attributes; + } else { + config.numAttrs = 0; + } + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes= dynamic_shared_memory; + config.stream = stream; + } + + __host__ inline LaunchConfig(dim3 grid, dim3 block, size_t dynamic_shared_memory, + cudaStream_t stream, dim3 cluster_preferred, dim3 cluster_minimum) noexcept requires(CLUSTER) { + attributes[0].id = cudaLaunchAttributePreferredClusterDimension; + attributes[0].val.preferredClusterDim.x = cluster_preferred.x; + attributes[0].val.preferredClusterDim.y = cluster_preferred.y; + attributes[0].val.preferredClusterDim.z = cluster_preferred.z; + attributes[1].id = cudaLaunchAttributeClusterDimension; + attributes[1].val.clusterDim.x = cluster_minimum.x; + attributes[1].val.clusterDim.y = cluster_minimum.y; + attributes[1].val.clusterDim.z = cluster_minimum.z; + if constexpr (PDL) { + attributes[2].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attributes[2].val.programmaticStreamSerializationAllowed = 1; + } + config.attrs = attributes; + config.numAttrs = num_attributes; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes= dynamic_shared_memory; + config.stream = stream; + } + + __host__ inline LaunchConfig(dim3 grid, dim3 block, size_t dynamic_shared_memory, + cudaStream_t stream, dim3 cluster) noexcept requires(CLUSTER) + : LaunchConfig(grid, block, dynamic_shared_memory, stream, cluster, cluster) { } + + __host__ inline LaunchConfig(const LaunchConfig& other) noexcept : config(other.config) { + std::copy_n(other.attributes, num_attributes, attributes); + config.attrs = attributes; + } + + __host__ inline LaunchConfig& operator=(const LaunchConfig& other) noexcept { + if (this != &other) { + std::copy_n(other.attributes, num_attributes, attributes); + config = other.config; + config.attrs = attributes; + } + return *this; + } + + __host__ LaunchConfig(LaunchConfig&&) = delete; + __host__ LaunchConfig& operator=(LaunchConfig&&) = delete; + + __host__ inline operator cudaLaunchConfig_t*() noexcept { return &config; } + __host__ inline operator const cudaLaunchConfig_t*() const noexcept { return &config; } +}; + +#endif + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/kittens.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/kittens.cuh new file mode 100644 index 0000000..6578fff --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/kittens.cuh @@ -0,0 +1,73 @@ + + + + + +#pragma once + + +#if defined(KITTENS_AMPERE) +#define KITTENS_SM80 +#endif +#if defined(KITTENS_HOPPER) +#define KITTENS_SM90 +#endif +#if defined(KITTENS_BLACKWELL) +#define KITTENS_SM100 +#endif + + +#if defined(KITTENS_SM80) + defined(KITTENS_SM90) + defined(KITTENS_SM100) + defined(KITTENS_SM103) + defined(KITTENS_SM120) != 1 +#error "Define exactly one of: KITTENS_SM80, KITTENS_SM90, KITTENS_SM100, KITTENS_SM103, KITTENS_SM120" +#endif + + +#if defined(KITTENS_SM100) || defined(KITTENS_SM103) +#define KITTENS_SM10X +#endif + + +#include +#include +#include +#include +#include +#ifdef KITTENS_NO_HOST +namespace std { using namespace cuda::std; } +using uint = uint32_t; +#if defined(__CUDACC_VER_MAJOR__) && __CUDACC_VER_MAJOR__ >= 13 +struct alignas(128) CUtensorMap { char __opaque[128]; }; +#else +struct alignas(64) CUtensorMap { char __opaque[128]; }; +#endif +#endif + + +#include +#include +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +#include +#endif +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) +#include +#endif + + +#ifndef KITTENS_NO_HOST +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#endif + + +#include "common/common.cuh" +#include "types/types.cuh" +#include "ops/ops.cuh" diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/group.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/group.cuh new file mode 100644 index 0000000..bb48e8d --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/group.cuh @@ -0,0 +1,117 @@ + + + + + +#pragma once + +#include "../../common/common.cuh" +#include "../../types/types.cuh" +#include "../thread/thread.cuh" + +#define KITTENS_CHECK_WARP static_assert(GROUP_WARPS==1, "Warp (GROUP_WARPS=1) function called from a non-warp group."); + +#define KITTENS_CHECK_WARPGROUP static_assert(GROUP_WARPS==4, "Warpgroup (GROUP_WARPS=4) function called from a non-warpgroup group."); + + +#if defined(KITTENS_SM90) +#include "mma/base/base.cuh" +#endif + +namespace kittens { + + + +template +struct group { +static constexpr int GROUP_WARPS = _GROUP_WARPS; +static constexpr int GROUP_THREADS = GROUP_WARPS * kittens::WARP_THREADS; +__device__ static inline int laneid() { return threadIdx.x % GROUP_THREADS; } +__device__ static inline int warpid() { return laneid() / kittens::WARP_THREADS; } +__device__ static inline int groupid() { return threadIdx.x / GROUP_THREADS; } + +__device__ static inline void sync(int id) { + asm volatile("bar.sync %0, %1;\n" :: "r"(id), "n"(GROUP_THREADS)); +} +template __device__ static inline void sync() { + static_assert(GROUP_WARPS==1, "barrier-less sync() can only be called by a single warp!"); + asm volatile("bar.warp.sync %0;\n" :: "n"(MASK)); +} +__device__ static inline void arrive(int id) { + asm volatile("bar.arrive %0, %1;\n" :: "r"(id), "n"(GROUP_THREADS)); +} + +#include "memory/memory.cuh" +#include "shared/shared.cuh" +#include "register/register.cuh" +#include "mma/mma.cuh" +#include "util/util.cuh" + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + +template __device__ static inline void increase_registers() { + static_assert(n_reg % 8 == 0, "n_reg must be a multiple of 8"); + asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" :: "n"(n_reg)); +} +template __device__ static inline void decrease_registers() { + static_assert(n_reg % 8 == 0, "n_reg must be a multiple of 8"); + asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" :: "n"(n_reg)); +} +__device__ static inline void producer_registers() { decrease_registers<24>(); } +template __device__ static inline void consumer_registers() { increase_registers<480/NCWG - 8*(NCWG>3) - 224*(NCWG==1)>(); } + + + + + + +struct tma { +#include "memory/tile/tma.cuh" +#include "memory/vec/tma.cuh" +#include "util/tma.cuh" +struct cluster { +#include "memory/tile/tma_cluster.cuh" +#include "memory/vec/tma_cluster.cuh" +#include "util/tma_cluster.cuh" +}; +}; + +#endif + +}; + +namespace everyone { + + +__device__ static inline void sync(int id) { + asm volatile("bar.sync %0;\n" :: "r"(id)); +} + + +namespace tma { +namespace cluster { +__device__ static inline void arrive() { + asm volatile ("barrier.cluster.arrive.release;\n"); +} +__device__ static inline void arrive_aligned() { + asm volatile ("barrier.cluster.arrive.release.aligned;\n"); +} +__device__ static inline void wait() { + asm volatile ("barrier.cluster.wait.acquire;\n"); +} +__device__ static inline void wait_aligned() { + asm volatile ("barrier.cluster.wait.acquire.aligned;\n"); +} +__device__ static inline void sync() { + arrive_aligned(); + wait_aligned(); +} +} +} + +}; + +using warp = group<1>; +using warpgroup = group<4>; + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/memory.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/memory.cuh new file mode 100644 index 0000000..404f8af --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/memory.cuh @@ -0,0 +1,7 @@ + + + + + +#include "tile/tile.cuh" +#include "vec/vec.cuh" diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/complex_global_to_register.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/complex_global_to_register.cuh new file mode 100644 index 0000000..989b677 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/complex_global_to_register.cuh @@ -0,0 +1,42 @@ + + + + + + + + + + + + + + +template>> +__device__ inline static void load(CRT &dst, const CGL &src, const COORD &idx) { + load(dst.real, src.real, idx); + load(dst.imag, src.imag, idx); +} +template>> +__device__ inline static void load(CRT &dst, const CGL &src, const COORD &idx) { + load<2, CRT, CGL>(dst, src, idx); +} + + + + + + + + + + +template>> +__device__ inline static void store(CGL &dst, const CRT &src, const COORD &idx) { + store(dst.real, src.real, idx); + store(dst.imag, src.imag, idx); +} +template>> +__device__ inline static void store(CGL &dst, const CRT &src, const COORD &idx) { + store<2, CRT, CGL>(dst, src, idx); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/complex_global_to_shared.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/complex_global_to_shared.cuh new file mode 100644 index 0000000..28c2565 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/complex_global_to_shared.cuh @@ -0,0 +1,69 @@ + + + + + + + + + + + + + + + +template> +__device__ static inline void load(CST &dst, const CGL &src, const COORD &idx) { + load(dst.real, src.real, idx); + load(dst.imag, src.imag, idx); +} +template> +__device__ static inline void load(CST &dst, const CGL &src, const COORD &idx) { + load<2, false, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx); + load<2, false, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx); +} + + + + + + + + + + + +template> +__device__ static inline void store(CGL &dst, const CST &src, const COORD &idx) { + store(dst.real, src.real, idx); + store(dst.imag, src.imag, idx); +} +template> +__device__ static inline void store(CGL &dst, const CST &src, const COORD &idx) { + store<2, false, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx); + store<2, false, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx); +} + + + + + + + + + + + + + +template> +__device__ static inline void load_async(CST &dst, const CGL &src, const COORD &idx) { + load_async(dst.real, src.real, idx); + load_async(dst.imag, src.imag, idx); +} +template> +__device__ static inline void load_async(CST &dst, const CGL &src, const COORD &idx) { + load_async<2, false, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx); + load_async<2, false, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/complex_shared_to_register.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/complex_shared_to_register.cuh new file mode 100644 index 0000000..63588be --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/complex_shared_to_register.cuh @@ -0,0 +1,32 @@ + + + + + + + + + + + + + +template +__device__ inline static void load(CRT &dst, const CST &src) { + load(dst.real, src.real); + load(dst.imag, src.imag); +} + + + + + + + + + +template +__device__ inline static void store(CST &dst, const CRT &src) { + store(dst.real, src.real); + store(dst.imag, src.imag); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/global_to_register.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/global_to_register.cuh new file mode 100644 index 0000000..573c31a --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/global_to_register.cuh @@ -0,0 +1,199 @@ + + + + + + + + + + + + + +template>> +__device__ inline static void load(RT &dst, const GL &src, const COORD &idx) { + using T2 = RT::dtype; + using U = typename GL::dtype; + + U *src_ptr = (U*)&src[(idx.template unit_coord())]; + const int row_stride = src.template stride(); + using U2 = base_types::packing::packed_type; + constexpr int pack_size = base_types::packing::num(); + constexpr int half_col = RT::tile_size_col / 2; + int warp_laneid = threadIdx.x % WARP_THREADS; + int local_warpid; + if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4)); + else local_warpid = warpid(); + const int row_offset = RT::rows*local_warpid; + #pragma unroll + for(int i = 0; i < RT::height; i++) { + int row = row_offset + i*RT::tile_size_row + (warp_laneid / 4); + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + pack_size*(warp_laneid % 4); + dst.tiles[i][j].data[0] = base_types::convertor::convert(*(U2*)(&src_ptr[(row+0)*row_stride + (col+0)])); + dst.tiles[i][j].data[2] = base_types::convertor::convert(*(U2*)(&src_ptr[(row+0)*row_stride + (col+half_col)])); + } + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + pack_size*(warp_laneid % 4); + dst.tiles[i][j].data[1] = base_types::convertor::convert(*(U2*)(&src_ptr[(row+8)*row_stride + (col+0)])); + dst.tiles[i][j].data[3] = base_types::convertor::convert(*(U2*)(&src_ptr[(row+8)*row_stride + (col+half_col)])); + } + } +} + + + + + + + + +template>> +__device__ inline static void load(RT &dst, const GL &src, const COORD &idx) { + using T = typename RT::T; + using U = typename GL::dtype; + + #if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert(!std::is_same_v && !std::is_same_v, "Unsupported type for load/store"); + #endif + + U *src_ptr = (U*)&src[(idx.template unit_coord())]; + const int row_stride = src.template stride(); + int warp_laneid = threadIdx.x % WARP_THREADS; + int local_warpid; + if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4)); + else local_warpid = warpid(); + const int row_offset = RT::rows*local_warpid; + #pragma unroll + for(int i = 0; i < RT::height; i++) { + int row = row_offset + i*RT::tile_size_row + 2*(warp_laneid % 4); + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + (warp_laneid / 4); + dst.tiles[i][j].data[0].x = base_types::convertor::convert(src_ptr[(row+0)*row_stride + (col+0)]); + dst.tiles[i][j].data[1].x = base_types::convertor::convert(src_ptr[(row+0)*row_stride + (col+8)]); + } + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + (warp_laneid / 4); + dst.tiles[i][j].data[0].y = base_types::convertor::convert(src_ptr[(row+1)*row_stride + (col+0)]); + dst.tiles[i][j].data[1].y = base_types::convertor::convert(src_ptr[(row+1)*row_stride + (col+8)]); + } + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + (warp_laneid / 4); + dst.tiles[i][j].data[2].x = base_types::convertor::convert(src_ptr[(row+8)*row_stride + (col+0)]); + dst.tiles[i][j].data[3].x = base_types::convertor::convert(src_ptr[(row+8)*row_stride + (col+8)]); + } + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + (warp_laneid / 4); + dst.tiles[i][j].data[2].y = base_types::convertor::convert(src_ptr[(row+9)*row_stride + (col+0)]); + dst.tiles[i][j].data[3].y = base_types::convertor::convert(src_ptr[(row+9)*row_stride + (col+8)]); + } + } +} +template>> +__device__ inline static void load(RT &dst, const GL &src, const COORD &idx) { + load<2>(dst, src, idx); +} + + + + + + + + +template>> +__device__ inline static void store(const GL &dst, const RT &src, const COORD &idx) { + using T2 = RT::dtype; + using U = typename GL::dtype; + + U *dst_ptr = (U*)&dst[(idx.template unit_coord())]; + const int row_stride = dst.template stride(); + using U2 = base_types::packing::packed_type; + constexpr int pack_size = base_types::packing::num(); + constexpr int half_col = RT::tile_size_col / 2; + int warp_laneid = threadIdx.x % WARP_THREADS; + int local_warpid; + if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4)); + else local_warpid = warpid(); + const int row_offset = RT::rows*local_warpid; + #pragma unroll + for(int i = 0; i < RT::height; i++) { + int row = row_offset + i*RT::tile_size_row + (warp_laneid / 4); + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + pack_size*(warp_laneid % 4); + *(U2*)(&dst_ptr[(row+0)*row_stride + (col+0)]) = base_types::convertor::convert(src.tiles[i][j].data[0]); + *(U2*)(&dst_ptr[(row+0)*row_stride + (col+half_col)]) = base_types::convertor::convert(src.tiles[i][j].data[2]); + } + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + pack_size*(warp_laneid % 4); + *(U2*)(&dst_ptr[(row+8)*row_stride + (col+0)]) = base_types::convertor::convert(src.tiles[i][j].data[1]); + *(U2*)(&dst_ptr[(row+8)*row_stride + (col+half_col)]) = base_types::convertor::convert(src.tiles[i][j].data[3]); + } + } +} + + + + + + + + +template>> +__device__ inline static void store(const GL &dst, const RT &src, const COORD &idx) { + using T = base_types::packing::unpacked_type; + using U = typename GL::dtype; + + #if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert(!std::is_same_v && !std::is_same_v, "Unsupported type for load/store"); + #endif + + U *dst_ptr = (U*)&dst[(idx.template unit_coord())]; + const int row_stride = dst.template stride(); + int warp_laneid = threadIdx.x % WARP_THREADS; + int local_warpid; + if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4)); + else local_warpid = warpid(); + const int row_offset = RT::rows*local_warpid; + #pragma unroll + for(int i = 0; i < RT::height; i++) { + int row = row_offset + i*RT::tile_size_row + 2*(warp_laneid % 4); + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + (warp_laneid / 4); + dst_ptr[(row+0)*row_stride + (col+0)] = base_types::convertor::convert(src.tiles[i][j].data[0].x); + dst_ptr[(row+0)*row_stride + (col+8)] = base_types::convertor::convert(src.tiles[i][j].data[1].x); + } + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + (warp_laneid / 4); + dst_ptr[(row+1)*row_stride + (col+0)] = base_types::convertor::convert(src.tiles[i][j].data[0].y); + dst_ptr[(row+1)*row_stride + (col+8)] = base_types::convertor::convert(src.tiles[i][j].data[1].y); + } + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + (warp_laneid / 4); + dst_ptr[(row+8)*row_stride + (col+0)] = base_types::convertor::convert(src.tiles[i][j].data[2].x); + dst_ptr[(row+8)*row_stride + (col+8)] = base_types::convertor::convert(src.tiles[i][j].data[3].x); + } + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + (warp_laneid / 4); + dst_ptr[(row+9)*row_stride + (col+0)] = base_types::convertor::convert(src.tiles[i][j].data[2].y); + dst_ptr[(row+9)*row_stride + (col+8)] = base_types::convertor::convert(src.tiles[i][j].data[3].y); + } + } +} +template>> +__device__ inline static void store(const GL &dst, const RT &src, const COORD &idx) { + store<2>(dst, src, idx); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/global_to_shared.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/global_to_shared.cuh new file mode 100644 index 0000000..4f2ca78 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/global_to_shared.cuh @@ -0,0 +1,167 @@ + + + + + + + + + + + + + + +template> +__device__ static inline void load(ST &dst, const GL &src, const COORD &idx) { + using T = typename ST::dtype; + const int row_stride = src.template stride(); + + constexpr int elem_per_memcpy = sizeof(float4)/sizeof(typename ST::dtype); + constexpr int memcpy_per_row = ST::cols / elem_per_memcpy; + constexpr int total_calls = (ST::rows*ST::cols + GROUP_THREADS*elem_per_memcpy-1) / (GROUP_THREADS*elem_per_memcpy); + + coord<> unit_coord = idx.template unit_coord(); + typename GL::dtype *src_ptr = (typename GL::dtype*)&src[unit_coord]; + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst.data[0])); + int laneid = threadIdx.x % GROUP_THREADS; + + #pragma unroll + for(int i = 0; i < total_calls; i++) { + + int load_idx = i * GROUP_THREADS + laneid; + + int row = load_idx / memcpy_per_row; + int col = (load_idx*elem_per_memcpy) % ST::cols; + + if constexpr (assume_aligned) { + float4 tmp; + move::ldg(tmp, (float4*)&src_ptr[row*row_stride + col]); + move::sts(dst.idx(dst_ptr, {row, col}), tmp); + } + else { + if (row + unit_coord.template dim() < src.template shape()) { + float4 tmp; + move::ldg(tmp, (float4*)&src_ptr[row*row_stride + col]); + move::sts(dst.idx(dst_ptr, {row, col}), tmp); + } + else { + float4 zeros = {0.f,0.f,0.f,0.f}; + move::sts(dst.idx(dst_ptr, {row, col}), zeros); + } + } + } +} +template> +__device__ static inline void load(ST &dst, const GL &src, const COORD &idx) { + load<2, false, ST, GL, COORD>(dst, src, idx); +} + + + + + + + + + +template> +__device__ static inline void store(const GL &dst, const ST &src, const COORD &idx) { + using T = typename ST::dtype; + const int row_stride = dst.template stride(); + + constexpr int elem_per_memcpy = sizeof(float4)/sizeof(typename ST::dtype); + constexpr int memcpy_per_row = ST::cols / elem_per_memcpy; + constexpr int total_calls = (ST::rows*ST::cols + GROUP_THREADS*elem_per_memcpy-1) / (GROUP_THREADS*elem_per_memcpy); + + coord<> unit_coord = idx.template unit_coord(); + typename GL::dtype *dst_ptr = (typename GL::dtype*)&dst[unit_coord]; + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src.data[0])); + int laneid = threadIdx.x % GROUP_THREADS; + + #pragma unroll + for(int i = 0; i < total_calls; i++) { + + int load_idx = i * GROUP_THREADS + laneid; + + int row = load_idx / memcpy_per_row; + int col = (load_idx*elem_per_memcpy) % ST::cols; + + if constexpr (assume_aligned) { + float4 tmp; + move::lds(tmp, src.idx(src_ptr, {row, col})); + move::stg((float4*)&dst_ptr[row*row_stride + col], tmp); + } + else { + if (row + unit_coord.template dim() < dst.template shape()) { + float4 tmp; + move::lds(tmp, src.idx(src_ptr, {row, col})); + move::stg((float4*)&dst_ptr[row*row_stride + col], tmp); + } + } + } +} +template> +__device__ static inline void store(const GL &dst, const ST &src, const COORD &idx) { + store<2, false, ST, GL, COORD>(dst, src, idx); +} + + + + + + + + + + +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx) { + using T = typename ST::dtype; + const int row_stride = src.template stride(); + + constexpr int elem_per_memcpy = sizeof(float4)/sizeof(typename ST::dtype); + constexpr int memcpy_per_row = ST::cols / elem_per_memcpy; + constexpr int total_calls = (ST::rows*ST::cols + GROUP_THREADS*elem_per_memcpy-1) / (GROUP_THREADS*elem_per_memcpy); + + coord<> unit_coord = idx.template unit_coord(); + typename GL::dtype *src_ptr = (typename GL::dtype*)&src[unit_coord]; + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst.data[0])); + int laneid = threadIdx.x % GROUP_THREADS; + + #pragma unroll + for(int i = 0; i < total_calls; i++) { + + int load_idx = i * GROUP_THREADS + laneid; + + int row = load_idx / memcpy_per_row; + int col = (load_idx*elem_per_memcpy) % ST::cols; + + if constexpr (assume_aligned) { + asm volatile( + "cp.async.cg.shared.global.L2::128B [%0], [%1], 16;\n" + :: "r"(dst.idx(dst_ptr, {row, col})), "l"(&src_ptr[row*row_stride + col]) + : "memory" + ); + } + else { + if (row + unit_coord.template dim() < src.template shape()) { + asm volatile( + "cp.async.cg.shared.global.L2::128B [%0], [%1], 16;\n" + :: "r"(dst.idx(dst_ptr, {row, col})), "l"(&src_ptr[row*row_stride + col]) + : "memory" + ); + } + else { + + float4 zeros = {0.f,0.f,0.f,0.f}; + move::sts(dst.idx(dst_ptr, {row, col}), zeros); + } + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); +} +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx) { + load_async<2, false, ST, GL, COORD>(dst, src, idx); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/parallel_global_to_global.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/parallel_global_to_global.cuh new file mode 100644 index 0000000..1ce7a5c --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/parallel_global_to_global.cuh @@ -0,0 +1,51 @@ + + + + + + + +template +__device__ inline static void all_reduce(PGL &pgl, const coord &idx) { + static_assert(std::is_same_v, "Currently only bf16 is supported."); + static_assert(TILE_COLS > 0 && TILE_COLS % (WARP_THREADS * 2) == 0, "TILE_COLS must be a multiple of 64 for best performance."); + + const int warps_per_row = TILE_COLS / (WARP_THREADS * 2); + const int row_base = idx.r * TILE_ROWS; + const int col_base = idx.c * TILE_COLS; + const int warp_laneid = threadIdx.x % WARP_THREADS; + + for (int i = warpid(); i < TILE_ROWS * warps_per_row; i += GROUP_WARPS) { + int row_idx = row_base + i / warps_per_row; + int col_idx = col_base + (i % warps_per_row) * WARP_THREADS * 2 + warp_laneid * 2; + bf16_2 tmp; + bf16_2 *ptr = reinterpret_cast(pgl.mc_ptr_at({idx.b, idx.d, row_idx, col_idx})); + multimem::ld_reduce(tmp, ptr); + multimem::st(ptr, tmp); + } +} + +template +__device__ inline static void reduce(GL &dst, const coord &dst_idx, PGL &src, const coord &src_idx) { + static_assert(std::is_same_v, "Currently only bf16 is supported."); + static_assert(TILE_COLS > 0 && TILE_COLS % (WARP_THREADS * 2) == 0, "TILE_COLS must be a multiple of 64 for best performance."); + + const int warps_per_row = TILE_COLS / (WARP_THREADS * 2); + const int src_row_base = src_idx.r * TILE_ROWS; + const int src_col_base = src_idx.c * TILE_COLS; + const int dst_row_base = dst_idx.r * TILE_ROWS; + const int dst_col_base = dst_idx.c * TILE_COLS; + const int warp_laneid = threadIdx.x % WARP_THREADS; + + for (int i = warpid(); i < TILE_ROWS * warps_per_row; i += GROUP_WARPS) { + int src_row_idx = src_row_base + i / warps_per_row; + int src_col_idx = src_col_base + (i % warps_per_row) * WARP_THREADS * 2 + warp_laneid * 2; + int dst_row_idx = dst_row_base + i / warps_per_row; + int dst_col_idx = dst_col_base + (i % warps_per_row) * WARP_THREADS * 2 + warp_laneid * 2; + (void)dst_row_idx; + (void)dst_col_idx; + bf16_2 tmp; + multimem::ld_reduce(tmp, reinterpret_cast(src.mc_ptr_at({src_idx.b, src_idx.d, src_row_idx, src_col_idx}))); + move::stg(reinterpret_cast(&dst[{dst_idx.b, dst_idx.d, dst_row_idx, dst_col_idx}]), tmp); + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/shared_to_register.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/shared_to_register.cuh new file mode 100644 index 0000000..eb43d9e --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/shared_to_register.cuh @@ -0,0 +1,321 @@ + + + + + + + + + + + + + +template +__device__ inline static void load(RT &dst, const ST &src) { + constexpr int warp_height = RT::height; + static_assert(ST::rows/RT::rows == GROUP_WARPS, "Group load / store requires tile height to be a multiple of the RT height."); + static_assert(ST::rows%RT::rows == 0, "Group load / store requires tile height to be a multiple of the RT height."); + static_assert(ST::cols==RT::cols, "Group load / store requires tile widths to match."); + int local_warpid; + if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4)); + else local_warpid = warpid(); + using T2 = RT::dtype; + using U = ST::dtype; + using T = base_types::packing::unpacked_type; + using U2 = base_types::packing::packed_type; + int warp_laneid = ::kittens::laneid(); + + + uint32_t shared_addr = static_cast(__cvta_generic_to_shared(&src.data[0])); + + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + if constexpr (sizeof(typename ST::dtype) == 2) { + + U2 tmp[4]; + int row = (local_warpid*warp_height + i)*RT::tile_size_row + (warp_laneid % 16); + int col = j*RT::tile_size_col + (warp_laneid / 16) * 8; + if constexpr (std::is_same_v) { + move::ldsm4(tmp[0], tmp[1], tmp[2], tmp[3], src.idx(shared_addr, {row, col})); + } + else { + move::ldsm4t(tmp[0], tmp[2], tmp[1], tmp[3], src.idx(shared_addr, {row, col})); + } + dst.tiles[i][j].data[0] = base_types::convertor::convert(tmp[0]); + dst.tiles[i][j].data[1] = base_types::convertor::convert(tmp[1]); + dst.tiles[i][j].data[2] = base_types::convertor::convert(tmp[2]); + dst.tiles[i][j].data[3] = base_types::convertor::convert(tmp[3]); + } + else if constexpr (std::is_same_v && sizeof(typename ST::dtype) == 1) { + + int warp_group_16 = (warp_laneid / 16); + int lane_in_16 = warp_laneid % 16; + int row = (local_warpid*warp_height + i)*RT::tile_size_row + (lane_in_16 % 16); + int col = j*RT::tile_size_col + warp_group_16 * 16; + + U2 tmp[4]; + if constexpr (std::is_same_v) { + move::ldsm4(tmp[0], tmp[1], tmp[2], tmp[3], src.idx(shared_addr, {row, col})); + } + else { + move::ldsm4t(tmp[0], tmp[2], tmp[1], tmp[3], src.idx(shared_addr, {row, col})); + } + dst.tiles[i][j].data[0] = base_types::convertor::convert(tmp[0]); + dst.tiles[i][j].data[1] = base_types::convertor::convert(tmp[1]); + dst.tiles[i][j].data[2] = base_types::convertor::convert(tmp[2]); + dst.tiles[i][j].data[3] = base_types::convertor::convert(tmp[3]); + } + else if constexpr (std::is_same_v && sizeof(typename ST::dtype) == 4) { + + int row = (local_warpid*warp_height + i)*RT::tile_size_row + (warp_laneid / 4); + int col = j*RT::tile_size_col + 2*(warp_laneid % 4); + if constexpr (ST::rows != ST::underlying_rows || ST::cols != ST::underlying_cols) { + row += src.row_offset; + col += src.col_offset; + } + int blit = sizeof(typename ST::dtype) * ((warp_laneid%4) / 2); + U2 tmp[4]; + static constexpr int swizzle_repeat = ST::swizzle_bytes * 8; + static constexpr int subtile_cols = ST::swizzle_bytes / sizeof(U); + const int outer_idx = col/subtile_cols; + const uint32_t addr_1 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+0)*subtile_cols + col%subtile_cols); + const uint32_t addr_2 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+8)*subtile_cols + col%subtile_cols); + const int swizzle_1 = blit ^ ((addr_1 % swizzle_repeat) >> 7) << 4; + const int swizzle_2 = blit ^ ((addr_2 % swizzle_repeat) >> 7) << 4; + move::lds(tmp[0].x, (addr_1+ 0)^swizzle_1); + move::lds(tmp[0].y, (addr_1+ 4)^swizzle_1); + move::lds(tmp[2].x, (addr_1+32)^swizzle_1); + move::lds(tmp[2].y, (addr_1+36)^swizzle_1); + move::lds(tmp[1].x, (addr_2+ 0)^swizzle_2); + move::lds(tmp[1].y, (addr_2+ 4)^swizzle_2); + move::lds(tmp[3].x, (addr_2+32)^swizzle_2); + move::lds(tmp[3].y, (addr_2+36)^swizzle_2); + dst.tiles[i][j].data[0] = base_types::convertor::convert(tmp[0]); + dst.tiles[i][j].data[1] = base_types::convertor::convert(tmp[1]); + dst.tiles[i][j].data[2] = base_types::convertor::convert(tmp[2]); + dst.tiles[i][j].data[3] = base_types::convertor::convert(tmp[3]); + if(blit) { + #pragma unroll + for(int k = 0; k < 4; k++) { + dst.tiles[i][j].data[k] = T2{dst.tiles[i][j].data[k].y, dst.tiles[i][j].data[k].x}; + } + } + } + else { + + int row = (local_warpid*warp_height + i)*RT::tile_size_row + 2*(warp_laneid % 4); + int col = j*RT::tile_size_col + (warp_laneid / 4); + U2 tmp[4]; + move::lds(tmp[0].x, src.idx(shared_addr, {row+0, col+0})); + move::lds(tmp[0].y, src.idx(shared_addr, {row+1, col+0})); + move::lds(tmp[1].x, src.idx(shared_addr, {row+0, col+8})); + move::lds(tmp[1].y, src.idx(shared_addr, {row+1, col+8})); + move::lds(tmp[2].x, src.idx(shared_addr, {row+8, col+0})); + move::lds(tmp[2].y, src.idx(shared_addr, {row+9, col+0})); + move::lds(tmp[3].x, src.idx(shared_addr, {row+8, col+8})); + move::lds(tmp[3].y, src.idx(shared_addr, {row+9, col+8})); + dst.tiles[i][j].data[0] = base_types::convertor::convert(tmp[0]); + dst.tiles[i][j].data[1] = base_types::convertor::convert(tmp[1]); + dst.tiles[i][j].data[2] = base_types::convertor::convert(tmp[2]); + dst.tiles[i][j].data[3] = base_types::convertor::convert(tmp[3]); + } + } + } +} + + + + + + + + + + +template +__device__ inline static void store(ST &dst, const RT &src) { + constexpr int warp_height = RT::height; + static_assert(ST::rows/RT::rows == GROUP_WARPS, "Group load / store requires tile height to be a multiple of the RT height."); + static_assert(ST::rows%RT::rows == 0, "Group load / store requires tile height to be a multiple of the RT height."); + static_assert(ST::cols==RT::cols, "Group load / store requires tile widths to match."); + int local_warpid; + if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4)); + else local_warpid = warpid(); + using T2 = RT::dtype; + using U = ST::dtype; + using T = base_types::packing::unpacked_type; + using U2 = base_types::packing::packed_type; + int warp_laneid = ::kittens::laneid(); + + + uint32_t shared_addr = static_cast(__cvta_generic_to_shared(&dst.data[0])); + + #pragma unroll + for(int i = 0; i < warp_height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + if constexpr (sizeof(typename ST::dtype) == 2) { + + U2 tmp[4]; + tmp[0] = base_types::convertor::convert(src.tiles[i][j].data[0]); + tmp[1] = base_types::convertor::convert(src.tiles[i][j].data[1]); + tmp[2] = base_types::convertor::convert(src.tiles[i][j].data[2]); + tmp[3] = base_types::convertor::convert(src.tiles[i][j].data[3]); +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + int row = (local_warpid*warp_height + i)*RT::tile_size_row + (warp_laneid % 16); + int col = j*RT::tile_size_col + (warp_laneid / 16) * 8; + if constexpr (std::is_same_v) { + move::stsm4(dst.idx(shared_addr, {row, col}), tmp[0], tmp[1], tmp[2], tmp[3]); + } + else { + move::stsm4t(dst.idx(shared_addr, {row, col}), tmp[0], tmp[2], tmp[1], tmp[3]); + } +#else + if constexpr (std::is_same_v) { + int row = (local_warpid*warp_height + i)*RT::tile_size_row + (warp_laneid / 4); + int col = j*RT::tile_size_col + 2*(warp_laneid % 4); + move::sts(dst.idx(shared_addr, {row+0, col+0}), tmp[0]); + move::sts(dst.idx(shared_addr, {row+8, col+0}), tmp[1]); + move::sts(dst.idx(shared_addr, {row+0, col+8}), tmp[2]); + move::sts(dst.idx(shared_addr, {row+8, col+8}), tmp[3]); + } + else { + int row = (local_warpid*warp_height + i)*RT::tile_size_row + 2*(warp_laneid % 4); + int col = j*RT::tile_size_col + (warp_laneid / 4); + move::sts(dst.idx(shared_addr, {row+0, col+0}), tmp[0].x); + move::sts(dst.idx(shared_addr, {row+1, col+0}), tmp[0].y); + move::sts(dst.idx(shared_addr, {row+0, col+8}), tmp[1].x); + move::sts(dst.idx(shared_addr, {row+1, col+8}), tmp[1].y); + move::sts(dst.idx(shared_addr, {row+8, col+0}), tmp[2].x); + move::sts(dst.idx(shared_addr, {row+9, col+0}), tmp[2].y); + move::sts(dst.idx(shared_addr, {row+8, col+8}), tmp[3].x); + move::sts(dst.idx(shared_addr, {row+9, col+8}), tmp[3].y); + } +#endif + } + else if constexpr (std::is_same_v && sizeof(typename ST::dtype) == 1) { + + + int warp_group_16 = (warp_laneid / 16); + int lane_in_16 = warp_laneid % 16; + int row = (local_warpid*warp_height + i)*RT::tile_size_row + (lane_in_16 % 16); + int col = j*RT::tile_size_col + warp_group_16 * 16; + + U2 tmp[4]; + tmp[0] = base_types::convertor::convert(src.tiles[i][j].data[0]); + tmp[1] = base_types::convertor::convert(src.tiles[i][j].data[1]); + tmp[2] = base_types::convertor::convert(src.tiles[i][j].data[2]); + tmp[3] = base_types::convertor::convert(src.tiles[i][j].data[3]); + if constexpr (std::is_same_v) { + move::stsm4(dst.idx(shared_addr, {row, col}), tmp[0], tmp[1], tmp[2], tmp[3]); + } + else { + move::stsm4t(dst.idx(shared_addr, {row, col}), tmp[0], tmp[2], tmp[1], tmp[3]); + } + } + else if constexpr (std::is_same_v && sizeof(typename ST::dtype) == 4) { + + int row = (local_warpid*warp_height + i)*RT::tile_size_row + (warp_laneid / 4); + int col = j*RT::tile_size_col + 2*(warp_laneid % 4); + if constexpr (ST::rows != ST::underlying_rows || ST::cols != ST::underlying_cols) { + row += dst.row_offset; + col += dst.col_offset; + } + int blit = sizeof(typename ST::dtype) * ((warp_laneid%4) / 2); + T2 reg_tmp[4]; + if(blit) { + #pragma unroll + for(int k = 0; k < 4; k++) { + reg_tmp[k] = T2{src.tiles[i][j].data[k].y, src.tiles[i][j].data[k].x}; + } + } + else { + #pragma unroll + for(int k = 0; k < 4; k++) { + reg_tmp[k] = src.tiles[i][j].data[k]; + } + } + U2 tmp[4]; + tmp[0] = base_types::convertor::convert(reg_tmp[0]); + tmp[1] = base_types::convertor::convert(reg_tmp[1]); + tmp[2] = base_types::convertor::convert(reg_tmp[2]); + tmp[3] = base_types::convertor::convert(reg_tmp[3]); + static constexpr int swizzle_repeat = ST::swizzle_bytes * 8; + static constexpr int subtile_cols = ST::swizzle_bytes / sizeof(U); + const int outer_idx = col/subtile_cols; + const uint32_t addr_1 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+0)*subtile_cols + col%subtile_cols); + const uint32_t addr_2 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+8)*subtile_cols + col%subtile_cols); + const int swizzle_1 = blit ^ ((addr_1 % swizzle_repeat) >> 7) << 4; + const int swizzle_2 = blit ^ ((addr_2 % swizzle_repeat) >> 7) << 4; + move::sts((addr_1+ 0)^swizzle_1, tmp[0].x); + move::sts((addr_1+ 4)^swizzle_1, tmp[0].y); + move::sts((addr_1+32)^swizzle_1, tmp[2].x); + move::sts((addr_1+36)^swizzle_1, tmp[2].y); + move::sts((addr_2+ 0)^swizzle_2, tmp[1].x); + move::sts((addr_2+ 4)^swizzle_2, tmp[1].y); + move::sts((addr_2+32)^swizzle_2, tmp[3].x); + move::sts((addr_2+36)^swizzle_2, tmp[3].y); + } + else { + + int row = (local_warpid*warp_height + i)*RT::tile_size_row + 2*(warp_laneid % 4); + int col = j*RT::tile_size_col + (warp_laneid / 4); + U2 tmp[4]; + tmp[0] = base_types::convertor::convert(src.tiles[i][j].data[0]); + tmp[1] = base_types::convertor::convert(src.tiles[i][j].data[1]); + tmp[2] = base_types::convertor::convert(src.tiles[i][j].data[2]); + tmp[3] = base_types::convertor::convert(src.tiles[i][j].data[3]); + move::sts(dst.idx(shared_addr, {row+0, col+0}), tmp[0].x); + move::sts(dst.idx(shared_addr, {row+1, col+0}), tmp[0].y); + move::sts(dst.idx(shared_addr, {row+0, col+8}), tmp[1].x); + move::sts(dst.idx(shared_addr, {row+1, col+8}), tmp[1].y); + move::sts(dst.idx(shared_addr, {row+8, col+0}), tmp[2].x); + move::sts(dst.idx(shared_addr, {row+9, col+0}), tmp[2].y); + move::sts(dst.idx(shared_addr, {row+8, col+8}), tmp[3].x); + move::sts(dst.idx(shared_addr, {row+9, col+8}), tmp[3].y); + } + } + } +} + + + +template +__device__ inline static auto load(RV &dst, const ST &src, int2 row_col) { + KITTENS_CHECK_WARP; + static_assert(ST::cols>=RV::length, "Shared tile must be at least as wide as the vector."); + using T = RV::T; + using U = ST::T; + int warp_laneid = ::kittens::laneid(); + + + uint32_t shared_addr = static_cast(__cvta_generic_to_shared(&src.data[0])); + + #pragma unroll + for(int col = warp_laneid; col < RV::length; col+=WARP_THREADS) { + U tmp; + move::lds(tmp, src.idx(shared_addr, {row_col.x, row_col.y + col})); + dst.data[col/WARP_THREADS][0] = base_types::convertor::convert(tmp); + } +} + +template +__device__ inline static auto store(ST &dst, const RV &src, int2 row_col) { + KITTENS_CHECK_WARP; + static_assert(ST::cols>=RV::length, "Shared tile must be at least as wide as the vector."); + using T = RV::T; + using U = ST::T; + int warp_laneid = ::kittens::laneid(); + + + uint32_t shared_addr = static_cast(__cvta_generic_to_shared(&dst.data[0])); + + #pragma unroll + for(int col = warp_laneid; col < RV::length; col+=WARP_THREADS) { + U tmp = base_types::convertor::convert(src.data[col/WARP_THREADS][0]); + move::sts(dst.idx(shared_addr, {row_col.x, row_col.y + col}), tmp); + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/tensor_to_register.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/tensor_to_register.cuh new file mode 100644 index 0000000..980c1c5 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/tensor_to_register.cuh @@ -0,0 +1,493 @@ + + + + + + + + + + + + + +template +__device__ inline static void load_async(RT &dst, const TM &src) { + if constexpr (GROUP_WARPS == 1) { + static_assert(RT::rows == TM::rows, "register tile and tensor tile must match rows"); + static_assert(RT::cols == TM::cols, "register tile and tensor tile must match cols"); + + using T2 = RT::dtype; + using U = typename TM::dtype; + using U2 = base_types::packing::packed_type; + + if constexpr (sizeof(typename TM::dtype) == 1) { + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + asm volatile( + "tcgen05.ld.sync.aligned.16x128b.x2.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(*(uint32_t*) &dst.tiles[i][j].data[0]), + "=r"(*(uint32_t*) &dst.tiles[i][j].data[1]), + "=r"(*(uint32_t*) &dst.tiles[i][j].data[2]), + "=r"(*(uint32_t*) &dst.tiles[i][j].data[3]) + : "r"(src.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))) + ); + } + } + } else if constexpr (sizeof(typename TM::dtype) == 2) { + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + asm volatile( + "tcgen05.ld.sync.aligned.16x128b.x2.pack::16b.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(*(uint32_t*) &dst.tiles[i][j].data[0]), + "=r"(*(uint32_t*) &dst.tiles[i][j].data[1]), + "=r"(*(uint32_t*) &dst.tiles[i][j].data[2]), + "=r"(*(uint32_t*) &dst.tiles[i][j].data[3]) + : "r"(src.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)) + ); + } + } + } + else if constexpr (sizeof(typename TM::dtype) == 4) { + if constexpr (std::is_same_v) { + #pragma unroll + for(int i = 0; i < RT::height; i++) { + if constexpr (RT::width%4 == 0) { + #pragma unroll + for(int j = 0; j < RT::width; j+=4) { + U2 data[16]; + 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];\n" + : "=r"(data[0].x), "=r"(data[0].y), + "=r"(data[1].x), "=r"(data[1].y), + "=r"(data[2].x), "=r"(data[2].y), + "=r"(data[3].x), "=r"(data[3].y), + "=r"(data[4].x), "=r"(data[4].y), + "=r"(data[5].x), "=r"(data[5].y), + "=r"(data[6].x), "=r"(data[6].y), + "=r"(data[7].x), "=r"(data[7].y), + "=r"(data[8].x), "=r"(data[8].y), + "=r"(data[9].x), "=r"(data[9].y), + "=r"(data[10].x), "=r"(data[10].y), + "=r"(data[11].x), "=r"(data[11].y), + "=r"(data[12].x), "=r"(data[12].y), + "=r"(data[13].x), "=r"(data[13].y), + "=r"(data[14].x), "=r"(data[14].y), + "=r"(data[15].x), "=r"(data[15].y) + : "r"(src.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))) + ); + #pragma unroll + for(int k = 0; k < 4; k++) { + dst.tiles[i][j+0].data[k] = base_types::convertor::convert(data[k]); + dst.tiles[i][j+1].data[k] = base_types::convertor::convert(data[k+4]); + dst.tiles[i][j+2].data[k] = base_types::convertor::convert(data[k+8]); + dst.tiles[i][j+3].data[k] = base_types::convertor::convert(data[k+12]); + } + } + } + else if constexpr (RT::width%2 == 0) { + #pragma unroll + for(int j = 0; j < RT::width; j+=2) { + U2 data[8]; + 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];\n" + : "=r"(data[0].x), "=r"(data[0].y), + "=r"(data[1].x), "=r"(data[1].y), + "=r"(data[2].x), "=r"(data[2].y), + "=r"(data[3].x), "=r"(data[3].y), + "=r"(data[4].x), "=r"(data[4].y), + "=r"(data[5].x), "=r"(data[5].y), + "=r"(data[6].x), "=r"(data[6].y), + "=r"(data[7].x), "=r"(data[7].y) + : "r"(src.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))) + ); + #pragma unroll + for(int k = 0; k < 4; k++) { + dst.tiles[i][j+0].data[k] = base_types::convertor::convert(data[k]); + dst.tiles[i][j+1].data[k] = base_types::convertor::convert(data[k+4]); + } + } + } + else { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + U2 data[4]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n" + : "=r"(data[0].x), "=r"(data[0].y), + "=r"(data[1].x), "=r"(data[1].y), + "=r"(data[2].x), "=r"(data[2].y), + "=r"(data[3].x), "=r"(data[3].y) + : "r"(src.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))) + ); + #pragma unroll + for(int k = 0; k < 4; k++) { + dst.tiles[i][j].data[k] = base_types::convertor::convert(data[k]); + } + } + } + } + } + else if constexpr (std::is_same_v) { + #pragma unroll + for(int i = 0; i < RT::height; i++) { + if constexpr (RT::width%4 == 0) { + #pragma unroll + for(int j = 0; j < RT::width; j+=4) { + U2 data[16]; + 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];\n" + : "=f"(data[0].x), "=f"(data[0].y), + "=f"(data[1].x), "=f"(data[1].y), + "=f"(data[2].x), "=f"(data[2].y), + "=f"(data[3].x), "=f"(data[3].y), + "=f"(data[4].x), "=f"(data[4].y), + "=f"(data[5].x), "=f"(data[5].y), + "=f"(data[6].x), "=f"(data[6].y), + "=f"(data[7].x), "=f"(data[7].y), + "=f"(data[8].x), "=f"(data[8].y), + "=f"(data[9].x), "=f"(data[9].y), + "=f"(data[10].x), "=f"(data[10].y), + "=f"(data[11].x), "=f"(data[11].y), + "=f"(data[12].x), "=f"(data[12].y), + "=f"(data[13].x), "=f"(data[13].y), + "=f"(data[14].x), "=f"(data[14].y), + "=f"(data[15].x), "=f"(data[15].y) + : "r"(src.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))) + ); + #pragma unroll + for(int k = 0; k < 4; k++) { + dst.tiles[i][j+0].data[k] = base_types::convertor::convert(data[k]); + dst.tiles[i][j+1].data[k] = base_types::convertor::convert(data[k+4]); + dst.tiles[i][j+2].data[k] = base_types::convertor::convert(data[k+8]); + dst.tiles[i][j+3].data[k] = base_types::convertor::convert(data[k+12]); + } + } + } + else if constexpr (RT::width%2 == 0) { + #pragma unroll + for(int j = 0; j < RT::width; j+=2) { + U2 data[8]; + 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];\n" + : "=f"(data[0].x), "=f"(data[0].y), + "=f"(data[1].x), "=f"(data[1].y), + "=f"(data[2].x), "=f"(data[2].y), + "=f"(data[3].x), "=f"(data[3].y), + "=f"(data[4].x), "=f"(data[4].y), + "=f"(data[5].x), "=f"(data[5].y), + "=f"(data[6].x), "=f"(data[6].y), + "=f"(data[7].x), "=f"(data[7].y) + : "r"(src.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))) + ); + #pragma unroll + for(int k = 0; k < 4; k++) { + dst.tiles[i][j+0].data[k] = base_types::convertor::convert(data[k]); + dst.tiles[i][j+1].data[k] = base_types::convertor::convert(data[k+4]); + } + } + } + else { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + U2 data[4]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n" + : "=f"(data[0].x), "=f"(data[0].y), + "=f"(data[1].x), "=f"(data[1].y), + "=f"(data[2].x), "=f"(data[2].y), + "=f"(data[3].x), "=f"(data[3].y) + : "r"(src.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))) + ); + #pragma unroll + for(int k = 0; k < 4; k++) { + dst.tiles[i][j].data[k] = base_types::convertor::convert(data[k]); + } + } + } + } + } + else { + static_assert(sizeof(U) == 999, "Unsupported 4-byte tensor memory type."); + } + } + } + else { + static_assert(GROUP_WARPS==4 || GROUP_WARPS==8); + constexpr int warp_rows = TM::rows/GROUP_WARPS; + static_assert(TM::cols==RT::cols); + static_assert(warp_rows==RT::rows); + if constexpr (GROUP_WARPS == 4) { + auto src_subtile = src.template subtile>(32*warpid(), 0); + ::kittens::group<1>::load_async(dst, src_subtile); + } + else { + auto src_subtile = src.template subtile>(32*(warpid()%4)+16*(warpid()/4), 0); + ::kittens::group<1>::load_async(dst, src_subtile); + } + } +} + + + + + + + + + + +template +__device__ inline static void store_async(TM &dst, const RT &src) { + if constexpr (GROUP_WARPS == 1) { + static_assert(RT::rows == TM::rows, "register tile and tensor tile must match rows"); + static_assert(RT::cols == TM::cols, "register tile and tensor tile must match cols"); + + using T2 = RT::dtype; + using T = base_types::packing::unpacked_type; + using U = TM::dtype; + using U2 = base_types::packing::packed_type; + + if constexpr (sizeof(typename TM::dtype) <= 2) { + #pragma unroll + for(int i = 0; i < RT::height; i++) { + if constexpr (RT::width%4 == 0) { + #pragma unroll + for(int j = 0; j < RT::width; j+=4) { + 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};\n" + :: "r"(dst.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))), + "r"(*(uint32_t*)&src.tiles[i][j+0].data[0]), + "r"(*(uint32_t*)&src.tiles[i][j+0].data[1]), + "r"(*(uint32_t*)&src.tiles[i][j+0].data[2]), + "r"(*(uint32_t*)&src.tiles[i][j+0].data[3]), + "r"(*(uint32_t*)&src.tiles[i][j+1].data[0]), + "r"(*(uint32_t*)&src.tiles[i][j+1].data[1]), + "r"(*(uint32_t*)&src.tiles[i][j+1].data[2]), + "r"(*(uint32_t*)&src.tiles[i][j+1].data[3]), + "r"(*(uint32_t*)&src.tiles[i][j+2].data[0]), + "r"(*(uint32_t*)&src.tiles[i][j+2].data[1]), + "r"(*(uint32_t*)&src.tiles[i][j+2].data[2]), + "r"(*(uint32_t*)&src.tiles[i][j+2].data[3]), + "r"(*(uint32_t*)&src.tiles[i][j+3].data[0]), + "r"(*(uint32_t*)&src.tiles[i][j+3].data[1]), + "r"(*(uint32_t*)&src.tiles[i][j+3].data[2]), + "r"(*(uint32_t*)&src.tiles[i][j+3].data[3]) + ); + } + } + else if constexpr (RT::width%2 == 0) { + #pragma unroll + for(int j = 0; j < RT::width; j+=2) { + asm volatile( + "tcgen05.st.sync.aligned.16x128b.x4.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};\n" + :: "r"(dst.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))), + "r"(*(uint32_t*)&src.tiles[i][j+0].data[0]), + "r"(*(uint32_t*)&src.tiles[i][j+0].data[1]), + "r"(*(uint32_t*)&src.tiles[i][j+0].data[2]), + "r"(*(uint32_t*)&src.tiles[i][j+0].data[3]), + "r"(*(uint32_t*)&src.tiles[i][j+1].data[0]), + "r"(*(uint32_t*)&src.tiles[i][j+1].data[1]), + "r"(*(uint32_t*)&src.tiles[i][j+1].data[2]), + "r"(*(uint32_t*)&src.tiles[i][j+1].data[3]) + ); + } + } + else { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + asm volatile( + "tcgen05.st.sync.aligned.16x128b.x2.b32 [%0], {%1, %2, %3, %4};\n" + :: "r"(dst.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))), + "r"(*(uint32_t*)&src.tiles[i][j].data[0]), + "r"(*(uint32_t*)&src.tiles[i][j].data[1]), + "r"(*(uint32_t*)&src.tiles[i][j].data[2]), + "r"(*(uint32_t*)&src.tiles[i][j].data[3]) + ); + } + } + } + } + else if constexpr (sizeof(typename TM::dtype) == 4) { + if constexpr (std::is_same_v) { + #pragma unroll + for(int i = 0; i < RT::height; i++) { + if constexpr(RT::width%4 == 0) { + #pragma unroll + for(int j = 0; j < RT::width; j+=4) { + U2 data[16]; + #pragma unroll + for(int k = 0; k < 4; k++) { + data[k] = base_types::convertor::convert(src.tiles[i][j].data[k]); + data[k+4] = base_types::convertor::convert(src.tiles[i][j+1].data[k]); + data[k+8] = base_types::convertor::convert(src.tiles[i][j+2].data[k]); + data[k+12] = base_types::convertor::convert(src.tiles[i][j+3].data[k]); + } + 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};\n" + :: "r"(dst.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))), + "r"(data[0].x), "r"(data[0].y), + "r"(data[1].x), "r"(data[1].y), + "r"(data[2].x), "r"(data[2].y), + "r"(data[3].x), "r"(data[3].y), + "r"(data[4].x), "r"(data[4].y), + "r"(data[5].x), "r"(data[5].y), + "r"(data[6].x), "r"(data[6].y), + "r"(data[7].x), "r"(data[7].y), + "r"(data[8].x), "r"(data[8].y), + "r"(data[9].x), "r"(data[9].y), + "r"(data[10].x), "r"(data[10].y), + "r"(data[11].x), "r"(data[11].y), + "r"(data[12].x), "r"(data[12].y), + "r"(data[13].x), "r"(data[13].y), + "r"(data[14].x), "r"(data[14].y), + "r"(data[15].x), "r"(data[15].y) + ); + } + } + else if constexpr(RT::width%2 == 0) { + #pragma unroll + for(int j = 0; j < RT::width; j+=2) { + U2 data[8]; + #pragma unroll + for(int k = 0; k < 4; k++) { + data[k] = base_types::convertor::convert(src.tiles[i][j].data[k]); + data[k+4] = base_types::convertor::convert(src.tiles[i][j+1].data[k]); + } + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x4.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};\n" + :: "r"(dst.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))), + "r"(data[0].x), "r"(data[0].y), + "r"(data[1].x), "r"(data[1].y), + "r"(data[2].x), "r"(data[2].y), + "r"(data[3].x), "r"(data[3].y), + "r"(data[4].x), "r"(data[4].y), + "r"(data[5].x), "r"(data[5].y), + "r"(data[6].x), "r"(data[6].y), + "r"(data[7].x), "r"(data[7].y) + ); + } + } + else { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + U2 data[4]; + #pragma unroll + for(int k = 0; k < 4; k++) { + data[k] = base_types::convertor::convert(src.tiles[i][j].data[k]); + } + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x2.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};\n" + :: "r"(dst.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))), + "r"(data[0].x), "r"(data[0].y), + "r"(data[1].x), "r"(data[1].y), + "r"(data[2].x), "r"(data[2].y), + "r"(data[3].x), "r"(data[3].y) + ); + } + } + } + } + else if constexpr (std::is_same_v) { + #pragma unroll + for(int i = 0; i < RT::height; i++) { + if constexpr(RT::width%4 == 0) { + #pragma unroll + for(int j = 0; j < RT::width; j+=4) { + U2 data[16]; + #pragma unroll + for(int k = 0; k < 4; k++) { + data[k] = base_types::convertor::convert(src.tiles[i][j].data[k]); + data[k+4] = base_types::convertor::convert(src.tiles[i][j+1].data[k]); + data[k+8] = base_types::convertor::convert(src.tiles[i][j+2].data[k]); + data[k+12] = base_types::convertor::convert(src.tiles[i][j+3].data[k]); + } + 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};\n" + :: "r"(dst.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))), + "f"(data[0].x), "f"(data[0].y), + "f"(data[1].x), "f"(data[1].y), + "f"(data[2].x), "f"(data[2].y), + "f"(data[3].x), "f"(data[3].y), + "f"(data[4].x), "f"(data[4].y), + "f"(data[5].x), "f"(data[5].y), + "f"(data[6].x), "f"(data[6].y), + "f"(data[7].x), "f"(data[7].y), + "f"(data[8].x), "f"(data[8].y), + "f"(data[9].x), "f"(data[9].y), + "f"(data[10].x), "f"(data[10].y), + "f"(data[11].x), "f"(data[11].y), + "f"(data[12].x), "f"(data[12].y), + "f"(data[13].x), "f"(data[13].y), + "f"(data[14].x), "f"(data[14].y), + "f"(data[15].x), "f"(data[15].y) + ); + } + } + else if constexpr(RT::width%2 == 0) { + #pragma unroll + for(int j = 0; j < RT::width; j+=2) { + U2 data[8]; + #pragma unroll + for(int k = 0; k < 4; k++) { + data[k] = base_types::convertor::convert(src.tiles[i][j].data[k]); + data[k+4] = base_types::convertor::convert(src.tiles[i][j+1].data[k]); + } + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x4.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};\n" + :: "r"(dst.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))), + "f"(data[0].x), "f"(data[0].y), + "f"(data[1].x), "f"(data[1].y), + "f"(data[2].x), "f"(data[2].y), + "f"(data[3].x), "f"(data[3].y), + "f"(data[4].x), "f"(data[4].y), + "f"(data[5].x), "f"(data[5].y), + "f"(data[6].x), "f"(data[6].y), + "f"(data[7].x), "f"(data[7].y) + ); + } + } + else { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + U2 data[4]; + #pragma unroll + for(int k = 0; k < 4; k++) { + data[k] = base_types::convertor::convert(src.tiles[i][j].data[k]); + } + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x2.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};\n" + :: "r"(dst.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))), + "f"(data[0].x), "f"(data[0].y), + "f"(data[1].x), "f"(data[1].y), + "f"(data[2].x), "f"(data[2].y), + "f"(data[3].x), "f"(data[3].y) + ); + } + } + } + } + else { + static_assert(sizeof(U) == 999, "Unsupported 4-byte tensor memory type."); + } + } + } + else { + static_assert(GROUP_WARPS==4 || GROUP_WARPS==8); + constexpr int warp_rows = TM::rows/GROUP_WARPS; + static_assert(TM::cols==RT::cols); + static_assert(warp_rows==RT::rows); + if constexpr (GROUP_WARPS == 4) { + auto dst_subtile = dst.template subtile>(32*warpid(), 0); + ::kittens::group<1>::store_async(dst_subtile, src); + } + else { + auto dst_subtile = dst.template subtile>(32*(warpid()%4)+16*(warpid()/4), 0); + ::kittens::group<1>::store_async(dst_subtile, src); + } + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/tile.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/tile.cuh new file mode 100644 index 0000000..b67fee5 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/tile.cuh @@ -0,0 +1,20 @@ + + + + + +#include "global_to_register.cuh" +#include "global_to_shared.cuh" +#include "shared_to_register.cuh" +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +#include "parallel_global_to_global.cuh" +#endif +#ifdef KITTENS_SM10X +#include "tensor_to_register.cuh" +#endif + +#include "complex_shared_to_register.cuh" +#include "complex_global_to_register.cuh" +#include "complex_global_to_shared.cuh" + + diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/tma.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/tma.cuh new file mode 100644 index 0000000..4f0a208 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/tma.cuh @@ -0,0 +1,134 @@ + + + + + +template> +__device__ static inline void prefetch(ST &dst, const GL &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::prefetch(dst, src, idx); + } +} +template> +__device__ static inline void prefetch(ST &dst, const GL &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::prefetch(dst, src, idx); + } +} + +template> +__device__ static inline void store_async(const GL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_async(dst, src, idx); + } +} +template> +__device__ static inline void store_async(const GL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_async(dst, src, idx); + } +} + +template> +__device__ static inline void store_async(const PGL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_async(dst, src, idx); + } +} +template> +__device__ static inline void store_async(const PGL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_async(dst, src, idx); + } +} + +template> +__device__ static inline void store_add_async(const GL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_add_async(dst, src, idx); + } +} +template> +__device__ static inline void store_add_async(const GL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_add_async(dst, src, idx); + } +} + +template> +__device__ static inline void store_add_async(const PGL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_add_async(dst, src, idx); + } +} +template> +__device__ static inline void store_add_async(const PGL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_add_async(dst, src, idx); + } +} + +template> +__device__ static inline void store_min_async(const GL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_min_async(dst, src, idx); + } +} +template> +__device__ static inline void store_min_async(const GL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_min_async(dst, src, idx); + } +} + +template> +__device__ static inline void store_min_async(const PGL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_min_async(dst, src, idx); + } +} +template> +__device__ static inline void store_min_async(const PGL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_min_async(dst, src, idx); + } +} + +template> +__device__ static inline void store_max_async(const GL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_max_async(dst, src, idx); + } +} +template> +__device__ static inline void store_max_async(const GL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_max_async(dst, src, idx); + } +} + +template> +__device__ static inline void store_max_async(const PGL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_max_async(dst, src, idx); + } +} +template> +__device__ static inline void store_max_async(const PGL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_max_async(dst, src, idx); + } +} + +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar) { + if(laneid() == 0) { + ::kittens::tma::load_async(dst, src, idx, bar); + } +} +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar) { + if(laneid() == 0) { + ::kittens::tma::load_async(dst, src, idx, bar); + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/tma_cluster.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/tma_cluster.cuh new file mode 100644 index 0000000..6982df8 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/tma_cluster.cuh @@ -0,0 +1,32 @@ + + + + + +#ifdef KITTENS_SM10X +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) { + if(laneid() == 0) { + ::kittens::tma::cluster::load_async(dst, src, idx, bar, cluster_mask, dst_mbar_cta); + } +} +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) { + if(laneid() == 0) { + ::kittens::tma::cluster::load_async(dst, src, idx, bar, cluster_mask, dst_mbar_cta); + } +} +#else +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask) { + if(laneid() == 0) { + ::kittens::tma::cluster::load_async(dst, src, idx, bar, cluster_mask); + } +} +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask) { + if(laneid() == 0) { + ::kittens::tma::cluster::load_async(dst, src, idx, bar, cluster_mask); + } +} +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/global_to_register.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/global_to_register.cuh new file mode 100644 index 0000000..1540cd8 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/global_to_register.cuh @@ -0,0 +1,138 @@ + + + + + + + + + + + + + +template +__device__ inline static void load(RV &dst, const GL &src, const coord> &idx) { + if constexpr (GROUP_WARPS == 1) { + using T2 = RV::dtype; + using U = typename GL::dtype; + using U2 = base_types::packing::packed_type; + using T = base_types::packing::unpacked_type; + + U *src_ptr = (U*)&src[(idx.template unit_coord<-1, 3>())]; + int laneid = ::kittens::laneid(); + + if constexpr (std::is_same_v) { + #pragma unroll + for(auto w = 0; w < (RV::outer_dim+3)/4; w++) { + int idx = w*64 + (laneid/4)*8 + 2*(laneid%4); + int o_dim = w*4 + (laneid/4) / 2; + int i_dim = (laneid/4) % 2; + + if(idx < RV::outer_dim*16) + dst[o_dim][i_dim] = base_types::convertor::convert(*(U2*)&src_ptr[idx]); + } + + #pragma unroll + for(auto w = 0; w < RV::outer_dim; w++) { + int leader = 8*(w%4) + (laneid%4); + dst[w][0] = packed_shfl_sync(MASK_ALL, dst[w][0], leader); + dst[w][1] = packed_shfl_sync(MASK_ALL, dst[w][1], leader+4); + } + } + else if constexpr (std::is_same_v) { + + + #pragma unroll + for(auto w = 0; w < (RV::outer_dim+1)/2; w++) { + int idx = w*32 + (laneid%4)*8 + (laneid/4); + int o_dim = w*2 + (laneid%4) / 2; + + if(idx < RV::outer_dim*16) { + T tmp = base_types::convertor::convert(src_ptr[idx]); + if(laneid%2==0) dst[o_dim][0].x = tmp; + else dst[o_dim][0].y = tmp; + } + } + + #pragma unroll + for(auto w = 0; w < RV::outer_dim; w++) { + int leader = (laneid/4)*4 + 2*(w%2); + dst[w][0].x = __shfl_sync(MASK_ALL, dst[w][0].x, leader); + dst[w][0].y = __shfl_sync(MASK_ALL, dst[w][0].y, leader+1); + } + } + else if constexpr (std::is_same_v) { + #pragma unroll + for(auto w = 0; w < RV::outer_dim; w++) { + if(w < RV::outer_dim-1 || RV::length%32 == 0 || laneid<16) { + dst[w][0] = base_types::convertor::convert(src_ptr[w*32 + laneid]); + } + } + } + } + else { + + ::kittens::group<1>::load(dst, src, coord(idx.b, idx.d, idx.r, idx.c*GROUP_WARPS+warpid())); + } +} + + + + + + + + +template +__device__ inline static void store(GL &dst, const RV &src, const coord> &idx) { + if constexpr (GROUP_WARPS == 1) { + using T2 = RV::dtype; + using U = typename GL::dtype; + using U2 = base_types::packing::packed_type; + using T = base_types::packing::unpacked_type; + + U *dst_ptr = (U*)&dst[(idx.template unit_coord<-1, 3>())]; + int laneid = ::kittens::laneid(); + + if constexpr (std::is_same_v) { + #pragma unroll + for(auto w = 0; w < (RV::outer_dim+3)/4; w++) { + int idx = w*64 + (laneid/4)*8 + 2*(laneid%4); + int o_dim = w*4 + (laneid/4) / 2; + int i_dim = (laneid/4) % 2; + + if(idx < RV::outer_dim*16) + *(U2*)&dst_ptr[idx] = base_types::convertor::convert(src[o_dim][i_dim]); + } + } + else if constexpr (std::is_same_v) { + + + #pragma unroll + for(auto w = 0; w < (RV::outer_dim+1)/2; w++) { + int idx = w*32 + (laneid%4)*8 + (laneid/4); + int o_dim = w*2 + (laneid%4) / 2; + + if(idx < RV::outer_dim*16) { + U tmp; + if(laneid%2==0) tmp = base_types::convertor::convert(src[o_dim][0].x); + else tmp = base_types::convertor::convert(src[o_dim][0].y); + dst_ptr[idx] = tmp; + } + } + } + else if constexpr (std::is_same_v) { + #pragma unroll + for(auto w = 0; w < RV::outer_dim; w++) { + if(w < RV::outer_dim-1 || RV::length%32 == 0 || laneid<16) { + dst_ptr[w*32 + laneid] = base_types::convertor::convert(src[w][0]); + } + } + } + } + else { + + ::kittens::group<1>::store(dst, src, coord(idx.b, idx.d, idx.r, idx.c*GROUP_WARPS+warpid())); + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/global_to_shared.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/global_to_shared.cuh new file mode 100644 index 0000000..cf265b6 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/global_to_shared.cuh @@ -0,0 +1,113 @@ + + + + + + + + + + + + + + + + +template> +__device__ static inline void load(SV &dst, const GL &src, const COORD &idx) { + constexpr uint32_t elem_per_transfer = sizeof(float4) / sizeof(typename SV::dtype); + constexpr uint32_t total_calls = SV::length / elem_per_transfer; + typename GL::dtype *src_ptr = (typename GL::dtype*)&src[(idx.template unit_coord<-1, 3>())]; + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst.data[0])); + #pragma unroll + for(uint32_t i = threadIdx.x%GROUP_THREADS; i < total_calls; i+=GROUP_THREADS) { + if(i * elem_per_transfer < SV::length) { + float4 tmp; + move::ldg(tmp, (float4*)&src_ptr[i*elem_per_transfer]); + move::sts(dst_ptr + sizeof(typename SV::dtype)*i*elem_per_transfer, tmp); + } + } +} + + + + + + + + + + + + +template> +__device__ static inline void store(GL &dst, const SV &src, const COORD &idx) { + constexpr uint32_t elem_per_transfer = sizeof(float4) / sizeof(typename SV::dtype); + constexpr uint32_t total_calls = SV::length / elem_per_transfer; + typename GL::dtype *dst_ptr = (typename GL::dtype*)&dst[(idx.template unit_coord<-1, 3>())]; + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src.data[0])); + #pragma unroll + for(uint32_t i = threadIdx.x%GROUP_THREADS; i < total_calls; i+=GROUP_THREADS) { + if(i * elem_per_transfer < SV::length) { + float4 tmp; + move::lds(tmp, src_ptr + sizeof(typename SV::dtype)*i*elem_per_transfer); + move::stg((float4*)&dst_ptr[i*elem_per_transfer], tmp); + } + } +} + + + + + + + + + +template> +__device__ static inline void load_async(SV &dst, const GL &src, const COORD &idx) { + constexpr uint32_t elem_per_transfer = sizeof(float4) / sizeof(typename SV::dtype); + constexpr uint32_t total_calls = SV::length / elem_per_transfer; + typename GL::dtype *src_ptr = (typename GL::dtype*)&src[(idx.template unit_coord<-1, 3>())]; + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst.data[0])); + #pragma unroll + for(uint32_t i = threadIdx.x%GROUP_THREADS; i < total_calls; i+=GROUP_THREADS) { + if(i * elem_per_transfer < SV::length) { + asm volatile( + "cp.async.cg.shared.global.L2::128B [%0], [%1], 16;\n" + :: "r"(dst_ptr + (uint32_t)sizeof(typename SV::dtype)*i*elem_per_transfer), "l"((uint64_t)&src_ptr[i*elem_per_transfer]) + : "memory" + ); + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); +} + + + + + + + + + + +template> +__device__ static inline void load_async_small(SV &dst, const GL &src, const COORD &idx) { + constexpr uint32_t elem_per_transfer = sizeof(float) / sizeof(typename SV::dtype); + constexpr uint32_t total_calls = SV::length / elem_per_transfer; + typename GL::dtype *src_ptr = (typename GL::dtype*)&src[(idx.template unit_coord<-1, 3>())]; + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst.data[0])); + #pragma unroll + for(uint32_t i = threadIdx.x%GROUP_THREADS; i < total_calls; i+=GROUP_THREADS) { + if(i * elem_per_transfer < SV::length) { + asm volatile( + "cp.async.ca.shared.global.L2::128B [%0], [%1], 4;\n" + :: "r"(dst_ptr + (uint32_t)sizeof(typename SV::dtype)*i*elem_per_transfer), "l"((uint64_t)&src_ptr[i*elem_per_transfer]) + : "memory" + ); + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/shared_to_register.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/shared_to_register.cuh new file mode 100644 index 0000000..1a279ae --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/shared_to_register.cuh @@ -0,0 +1,159 @@ + + + + + + + + + + + + + +template +__device__ inline static void load(RV &dst, const SV &src) { + using T2 = RV::dtype; + using U = SV::dtype; + using U2 = base_types::packing::packed_type; + using T = base_types::packing::unpacked_type; + if constexpr (GROUP_WARPS == 1) { + static_assert(SV::length == RV::length); + + int laneid = ::kittens::laneid(); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src.data[0])); + + __syncwarp(); + if constexpr (std::is_same_v) { + #pragma unroll + for(auto w = 0; w < (RV::outer_dim+3)/4; w++) { + int idx = w*64 + (laneid/4)*8 + 2*(laneid%4); + int o_dim = w*4 + (laneid/4) / 2; + int i_dim = (laneid/4) % 2; + + if(idx < RV::outer_dim*16) { + U2 tmp; + move::lds(tmp, src_ptr + sizeof(typename SV::dtype)*idx); + dst[o_dim][i_dim] = base_types::convertor::convert(tmp); + } + } + __syncwarp(); + + #pragma unroll + for(auto w = 0; w < RV::outer_dim; w++) { + int leader = 8*(w%4) + (laneid%4); + dst[w][0] = packed_shfl_sync(MASK_ALL, dst[w][0], leader); + dst[w][1] = packed_shfl_sync(MASK_ALL, dst[w][1], leader+4); + } + } + else if constexpr (std::is_same_v) { + + + #pragma unroll + for(auto w = 0; w < (RV::outer_dim+1)/2; w++) { + int idx = w*32 + (laneid%4)*8 + (laneid/4); + int o_dim = w*2 + (laneid%4) / 2; + + if(idx < RV::outer_dim*16) { + U tmp; + move::lds(tmp, src_ptr + sizeof(typename SV::dtype)*idx); + if(laneid%2==0) dst[o_dim][0].x = base_types::convertor::convert(tmp); + else dst[o_dim][0].y = base_types::convertor::convert(tmp); + } + } + __syncwarp(); + + #pragma unroll + for(auto w = 0; w < RV::outer_dim; w++) { + int leader = (laneid/4)*4 + 2*(w%2); + dst[w][0].x = __shfl_sync(MASK_ALL, dst[w][0].x, leader); + dst[w][0].y = __shfl_sync(MASK_ALL, dst[w][0].y, leader+1); + } + } + else if constexpr (std::is_same_v) { + #pragma unroll + for(auto w = 0; w < RV::outer_dim; w++) { + if(w < RV::outer_dim-1 || RV::length%32 == 0 || laneid<16) { + U tmp; + move::lds(tmp, src_ptr + sizeof(typename SV::dtype)*(w*32 + laneid)); + dst[w][0] = base_types::convertor::convert(tmp); + } + } + } + } + else { + static_assert(SV::length == RV::length*GROUP_WARPS); + auto &_src = src.template subvec(warpid()); + + ::kittens::group<1>::load(dst, _src); + } +} + + + + + + + + + +template +__device__ inline static void store(SV &dst, const RV &src) { + using T2 = RV::dtype; + using U = SV::dtype; + using U2 = base_types::packing::packed_type; + using T = base_types::packing::unpacked_type; + + if constexpr (GROUP_WARPS == 1) { + static_assert(SV::length == RV::length); + + int laneid = ::kittens::laneid(); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst.data[0])); + + __syncwarp(); + if constexpr (std::is_same_v) { + #pragma unroll + for(auto w = 0; w < (RV::outer_dim+3)/4; w++) { + int idx = w*64 + (laneid/4)*8 + 2*(laneid%4); + int o_dim = w*4 + (laneid/4) / 2; + int i_dim = (laneid/4) % 2; + + if(idx < RV::outer_dim*16) { + U2 tmp = base_types::convertor::convert(src[o_dim][i_dim]); + move::sts(dst_ptr + sizeof(typename SV::dtype)*idx, tmp); + } + } + } + else if constexpr (std::is_same_v) { + + + #pragma unroll + for(auto w = 0; w < (RV::outer_dim+1)/2; w++) { + int idx = w*32 + (laneid%4)*8 + (laneid/4); + int o_dim = w*2 + (laneid%4) / 2; + + if(idx < RV::outer_dim*16) { + U tmp; + if(laneid%2==0) tmp = base_types::convertor::convert(src[o_dim][0].x); + else tmp = base_types::convertor::convert(src[o_dim][0].y); + move::sts(dst_ptr + sizeof(typename SV::dtype)*idx, tmp); + } + } + } + else if constexpr (std::is_same_v) { + #pragma unroll + for(auto w = 0; w < RV::outer_dim; w++) { + if(w < RV::outer_dim-1 || RV::length%32 == 0 || laneid<16) { + U tmp = base_types::convertor::convert(src[w][0]); + move::sts(dst_ptr + sizeof(typename SV::dtype)*(w*32 + laneid), tmp); + } + } + } + } + else { + static_assert(SV::length == RV::length*GROUP_WARPS); + auto &_dst = dst.template subvec(warpid()); + + ::kittens::group<1>::store(_dst, src); + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/tma.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/tma.cuh new file mode 100644 index 0000000..2f9ddaa --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/tma.cuh @@ -0,0 +1,221 @@ + + + + + + + + + + + + + + + +template> +__device__ static inline void prefetch(SV &dst, const GL &src, const COORD &idx) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(src.template get_tma()); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + ::kittens::detail::tma::vec_prefetch_tma_internal(tma_ptr, tma_coord); + } +} +__KITTENS_TMA_DEFINE_DEFAULT_LOAD_CACHE_VEC__(prefetch) + + + + + + + + + + + + + + +template> +__device__ static inline void store_async(const GL &dst, const SV &src, const COORD &idx) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + store_commit_group(); +} +__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_async) + +template> +__device__ static inline void store_async(const PGL &dst, const SV &src, const COORD &idx) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + store_commit_group(); +} +__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_async) + + + + + + + + + + + + +template> +__device__ static inline void store_add_async(const GL &dst, const SV &src, const COORD &idx) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_add_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + store_commit_group(); +} +__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_add_async) + +template> +__device__ static inline void store_add_async(const PGL &dst, const SV &src, const COORD &idx) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_add_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + store_commit_group(); +} +__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_add_async) + + + + + + + + + + + + +template> +__device__ static inline void store_min_async(const GL &dst, const SV &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_min_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + store_commit_group(); +} +__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_min_async) + +template> +__device__ static inline void store_min_async(const PGL &dst, const SV &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_min_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + store_commit_group(); +} +__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_min_async) + + + + + + + + + + + +template> +__device__ static inline void store_max_async(const GL &dst, const SV &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_max_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + store_commit_group(); +} +__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_max_async) + +template> +__device__ static inline void store_max_async(const PGL &dst, const SV &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_max_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + store_commit_group(); +} +__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_max_async) + + + + + + + + + + + + +template> +__device__ static inline void load_async(SV &dst, const GL &src, const COORD &idx, semaphore& bar) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(src.template get_tma()); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(&bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst)); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t dst_i_ptr = dst_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_load_async_tma_internal(tma_ptr, dst_i_ptr, mbar_ptr, tma_coord); + } +} +__KITTENS_TMA_DEFINE_SEMAPHORE_CACHE_VEC__(load_async) diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/tma_cluster.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/tma_cluster.cuh new file mode 100644 index 0000000..3f45de7 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/tma_cluster.cuh @@ -0,0 +1,31 @@ + + + + + + + + + + + + + + + + + + template> + __device__ static inline void load_async(SV &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(src.template get_tma()); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(&bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst)); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t dst_i_ptr = dst_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::cluster::vec_load_async_tma_internal(tma_ptr, dst_i_ptr, mbar_ptr, tma_coord, cluster_mask, dst_mbar_cta); + } + } + __KITTENS_TMA_DEFINE_CLUSTER_SEMAPHORE_CACHE_VEC__(load_async) diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/vec.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/vec.cuh new file mode 100644 index 0000000..bf995b5 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/vec.cuh @@ -0,0 +1,10 @@ + + + + + +#include "shared_to_register.cuh" +#include "global_to_register.cuh" +#include "global_to_shared.cuh" + + diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x112.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x112.impl new file mode 100644 index 0000000..bf205aa --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x112.impl @@ -0,0 +1,334 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %61, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n112k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55}, " \ + "{%56, %57, %58, %59}, " \ + "%60, " \ + "p, 1, %63, %62;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %61, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n112k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55}, " \ + "{%56, %57, %58, %59}, " \ + "%60, " \ + "p, 1, %63, %62;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %33, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n112k16.f16.f16.f16 " \ + "{%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, " \ + "p, 1, %35, %34;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %58, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n112k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55}, " \ + "%56, " \ + "%57, " \ + "p, 1, %61, %59, %60;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %58, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n112k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55}, " \ + "%56, " \ + "%57, " \ + "p, 1, %61, %59, %60;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %30, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n112k16.f16.f16.f16 " \ + "{%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, " \ + "p, 1, %33, %31, %32;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x128.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x128.impl new file mode 100644 index 0000000..5690b2c --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x128.impl @@ -0,0 +1,1029 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %69, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "{%64, %65, %66, %67}, " \ + "%68, " \ + "p, 1, %71, %70;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y), + "+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y), + "+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y), + "+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y), + "+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %69, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "{%64, %65, %66, %67}, " \ + "%68, " \ + "p, 1, %71, %70;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y), + "+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y), + "+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y), + "+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y), + "+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %37, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "{%32, %33, %34, %35}, " \ + "%36, " \ + "p, 1, %39, %38;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %69, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "{%64, %65, %66, %67}, " \ + "%68, " \ + "p, 1, %70;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y), + "+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y), + "+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y), + "+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y), + "+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %69, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "{%64, %65, %66, %67}, " \ + "%68, " \ + "p, 1, %70;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y), + "+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y), + "+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y), + "+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y), + "+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %37, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.f16.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "{%32, %33, %34, %35}, " \ + "%36, " \ + "p, 1, %38;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %37, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.f16.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "{%32, %33, %34, %35}, " \ + "%36, " \ + "p, 1, %38;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %69, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "{%64, %65, %66, %67}, " \ + "%68, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %69, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "{%64, %65, %66, %67}, " \ + "%68, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %66, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "%64, " \ + "%65, " \ + "p, 1, %69, %67, %68;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y), + "+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y), + "+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y), + "+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y), + "+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %66, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "%64, " \ + "%65, " \ + "p, 1, %69, %67, %68;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y), + "+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y), + "+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y), + "+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y), + "+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %34, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "%32, " \ + "%33, " \ + "p, 1, %37, %35, %36;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %66, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "%64, " \ + "%65, " \ + "p, 1, %67;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y), + "+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y), + "+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y), + "+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y), + "+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %66, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "%64, " \ + "%65, " \ + "p, 1, %67;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y), + "+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y), + "+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y), + "+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y), + "+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %34, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.f16.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "%32, " \ + "%33, " \ + "p, 1, %35;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %34, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.f16.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "%32, " \ + "%33, " \ + "p, 1, %35;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %66, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "%64, " \ + "%65, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %66, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "%64, " \ + "%65, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x144.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x144.impl new file mode 100644 index 0000000..7f6756a --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x144.impl @@ -0,0 +1,382 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %77, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n144k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71}, " \ + "{%72, %73, %74, %75}, " \ + "%76, " \ + "p, 1, %79, %78;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %77, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n144k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71}, " \ + "{%72, %73, %74, %75}, " \ + "%76, " \ + "p, 1, %79, %78;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %41, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n144k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35}, " \ + "{%36, %37, %38, %39}, " \ + "%40, " \ + "p, 1, %43, %42;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %74, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n144k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71}, " \ + "%72, " \ + "%73, " \ + "p, 1, %77, %75, %76;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %74, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n144k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71}, " \ + "%72, " \ + "%73, " \ + "p, 1, %77, %75, %76;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %38, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n144k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35}, " \ + "%36, " \ + "%37, " \ + "p, 1, %41, %39, %40;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x16.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x16.impl new file mode 100644 index 0000000..34cf64f --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x16.impl @@ -0,0 +1,190 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %13, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n16k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7}, " \ + "{%8, %9, %10, %11}, " \ + "%12, " \ + "p, 1, %15, %14;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %13, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7}, " \ + "{%8, %9, %10, %11}, " \ + "%12, " \ + "p, 1, %15, %14;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %9, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n16k16.f16.f16.f16 " \ + "{%0, %1, %2, %3}, " \ + "{%4, %5, %6, %7}, " \ + "%8, " \ + "p, 1, %11, %10;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %10, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n16k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7}, " \ + "%8, " \ + "%9, " \ + "p, 1, %13, %11, %12;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %10, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7}, " \ + "%8, " \ + "%9, " \ + "p, 1, %13, %11, %12;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %6, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n16k16.f16.f16.f16 " \ + "{%0, %1, %2, %3}, " \ + "%4, " \ + "%5, " \ + "p, 1, %9, %7, %8;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x160.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x160.impl new file mode 100644 index 0000000..cf9d4cd --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x160.impl @@ -0,0 +1,914 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %85, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "{%80, %81, %82, %83}, " \ + "%84, " \ + "p, 1, %87, %86;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %85, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "{%80, %81, %82, %83}, " \ + "%84, " \ + "p, 1, %87, %86;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %45, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39}, " \ + "{%40, %41, %42, %43}, " \ + "%44, " \ + "p, 1, %47, %46;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %85, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "{%80, %81, %82, %83}, " \ + "%84, " \ + "p, 1, %86;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %85, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "{%80, %81, %82, %83}, " \ + "%84, " \ + "p, 1, %86;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %85, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "{%80, %81, %82, %83}, " \ + "%84, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %85, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "{%80, %81, %82, %83}, " \ + "%84, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %82, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "%80, " \ + "%81, " \ + "p, 1, %85, %83, %84;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %82, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "%80, " \ + "%81, " \ + "p, 1, %85, %83, %84;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %42, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39}, " \ + "%40, " \ + "%41, " \ + "p, 1, %45, %43, %44;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %82, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "%80, " \ + "%81, " \ + "p, 1, %83;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %82, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "%80, " \ + "%81, " \ + "p, 1, %83;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %82, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "%80, " \ + "%81, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %82, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "%80, " \ + "%81, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x176.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x176.impl new file mode 100644 index 0000000..79fadbd --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x176.impl @@ -0,0 +1,430 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %93, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n176k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87}, " \ + "{%88, %89, %90, %91}, " \ + "%92, " \ + "p, 1, %95, %94;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %93, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n176k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87}, " \ + "{%88, %89, %90, %91}, " \ + "%92, " \ + "p, 1, %95, %94;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %49, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n176k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43}, " \ + "{%44, %45, %46, %47}, " \ + "%48, " \ + "p, 1, %51, %50;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %90, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n176k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87}, " \ + "%88, " \ + "%89, " \ + "p, 1, %93, %91, %92;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %90, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n176k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87}, " \ + "%88, " \ + "%89, " \ + "p, 1, %93, %91, %92;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %46, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n176k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43}, " \ + "%44, " \ + "%45, " \ + "p, 1, %49, %47, %48;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x192.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x192.impl new file mode 100644 index 0000000..d6c7068 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x192.impl @@ -0,0 +1,954 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %101, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "{%96, %97, %98, %99}, " \ + "%100, " \ + "p, 1, %103, %102;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %101, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "{%96, %97, %98, %99}, " \ + "%100, " \ + "p, 1, %103, %102;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %53, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "{%48, %49, %50, %51}, " \ + "%52, " \ + "p, 1, %55, %54;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %101, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "{%96, %97, %98, %99}, " \ + "%100, " \ + "p, 1, %102;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %101, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "{%96, %97, %98, %99}, " \ + "%100, " \ + "p, 1, %102;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %101, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "{%96, %97, %98, %99}, " \ + "%100, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %101, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "{%96, %97, %98, %99}, " \ + "%100, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %98, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "%96, " \ + "%97, " \ + "p, 1, %101, %99, %100;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %98, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "%96, " \ + "%97, " \ + "p, 1, %101, %99, %100;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %50, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "%48, " \ + "%49, " \ + "p, 1, %53, %51, %52;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %98, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "%96, " \ + "%97, " \ + "p, 1, %99;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %98, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "%96, " \ + "%97, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %98, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "%96, " \ + "%97, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x208.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x208.impl new file mode 100644 index 0000000..06caa1e --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x208.impl @@ -0,0 +1,478 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %109, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n208k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103}, " \ + "{%104, %105, %106, %107}, " \ + "%108, " \ + "p, 1, %111, %110;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %109, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n208k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103}, " \ + "{%104, %105, %106, %107}, " \ + "%108, " \ + "p, 1, %111, %110;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %57, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n208k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51}, " \ + "{%52, %53, %54, %55}, " \ + "%56, " \ + "p, 1, %59, %58;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %106, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n208k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103}, " \ + "%104, " \ + "%105, " \ + "p, 1, %109, %107, %108;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %106, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n208k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103}, " \ + "%104, " \ + "%105, " \ + "p, 1, %109, %107, %108;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %54, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n208k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51}, " \ + "%52, " \ + "%53, " \ + "p, 1, %57, %55, %56;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x224.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x224.impl new file mode 100644 index 0000000..dc342d1 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x224.impl @@ -0,0 +1,1138 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %117, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "{%112, %113, %114, %115}, " \ + "%116, " \ + "p, 1, %119, %118;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %117, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "{%112, %113, %114, %115}, " \ + "%116, " \ + "p, 1, %119, %118;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %61, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55}, " \ + "{%56, %57, %58, %59}, " \ + "%60, " \ + "p, 1, %63, %62;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %117, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "{%112, %113, %114, %115}, " \ + "%116, " \ + "p, 1, %118;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %117, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "{%112, %113, %114, %115}, " \ + "%116, " \ + "p, 1, %118;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %117, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "{%112, %113, %114, %115}, " \ + "%116, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y), + "+r"(dst.tiles[0][12].data[0].x), "+r"(dst.tiles[0][12].data[0].y), + "+r"(dst.tiles[0][12].data[1].x), "+r"(dst.tiles[0][12].data[1].y), + "+r"(dst.tiles[0][12].data[2].x), "+r"(dst.tiles[0][12].data[2].y), + "+r"(dst.tiles[0][12].data[3].x), "+r"(dst.tiles[0][12].data[3].y), + "+r"(dst.tiles[0][13].data[0].x), "+r"(dst.tiles[0][13].data[0].y), + "+r"(dst.tiles[0][13].data[1].x), "+r"(dst.tiles[0][13].data[1].y), + "+r"(dst.tiles[0][13].data[2].x), "+r"(dst.tiles[0][13].data[2].y), + "+r"(dst.tiles[0][13].data[3].x), "+r"(dst.tiles[0][13].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %117, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "{%112, %113, %114, %115}, " \ + "%116, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y), + "+r"(dst.tiles[0][12].data[0].x), "+r"(dst.tiles[0][12].data[0].y), + "+r"(dst.tiles[0][12].data[1].x), "+r"(dst.tiles[0][12].data[1].y), + "+r"(dst.tiles[0][12].data[2].x), "+r"(dst.tiles[0][12].data[2].y), + "+r"(dst.tiles[0][12].data[3].x), "+r"(dst.tiles[0][12].data[3].y), + "+r"(dst.tiles[0][13].data[0].x), "+r"(dst.tiles[0][13].data[0].y), + "+r"(dst.tiles[0][13].data[1].x), "+r"(dst.tiles[0][13].data[1].y), + "+r"(dst.tiles[0][13].data[2].x), "+r"(dst.tiles[0][13].data[2].y), + "+r"(dst.tiles[0][13].data[3].x), "+r"(dst.tiles[0][13].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %114, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "%112, " \ + "%113, " \ + "p, 1, %117, %115, %116;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %114, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "%112, " \ + "%113, " \ + "p, 1, %117, %115, %116;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %58, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55}, " \ + "%56, " \ + "%57, " \ + "p, 1, %61, %59, %60;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %114, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "%112, " \ + "%113, " \ + "p, 1, %117, %115, %116;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %114, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "%112, " \ + "%113, " \ + "p, 1, %117, %115, %116;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %114, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "%112, " \ + "%113, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y), + "+r"(dst.tiles[0][12].data[0].x), "+r"(dst.tiles[0][12].data[0].y), + "+r"(dst.tiles[0][12].data[1].x), "+r"(dst.tiles[0][12].data[1].y), + "+r"(dst.tiles[0][12].data[2].x), "+r"(dst.tiles[0][12].data[2].y), + "+r"(dst.tiles[0][12].data[3].x), "+r"(dst.tiles[0][12].data[3].y), + "+r"(dst.tiles[0][13].data[0].x), "+r"(dst.tiles[0][13].data[0].y), + "+r"(dst.tiles[0][13].data[1].x), "+r"(dst.tiles[0][13].data[1].y), + "+r"(dst.tiles[0][13].data[2].x), "+r"(dst.tiles[0][13].data[2].y), + "+r"(dst.tiles[0][13].data[3].x), "+r"(dst.tiles[0][13].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %114, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "%112, " \ + "%113, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y), + "+r"(dst.tiles[0][12].data[0].x), "+r"(dst.tiles[0][12].data[0].y), + "+r"(dst.tiles[0][12].data[1].x), "+r"(dst.tiles[0][12].data[1].y), + "+r"(dst.tiles[0][12].data[2].x), "+r"(dst.tiles[0][12].data[2].y), + "+r"(dst.tiles[0][12].data[3].x), "+r"(dst.tiles[0][12].data[3].y), + "+r"(dst.tiles[0][13].data[0].x), "+r"(dst.tiles[0][13].data[0].y), + "+r"(dst.tiles[0][13].data[1].x), "+r"(dst.tiles[0][13].data[1].y), + "+r"(dst.tiles[0][13].data[2].x), "+r"(dst.tiles[0][13].data[2].y), + "+r"(dst.tiles[0][13].data[3].x), "+r"(dst.tiles[0][13].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x240.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x240.impl new file mode 100644 index 0000000..aaf5336 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x240.impl @@ -0,0 +1,526 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %125, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n240k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119}, " \ + "{%120, %121, %122, %123}, " \ + "%124, " \ + "p, 1, %127, %126;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %125, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n240k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119}, " \ + "{%120, %121, %122, %123}, " \ + "%124, " \ + "p, 1, %127, %126;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %65, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n240k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59}, " \ + "{%60, %61, %62, %63}, " \ + "%64, " \ + "p, 1, %67, %66;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %122, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n240k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119}, " \ + "%120, " \ + "%121, " \ + "p, 1, %125, %123, %124;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %122, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n240k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119}, " \ + "%120, " \ + "%121, " \ + "p, 1, %125, %123, %124;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %62, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n240k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59}, " \ + "%60, " \ + "%61, " \ + "p, 1, %65, %63, %64;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x256.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x256.impl new file mode 100644 index 0000000..9c2e06e --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x256.impl @@ -0,0 +1,1604 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %133, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "{%128, %129, %130, %131}, " \ + "%132, " \ + "p, 1, %135, %134;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y), + "+f"(dst.tiles[0][15].data[0].x), "+f"(dst.tiles[0][15].data[0].y), + "+f"(dst.tiles[0][15].data[1].x), "+f"(dst.tiles[0][15].data[1].y), + "+f"(dst.tiles[0][15].data[2].x), "+f"(dst.tiles[0][15].data[2].y), + "+f"(dst.tiles[0][15].data[3].x), "+f"(dst.tiles[0][15].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %133, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "{%128, %129, %130, %131}, " \ + "%132, " \ + "p, 1, %135, %134;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y), + "+f"(dst.tiles[0][15].data[0].x), "+f"(dst.tiles[0][15].data[0].y), + "+f"(dst.tiles[0][15].data[1].x), "+f"(dst.tiles[0][15].data[1].y), + "+f"(dst.tiles[0][15].data[2].x), "+f"(dst.tiles[0][15].data[2].y), + "+f"(dst.tiles[0][15].data[3].x), "+f"(dst.tiles[0][15].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %69, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "{%64, %65, %66, %67}, " \ + "%68, " \ + "p, 1, %71, %70;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %133, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "{%128, %129, %130, %131}, " \ + "%132, " \ + "p, 1, %134;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y), + "+f"(dst.tiles[0][15].data[0].x), "+f"(dst.tiles[0][15].data[0].y), + "+f"(dst.tiles[0][15].data[1].x), "+f"(dst.tiles[0][15].data[1].y), + "+f"(dst.tiles[0][15].data[2].x), "+f"(dst.tiles[0][15].data[2].y), + "+f"(dst.tiles[0][15].data[3].x), "+f"(dst.tiles[0][15].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %133, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "{%128, %129, %130, %131}, " \ + "%132, " \ + "p, 1, %134;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y), + "+f"(dst.tiles[0][15].data[0].x), "+f"(dst.tiles[0][15].data[0].y), + "+f"(dst.tiles[0][15].data[1].x), "+f"(dst.tiles[0][15].data[1].y), + "+f"(dst.tiles[0][15].data[2].x), "+f"(dst.tiles[0][15].data[2].y), + "+f"(dst.tiles[0][15].data[3].x), "+f"(dst.tiles[0][15].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %69, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.f16.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "{%64, %65, %66, %67}, " \ + "%68, " \ + "p, 1, %70;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %69, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.f16.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "{%64, %65, %66, %67}, " \ + "%68, " \ + "p, 1, %70;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %133, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "{%128, %129, %130, %131}, " \ + "%132, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y), + "+r"(dst.tiles[0][12].data[0].x), "+r"(dst.tiles[0][12].data[0].y), + "+r"(dst.tiles[0][12].data[1].x), "+r"(dst.tiles[0][12].data[1].y), + "+r"(dst.tiles[0][12].data[2].x), "+r"(dst.tiles[0][12].data[2].y), + "+r"(dst.tiles[0][12].data[3].x), "+r"(dst.tiles[0][12].data[3].y), + "+r"(dst.tiles[0][13].data[0].x), "+r"(dst.tiles[0][13].data[0].y), + "+r"(dst.tiles[0][13].data[1].x), "+r"(dst.tiles[0][13].data[1].y), + "+r"(dst.tiles[0][13].data[2].x), "+r"(dst.tiles[0][13].data[2].y), + "+r"(dst.tiles[0][13].data[3].x), "+r"(dst.tiles[0][13].data[3].y), + "+r"(dst.tiles[0][14].data[0].x), "+r"(dst.tiles[0][14].data[0].y), + "+r"(dst.tiles[0][14].data[1].x), "+r"(dst.tiles[0][14].data[1].y), + "+r"(dst.tiles[0][14].data[2].x), "+r"(dst.tiles[0][14].data[2].y), + "+r"(dst.tiles[0][14].data[3].x), "+r"(dst.tiles[0][14].data[3].y), + "+r"(dst.tiles[0][15].data[0].x), "+r"(dst.tiles[0][15].data[0].y), + "+r"(dst.tiles[0][15].data[1].x), "+r"(dst.tiles[0][15].data[1].y), + "+r"(dst.tiles[0][15].data[2].x), "+r"(dst.tiles[0][15].data[2].y), + "+r"(dst.tiles[0][15].data[3].x), "+r"(dst.tiles[0][15].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %133, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "{%128, %129, %130, %131}, " \ + "%132, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y), + "+r"(dst.tiles[0][12].data[0].x), "+r"(dst.tiles[0][12].data[0].y), + "+r"(dst.tiles[0][12].data[1].x), "+r"(dst.tiles[0][12].data[1].y), + "+r"(dst.tiles[0][12].data[2].x), "+r"(dst.tiles[0][12].data[2].y), + "+r"(dst.tiles[0][12].data[3].x), "+r"(dst.tiles[0][12].data[3].y), + "+r"(dst.tiles[0][13].data[0].x), "+r"(dst.tiles[0][13].data[0].y), + "+r"(dst.tiles[0][13].data[1].x), "+r"(dst.tiles[0][13].data[1].y), + "+r"(dst.tiles[0][13].data[2].x), "+r"(dst.tiles[0][13].data[2].y), + "+r"(dst.tiles[0][13].data[3].x), "+r"(dst.tiles[0][13].data[3].y), + "+r"(dst.tiles[0][14].data[0].x), "+r"(dst.tiles[0][14].data[0].y), + "+r"(dst.tiles[0][14].data[1].x), "+r"(dst.tiles[0][14].data[1].y), + "+r"(dst.tiles[0][14].data[2].x), "+r"(dst.tiles[0][14].data[2].y), + "+r"(dst.tiles[0][14].data[3].x), "+r"(dst.tiles[0][14].data[3].y), + "+r"(dst.tiles[0][15].data[0].x), "+r"(dst.tiles[0][15].data[0].y), + "+r"(dst.tiles[0][15].data[1].x), "+r"(dst.tiles[0][15].data[1].y), + "+r"(dst.tiles[0][15].data[2].x), "+r"(dst.tiles[0][15].data[2].y), + "+r"(dst.tiles[0][15].data[3].x), "+r"(dst.tiles[0][15].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %130, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "%128, " \ + "%129, " \ + "p, 1, %133, %131, %132;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y), + "+f"(dst.tiles[0][15].data[0].x), "+f"(dst.tiles[0][15].data[0].y), + "+f"(dst.tiles[0][15].data[1].x), "+f"(dst.tiles[0][15].data[1].y), + "+f"(dst.tiles[0][15].data[2].x), "+f"(dst.tiles[0][15].data[2].y), + "+f"(dst.tiles[0][15].data[3].x), "+f"(dst.tiles[0][15].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %130, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "%128, " \ + "%129, " \ + "p, 1, %133, %131, %132;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y), + "+f"(dst.tiles[0][15].data[0].x), "+f"(dst.tiles[0][15].data[0].y), + "+f"(dst.tiles[0][15].data[1].x), "+f"(dst.tiles[0][15].data[1].y), + "+f"(dst.tiles[0][15].data[2].x), "+f"(dst.tiles[0][15].data[2].y), + "+f"(dst.tiles[0][15].data[3].x), "+f"(dst.tiles[0][15].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %66, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "%64, " \ + "%65, " \ + "p, 1, %69, %67, %68;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %130, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "%128, " \ + "%129, " \ + "p, 1, %131;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y), + "+f"(dst.tiles[0][15].data[0].x), "+f"(dst.tiles[0][15].data[0].y), + "+f"(dst.tiles[0][15].data[1].x), "+f"(dst.tiles[0][15].data[1].y), + "+f"(dst.tiles[0][15].data[2].x), "+f"(dst.tiles[0][15].data[2].y), + "+f"(dst.tiles[0][15].data[3].x), "+f"(dst.tiles[0][15].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %130, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.f32.e5m.e5m " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "%128, " \ + "%129, " \ + "p, 1, %131;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y), + "+f"(dst.tiles[0][15].data[0].x), "+f"(dst.tiles[0][15].data[0].y), + "+f"(dst.tiles[0][15].data[1].x), "+f"(dst.tiles[0][15].data[1].y), + "+f"(dst.tiles[0][15].data[2].x), "+f"(dst.tiles[0][15].data[2].y), + "+f"(dst.tiles[0][15].data[3].x), "+f"(dst.tiles[0][15].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %66, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.f16.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "%64, " \ + "%65, " \ + "p, 1, %67;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %66, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.f16.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "%64, " \ + "%65, " \ + "p, 1, %67;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %130, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "%128, " \ + "%129, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y), + "+r"(dst.tiles[0][12].data[0].x), "+r"(dst.tiles[0][12].data[0].y), + "+r"(dst.tiles[0][12].data[1].x), "+r"(dst.tiles[0][12].data[1].y), + "+r"(dst.tiles[0][12].data[2].x), "+r"(dst.tiles[0][12].data[2].y), + "+r"(dst.tiles[0][12].data[3].x), "+r"(dst.tiles[0][12].data[3].y), + "+r"(dst.tiles[0][13].data[0].x), "+r"(dst.tiles[0][13].data[0].y), + "+r"(dst.tiles[0][13].data[1].x), "+r"(dst.tiles[0][13].data[1].y), + "+r"(dst.tiles[0][13].data[2].x), "+r"(dst.tiles[0][13].data[2].y), + "+r"(dst.tiles[0][13].data[3].x), "+r"(dst.tiles[0][13].data[3].y), + "+r"(dst.tiles[0][14].data[0].x), "+r"(dst.tiles[0][14].data[0].y), + "+r"(dst.tiles[0][14].data[1].x), "+r"(dst.tiles[0][14].data[1].y), + "+r"(dst.tiles[0][14].data[2].x), "+r"(dst.tiles[0][14].data[2].y), + "+r"(dst.tiles[0][14].data[3].x), "+r"(dst.tiles[0][14].data[3].y), + "+r"(dst.tiles[0][15].data[0].x), "+r"(dst.tiles[0][15].data[0].y), + "+r"(dst.tiles[0][15].data[1].x), "+r"(dst.tiles[0][15].data[1].y), + "+r"(dst.tiles[0][15].data[2].x), "+r"(dst.tiles[0][15].data[2].y), + "+r"(dst.tiles[0][15].data[3].x), "+r"(dst.tiles[0][15].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %130, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "%128, " \ + "%129, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y), + "+r"(dst.tiles[0][12].data[0].x), "+r"(dst.tiles[0][12].data[0].y), + "+r"(dst.tiles[0][12].data[1].x), "+r"(dst.tiles[0][12].data[1].y), + "+r"(dst.tiles[0][12].data[2].x), "+r"(dst.tiles[0][12].data[2].y), + "+r"(dst.tiles[0][12].data[3].x), "+r"(dst.tiles[0][12].data[3].y), + "+r"(dst.tiles[0][13].data[0].x), "+r"(dst.tiles[0][13].data[0].y), + "+r"(dst.tiles[0][13].data[1].x), "+r"(dst.tiles[0][13].data[1].y), + "+r"(dst.tiles[0][13].data[2].x), "+r"(dst.tiles[0][13].data[2].y), + "+r"(dst.tiles[0][13].data[3].x), "+r"(dst.tiles[0][13].data[3].y), + "+r"(dst.tiles[0][14].data[0].x), "+r"(dst.tiles[0][14].data[0].y), + "+r"(dst.tiles[0][14].data[1].x), "+r"(dst.tiles[0][14].data[1].y), + "+r"(dst.tiles[0][14].data[2].x), "+r"(dst.tiles[0][14].data[2].y), + "+r"(dst.tiles[0][14].data[3].x), "+r"(dst.tiles[0][14].data[3].y), + "+r"(dst.tiles[0][15].data[0].x), "+r"(dst.tiles[0][15].data[0].y), + "+r"(dst.tiles[0][15].data[1].x), "+r"(dst.tiles[0][15].data[1].y), + "+r"(dst.tiles[0][15].data[2].x), "+r"(dst.tiles[0][15].data[2].y), + "+r"(dst.tiles[0][15].data[3].x), "+r"(dst.tiles[0][15].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x32.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x32.impl new file mode 100644 index 0000000..205773f --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x32.impl @@ -0,0 +1,564 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %21, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "{%16, %17, %18, %19}, " \ + "%20, " \ + "p, 1, %23, %22;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %21, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "{%16, %17, %18, %19}, " \ + "%20, " \ + "p, 1, %23, %22;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %13, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7}, " \ + "{%8, %9, %10, %11}, " \ + "%12, " \ + "p, 1, %15, %14;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %21, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "{%16, %17, %18, %19}, " \ + "%20, " \ + "p, 1, %22;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %21, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "{%16, %17, %18, %19}, " \ + "%20, " \ + "p, 1, %22;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %13, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.f16.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7}, " \ + "{%8, %9, %10, %11}, " \ + "%12, " \ + "p, 1, %14;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %21, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "{%16, %17, %18, %19}, " \ + "%20, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %21, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "{%16, %17, %18, %19}, " \ + "%20, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %18, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "%16, " \ + "%17, " \ + "p, 1, %21, %19, %20;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %18, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "%16, " \ + "%17, " \ + "p, 1, %21, %19, %20;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %10, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7}, " \ + "%8, " \ + "%9, " \ + "p, 1, %13, %11, %12;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %18, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "%16, " \ + "%17, " \ + "p, 1, %19;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %18, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "%16, " \ + "%17, " \ + "p, 1, %19;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %10, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.f16.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7}, " \ + "%8, " \ + "%9, " \ + "p, 1, %11;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %10, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.f16.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7}, " \ + "%8, " \ + "%9, " \ + "p, 1, %11;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %18, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "%16, " \ + "%17, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %18, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "%16, " \ + "%17, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x48.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x48.impl new file mode 100644 index 0000000..702ee70 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x48.impl @@ -0,0 +1,238 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %29, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n48k16.f32.bf16.bf16 " \ + "{%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, " \ + "p, 1, %31, %30;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %29, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n48k16.f32.f16.f16 " \ + "{%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, " \ + "p, 1, %31, %30;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %17, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n48k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11}, " \ + "{%12, %13, %14, %15}, " \ + "%16, " \ + "p, 1, %19, %18;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %26, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n48k16.f32.bf16.bf16 " \ + "{%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, " \ + "p, 1, %29, %27, %28;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %26, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n48k16.f32.f16.f16 " \ + "{%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, " \ + "p, 1, %29, %27, %28;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %14, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n48k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11}, " \ + "%12, " \ + "%13, " \ + "p, 1, %17, %15, %16;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x64.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x64.impl new file mode 100644 index 0000000..ff7983d --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x64.impl @@ -0,0 +1,739 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %37, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "{%32, %33, %34, %35}, " \ + "%36, " \ + "p, 1, %39, %38;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %37, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "{%32, %33, %34, %35}, " \ + "%36, " \ + "p, 1, %39, %38;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %21, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "{%16, %17, %18, %19}, " \ + "%20, " \ + "p, 1, %23, %22;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %37, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "{%32, %33, %34, %35}, " \ + "%36, " \ + "p, 1, %38;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %37, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "{%32, %33, %34, %35}, " \ + "%36, " \ + "p, 1, %38;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %21, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.f16.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "{%16, %17, %18, %19}, " \ + "%20, " \ + "p, 1, %22;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %21, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.f16.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "{%16, %17, %18, %19}, " \ + "%20, " \ + "p, 1, %22;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %37, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "{%32, %33, %34, %35}, " \ + "%36, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %37, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "{%32, %33, %34, %35}, " \ + "%36, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %34, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "%32, " \ + "%33, " \ + "p, 1, %37, %35, %36;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %34, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "%32, " \ + "%33, " \ + "p, 1, %37, %35, %36;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %18, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "%16, " \ + "%17, " \ + "p, 1, %21, %19, %20;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %34, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "%32, " \ + "%33, " \ + "p, 1, %35;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %34, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "%32, " \ + "%33, " \ + "p, 1, %35;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %18, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.f16.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "%16, " \ + "%17, " \ + "p, 1, %19;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %18, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.f16.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "%16, " \ + "%17, " \ + "p, 1, %19;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %34, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "%32, " \ + "%33, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %34, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "%32, " \ + "%33, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x80.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x80.impl new file mode 100644 index 0000000..7a33667 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x80.impl @@ -0,0 +1,286 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %45, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n80k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39}, " \ + "{%40, %41, %42, %43}, " \ + "%44, " \ + "p, 1, %47, %46;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %45, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n80k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39}, " \ + "{%40, %41, %42, %43}, " \ + "%44, " \ + "p, 1, %47, %46;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %25, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n80k16.f16.f16.f16 " \ + "{%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, " \ + "p, 1, %27, %26;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %42, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n80k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39}, " \ + "%40, " \ + "%41, " \ + "p, 1, %45, %43, %44;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %42, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n80k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39}, " \ + "%40, " \ + "%41, " \ + "p, 1, %45, %43, %44;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %22, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n80k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19}, " \ + "%20, " \ + "%21, " \ + "p, 1, %25, %23, %24;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x96.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x96.impl new file mode 100644 index 0000000..b7f955f --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x96.impl @@ -0,0 +1,887 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %53, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "{%48, %49, %50, %51}, " \ + "%52, " \ + "p, 1, %55, %54;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %53, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "{%48, %49, %50, %51}, " \ + "%52, " \ + "p, 1, %55, %54;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %29, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k16.f16.f16.f16 " \ + "{%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, " \ + "p, 1, %31, %30;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %53, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "{%48, %49, %50, %51}, " \ + "%52, " \ + "p, 1, %54;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %53, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "{%48, %49, %50, %51}, " \ + "%52, " \ + "p, 1, %54;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %29, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.f16.e4m3.e4m3 " \ + "{%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, " \ + "p, 1, %30;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %29, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.f16.e5m2.e5m2 " \ + "{%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, " \ + "p, 1, %30;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %53, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "{%48, %49, %50, %51}, " \ + "%52, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %53, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "{%48, %49, %50, %51}, " \ + "%52, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %50, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "%48, " \ + "%49, " \ + "p, 1, %53, %51, %52;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %50, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "%48, " \ + "%49, " \ + "p, 1, %53, %51, %52;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %26, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k16.f16.f16.f16 " \ + "{%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, " \ + "p, 1, %29, %27, %28;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %50, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "%48, " \ + "%49, " \ + "p, 1, %51;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %50, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "%48, " \ + "%49, " \ + "p, 1, %51;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %26, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.f16.e4m3.e4m3 " \ + "{%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, " \ + "p, 1, %27;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %26, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.f16.e5m2.e5m2 " \ + "{%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, " \ + "p, 1, %27;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %50, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "%48, " \ + "%49, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %50, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "%48, " \ + "%49, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/base.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/base.cuh new file mode 100644 index 0000000..1da1375 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/base.cuh @@ -0,0 +1,47 @@ +#pragma once + +#include "../../../../common/common.cuh" +#include "../../../../types/types.cuh" + +namespace kittens { +namespace detail { +namespace wgmma { + + +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ); + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ); +}; + + +#include "64x16.impl" +#include "64x32.impl" +#include "64x48.impl" +#include "64x64.impl" +#include "64x80.impl" +#include "64x96.impl" +#include "64x112.impl" +#include "64x128.impl" +#include "64x144.impl" +#include "64x160.impl" +#include "64x176.impl" +#include "64x192.impl" +#include "64x208.impl" +#include "64x224.impl" +#include "64x240.impl" +#include "64x256.impl" + +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/mma.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/mma.cuh new file mode 100644 index 0000000..2010402 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/mma.cuh @@ -0,0 +1,17 @@ + + + + + + +#include "warp.cuh" + + +#if defined(KITTENS_SM90) +#include "warpgroup.cuh" +#endif + + +#ifdef KITTENS_SM10X +#include "tcgen05.cuh" +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/tcgen05.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/tcgen05.cuh new file mode 100644 index 0000000..0c5db67 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/tcgen05.cuh @@ -0,0 +1,236 @@ + + + + +template +__device__ static inline void mma(D &d, const A &a, const B &b, semaphore &sem) { + if(laneid() == 0) ::kittens::mma(d, a, b, sem); +} +template +__device__ static inline void mma2(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm2(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mma(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + if(laneid() == 0) ::kittens::mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mma2(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mm(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mm2(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma2(d, a, b, sa, sb, sem); +} + +template +__device__ static inline void mma_AB(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mma2_AB(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mma_ABt(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mma2_ABt(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mma_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mma2_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma2(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mma_AtB(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mma2_AtB(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mma_AtBt(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mma2_AtBt(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} + +template +__device__ static inline void mm_AB(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm2_AB(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mm_ABt(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm2_ABt(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mm_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mm2_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma2(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mm_AtB(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm2_AtB(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mm_AtBt(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm2_AtBt(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} + + + + +template +__device__ static inline void mma(D &d, const A &a, const B &b) { + if(laneid() == 0) ::kittens::mma(d, a, b); +} +template +__device__ static inline void mma2(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm2(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mma(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + if(laneid() == 0) ::kittens::mma(d, a, b, sa, sb); +} +template +__device__ static inline void mma2(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma(d, a, b, sa, sb); +} +template +__device__ static inline void mm(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma(d, a, b, sa, sb); +} +template +__device__ static inline void mm2(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma2(d, a, b, sa, sb); +} + +template +__device__ static inline void mma_AB(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mma2_AB(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mma_ABt(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mma2_ABt(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mma_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma(d, a, b, sa, sb); +} +template +__device__ static inline void mma2_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma2(d, a, b, sa, sb); +} +template +__device__ static inline void mma_AtB(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mma2_AtB(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mma_AtBt(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mma2_AtBt(D &d, const A &a, const B &b) { + mma2(d, a, b); +} + +template +__device__ static inline void mm_AB(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm2_AB(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mm_ABt(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm2_ABt(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mm_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma(d, a, b, sa, sb); +} +template +__device__ static inline void mm2_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma2(d, a, b, sa, sb); +} +template +__device__ static inline void mm_AtB(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm2_AtB(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mm_AtBt(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm2_AtBt(D &d, const A &a, const B &b) { + mma2(d, a, b); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/warp.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/warp.cuh new file mode 100644 index 0000000..a94c57f --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/warp.cuh @@ -0,0 +1,1057 @@ + + + + + + + + + + + + + + + + + + + + + + +__device__ static inline void hmma16816( float2 &d0, float2 &d1, + const bf16_2 &a0, const bf16_2 &a1, const bf16_2 &a2, const bf16_2 &a3, + const bf16_2 &b0, const bf16_2 &b1, + const float2 &c0, const float2 &c1 ) { + asm volatile( + + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " \ + "{%0, %1, %2, %3}, " \ + "{%4, %5, %6, %7}, " \ + "{%8, %9}, " \ + "{%10, %11, %12, %13};" + + + : "+f"(d0.x), "+f"(d0.y), + "+f"(d1.x), "+f"(d1.y) + + + : "r"(*(uint32_t*)(&a0)), "r"(*(uint32_t*)(&a1)), + "r"(*(uint32_t*)(&a2)), "r"(*(uint32_t*)(&a3)), + + + "r"(*(uint32_t*)(&b0)), "r"(*(uint32_t*)(&b1)), + + + "f"(c0.x), "f"(c0.y), + "f"(c1.x), "f"(c1.y) + ); +} + + + + + + + + + + + + + + + + + +__device__ static inline void hmma16816( float2 &d0, float2 &d1, + const half_2 &a0, const half_2 &a1, const half_2 &a2, const half_2 &a3, + const half_2 &b0, const half_2 &b1, + const float2 &c0, const float2 &c1 ) { + asm volatile( + + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " \ + "{%0, %1, %2, %3}, " \ + "{%4, %5, %6, %7}, " \ + "{%8, %9}, " \ + "{%10, %11, %12, %13};" + + + : "+f"(d0.x), "+f"(d0.y), + "+f"(d1.x), "+f"(d1.y) + + + : "r"(*(uint32_t*)(&a0)), "r"(*(uint32_t*)(&a1)), + "r"(*(uint32_t*)(&a2)), "r"(*(uint32_t*)(&a3)), + + + "r"(*(uint32_t*)(&b0)), "r"(*(uint32_t*)(&b1)), + + + "f"(c0.x), "f"(c0.y), + "f"(c1.x), "f"(c1.y) + ); +} + + + + + + + + + + + + + + + + + +__device__ static inline void hmma16816( half_2 &d0, half_2 &d1, + const half_2 &a0, const half_2 &a1, const half_2 &a2, const half_2 &a3, + const half_2 &b0, const half_2 &b1, + const half_2 &c0, const half_2 &c1 ) { + asm volatile( + + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " \ + "{%0, %1}, " \ + "{%2, %3, %4, %5}, " \ + "{%6, %7}, " \ + "{%8, %9};" + + + : "=r"(*(uint32_t*)(&d0)), "=r"(*(uint32_t*)(&d1)) + + + : "r"(*(uint32_t*)(&a0)), "r"(*(uint32_t*)(&a1)), + "r"(*(uint32_t*)(&a2)), "r"(*(uint32_t*)(&a3)), + + + "r"(*(uint32_t*)(&b0)), "r"(*(uint32_t*)(&b1)), + + + "r"(*(uint32_t*)(&c0)), "r"(*(uint32_t*)(&c1)) + ); +} + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + + + + + + + + + + + +__device__ static inline void hmma16816( float2 &d0, float2 &d1, + const fp8e4m3_4 &a0, const fp8e4m3_4 &a1, + const fp8e4m3_4 &a2, const fp8e4m3_4 &a3, + const fp8e4m3_4 &b0, const fp8e4m3_4 &b1, + const float2 &c0, const float2 &c1) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, %1, %2, %3}, " + "{%4, %5, %6, %7}, " + "{%8, %9}, " + "{%10, %11, %12, %13};" + + + : "+f"(d0.x), "+f"(d0.y), + "+f"(d1.x), "+f"(d1.y) + + + : "r"(*(uint32_t*)(&a0)), "r"(*(uint32_t*)(&a1)), + "r"(*(uint32_t*)(&a2)), "r"(*(uint32_t*)(&a3)), + + + "r"(*(uint32_t*)(&b0)), "r"(*(uint32_t*)(&b1)), + + + "f"(c0.x), "f"(c0.y), + "f"(c1.x), "f"(c1.y) + ); +} +#endif + + + + + + +__device__ static inline void imma16832(int2 &d0, int2 &d1, + const int8_4 &a0, const int8_4 &a1, const int8_4 &a2, const int8_4 &a3, + const int8_4 &b0, const int8_4 &b1, + const int2 &c0, const int2 &c1) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.satfinite.s32.s8.s8.s32 " + "{%0, %1, %2, %3}, " + "{%4, %5, %6, %7}, " + "{%8, %9}, " + "{%10, %11, %12, %13};" + : "=r"(d0.x), "=r"(d0.y), + "=r"(d1.x), "=r"(d1.y) + : "r"(*(uint32_t*)(&a0)), "r"(*(uint32_t*)(&a1)), + "r"(*(uint32_t*)(&a2)), "r"(*(uint32_t*)(&a3)), + "r"(*(uint32_t*)(&b0)), "r"(*(uint32_t*)(&b1)), + "r"(c0.x), "r"(c0.y), + "r"(c1.x), "r"(c1.y) + ); +} +__device__ static inline void imma16832(int2 &d0, int2 &d1, + const int8_4 &a0, const int8_4 &a1, const int8_4 &a2, const int8_4 &a3, + const uint8_4 &b0, const uint8_4 &b1, + const int2 &c0, const int2 &c1) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.satfinite.s32.s8.u8.s32 " + "{%0, %1, %2, %3}, " + "{%4, %5, %6, %7}, " + "{%8, %9}, " + "{%10, %11, %12, %13};" + : "=r"(d0.x), "=r"(d0.y), + "=r"(d1.x), "=r"(d1.y) + : "r"(*(uint32_t*)(&a0)), "r"(*(uint32_t*)(&a1)), + "r"(*(uint32_t*)(&a2)), "r"(*(uint32_t*)(&a3)), + "r"(*(uint32_t*)(&b0)), "r"(*(uint32_t*)(&b1)), + "r"(c0.x), "r"(c0.y), + "r"(c1.x), "r"(c1.y) + ); +} +__device__ static inline void imma16832(int2 &d0, int2 &d1, + const uint8_4 &a0, const uint8_4 &a1, const uint8_4 &a2, const uint8_4 &a3, + const int8_4 &b0, const int8_4 &b1, + const int2 &c0, const int2 &c1) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.satfinite.s32.u8.s8.s32 " + "{%0, %1, %2, %3}, " + "{%4, %5, %6, %7}, " + "{%8, %9}, " + "{%10, %11, %12, %13};" + : "=r"(d0.x), "=r"(d0.y), + "=r"(d1.x), "=r"(d1.y) + : "r"(*(uint32_t*)(&a0)), "r"(*(uint32_t*)(&a1)), + "r"(*(uint32_t*)(&a2)), "r"(*(uint32_t*)(&a3)), + "r"(*(uint32_t*)(&b0)), "r"(*(uint32_t*)(&b1)), + "r"(c0.x), "r"(c0.y), + "r"(c1.x), "r"(c1.y) + ); +} +__device__ static inline void imma16832(int2 &d0, int2 &d1, + const uint8_4 &a0, const uint8_4 &a1, const uint8_4 &a2, const uint8_4 &a3, + const uint8_4 &b0, const uint8_4 &b1, + const int2 &c0, const int2 &c1) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.satfinite.s32.u8.u8.s32 " + "{%0, %1, %2, %3}, " + "{%4, %5, %6, %7}, " + "{%8, %9}, " + "{%10, %11, %12, %13};" + : "=r"(d0.x), "=r"(d0.y), + "=r"(d1.x), "=r"(d1.y) + : "r"(*(uint32_t*)(&a0)), "r"(*(uint32_t*)(&a1)), + "r"(*(uint32_t*)(&a2)), "r"(*(uint32_t*)(&a3)), + "r"(*(uint32_t*)(&b0)), "r"(*(uint32_t*)(&b1)), + "r"(c0.x), "r"(c0.y), + "r"(c1.x), "r"(c1.y) + ); +} + + + + + + + + + + + + +__device__ static inline void mma_AB_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} + + + + + + + + + + + + +__device__ static inline void mma_AB_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + + + + +__device__ static inline void mma_AB_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} +#endif + + + + + + + + + + + +__device__ static inline void mma_AB_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} + + + + + + + + + + + +__device__ static inline void mma_ABt_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} + + + + + + + + + + + + +__device__ static inline void mma_ABt_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + + + + +__device__ static inline void mma_ABt_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} +#endif +template +__device__ static inline void mma_ABt_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + static_assert((std::is_same_v || std::is_same_v) && + (std::is_same_v || std::is_same_v)); + imma16832( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + imma16832( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} + + + + + + + + + + + + + +__device__ static inline void mma_AtB_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} + + + + + + + + + + + + +__device__ static inline void mma_AtB_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + + + + +__device__ static inline void mma_AtB_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} +#endif + + + + + + + + + + + + +__device__ static inline void mma_AtBt_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} + + + + + + + + + + + + +__device__ static inline void mma_AtBt_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + + + + +__device__ static inline void mma_AtBt_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} +#endif + + + + + + + + + + + + + + + +template +__device__ static inline void mma_AB(D &d, + const A &a, + const B &b, + const C &c) { + KITTENS_CHECK_WARP + static_assert(D::rows == A::rows && D::cols == B::cols); + static_assert(A::cols == B::rows); + static_assert(D::rows == C::rows && D::cols == C::cols); + #if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert( + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) + ); + #else + static_assert( + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) + ); + #endif + #pragma unroll + for(int n = 0; n < D::height; n++) { + #pragma unroll + for(int m = 0; m < D::width; m++) { + mma_AB_base( + d.tiles[n][m], + a.tiles[n][0], + b.tiles[0][m], + c.tiles[n][m] + ); + #pragma unroll + for(int k = 1; k < A::width; k++) { + mma_AB_base( + d.tiles[n][m], + a.tiles[n][k], + b.tiles[k][m], + d.tiles[n][m] + ); + } + } + } +} + + + + + + + + + + + + + + +template +__device__ static inline void mma_ABt(D &d, + const A &a, + const B &b, + const C &c) { + KITTENS_CHECK_WARP + static_assert(D::rows == A::rows && D::cols == B::rows); + static_assert(A::cols == B::cols); + static_assert(D::rows == C::rows && D::cols == C::cols); + #if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert( + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && + (std::is_same_v || std::is_same_v) && + (std::is_same_v || std::is_same_v) && + std::is_same_v) + ); + #else + static_assert( + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && + (std::is_same_v || std::is_same_v) && + (std::is_same_v || std::is_same_v) && + std::is_same_v) + ); + #endif + #pragma unroll + for(int n = 0; n < D::height; n++) { + #pragma unroll + for(int m = 0; m < D::width; m++) { + mma_ABt_base( + d.tiles[n][m], + a.tiles[n][0], + b.tiles[m][0], + c.tiles[n][m] + ); + #pragma unroll + for(int k = 1; k < A::width; k++) { + mma_ABt_base( + d.tiles[n][m], + a.tiles[n][k], + b.tiles[m][k], + d.tiles[n][m] + ); + } + } + } +} + + + + + + + + + + + + + + +template +__device__ static inline void mma_AtB(D &d, + const A &a, + const B &b, + const C &c) { + KITTENS_CHECK_WARP + static_assert(D::rows == A::cols && D::cols == B::cols); + static_assert(A::rows == B::rows); + static_assert(D::rows == C::rows && D::cols == C::cols); + #if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert( + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) + ); + #else + static_assert( + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) + ); + #endif + #pragma unroll + for(int n = 0; n < D::height; n++) { + #pragma unroll + for(int m = 0; m < D::width; m++) { + mma_AtB_base( + d.tiles[n][m], + a.tiles[0][n], + b.tiles[0][m], + c.tiles[n][m] + ); + #pragma unroll + for(int k = 1; k < A::height; k++) { + mma_AtB_base( + d.tiles[n][m], + a.tiles[k][n], + b.tiles[k][m], + d.tiles[n][m] + ); + } + } + } +} + + + + + + + + + + + + + + +template +__device__ static inline void mma_AtBt(D &d, + const A &a, + const B &b, + const C &c) { + KITTENS_CHECK_WARP + static_assert(D::rows == A::cols && D::cols == B::rows); + static_assert(A::rows == B::cols); + static_assert(D::rows == C::rows && D::cols == C::cols); + #if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert( + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) + ); + #else + static_assert( + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) + ); + #endif + #pragma unroll + for(int n = 0; n < D::height; n++) { + #pragma unroll + for(int m = 0; m < D::width; m++) { + mma_AtBt_base( + d.tiles[n][m], + a.tiles[0][n], + b.tiles[m][0], + c.tiles[n][m] + ); + #pragma unroll + for(int k = 1; k < A::height; k++) { + mma_AtBt_base( + d.tiles[n][m], + a.tiles[k][n], + b.tiles[m][k], + d.tiles[n][m] + ); + } + } + } +} + +template +__device__ static inline void mma(D &d, + const A &a, + const B &b, + const C &c) { + KITTENS_CHECK_WARP + if constexpr(trans_A == transpose::T) { + if constexpr(trans_B == transpose::T) { + mma_AtBt(d, a, b, c); + } else { + mma_AtB(d, a, b, c); + } + } else { + if constexpr(trans_B == transpose::T) { + mma_ABt(d, a, b, c); + } else { + mma_AB(d, a, b, c); + } + } +} +template +__device__ static inline C mma(const A &a, + const B &b, + const C &c) { + KITTENS_CHECK_WARP + C d; + if constexpr(trans_A == transpose::T) { + if constexpr(trans_B == transpose::T) { + mma_AtBt(d, a, b, c); + } else { + mma_AtB(d, a, b, c); + } + } else { + if constexpr(trans_B == transpose::T) { + mma_ABt(d, a, b, c); + } else { + mma_AB(d, a, b, c); + } + } + return d; +} + + + + + + + + + + + + + + + + + + + + + + + +template +__device__ static inline void mma_AB(crt_hf &d, + const crt_hf &a, + const crt_hf &b, + const crt_hf &c) { + KITTENS_CHECK_WARP + + + ::kittens::group<1>::copy(d.real, c.real); + ::kittens::group<1>::copy(d.imag, c.imag); + + + rt_hf tmp; + + constexpr half factor = std::bit_cast<__half>(uint16_t(0xFB80)); + ::kittens::group<1>::mul(tmp, a.imag, factor); + mma_AB(d.real, a.real, b.real, d.real); + mma_AB(d.real, tmp, b.imag, d.real); + + mma_AB(d.imag, a.real, b.imag, d.imag); + mma_AB(d.imag, a.imag, b.real, d.imag); +} + + + + + + + + + + + + + + +template +__device__ static inline void mma_AB(crt_fl &d, + const crt_bf &a, + const crt_bf &b, + const crt_fl &c) { + KITTENS_CHECK_WARP + + + ::kittens::group<1>::copy(d.real, c.real); + ::kittens::group<1>::copy(d.imag, c.imag); + + + kittens::rt_bf tmp; + + constexpr bf16 factor = std::bit_cast<__nv_bfloat16>(uint16_t(0xBF80)); + ::kittens::group<1>::mul(tmp, a.imag, factor); + mma_AB(d.real, a.real, b.real, d.real); + mma_AB(d.real, tmp, b.imag, d.real); + + mma_AB(d.imag, a.real, b.imag, d.imag); + mma_AB(d.imag, a.imag, b.real, d.imag); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/warpgroup.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/warpgroup.cuh new file mode 100644 index 0000000..5b5a0a3 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/warpgroup.cuh @@ -0,0 +1,1167 @@ + + + + + + + + + + + + + + + + + + + + + +template +__device__ static inline void mma_fence(D &dst) { + KITTENS_CHECK_WARPGROUP + #pragma unroll + for(int i = 0; i < D::height; i++) { + #pragma unroll + for(int j = 0; j < D::width; j++) { + #pragma unroll + for(int k = 0; k < D::packed_per_tile; k++) { + if constexpr(std::is_same_v) { + asm volatile("" : "+f"(dst.tiles[i][j].data[k].x) :: "memory"); + asm volatile("" : "+f"(dst.tiles[i][j].data[k].y) :: "memory"); + } else if constexpr(std::is_same_v) { + asm volatile("" : "+r"(dst.tiles[i][j].data[k].x) :: "memory"); + asm volatile("" : "+r"(dst.tiles[i][j].data[k].y) :: "memory"); + } else { + asm volatile("" : "+r"(*(uint32_t*)&dst.tiles[i][j].data[k]) :: "memory"); + } + } + } + } + asm volatile ("wgmma.fence.sync.aligned;\n" ::: "memory"); + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); +} +template +__device__ static inline void mma_fence(D &dst) { + KITTENS_CHECK_WARPGROUP + #pragma unroll + for(int i = 0; i < D::height; i++) { + #pragma unroll + for(int j = 0; j < D::width; j++) { + #pragma unroll + for(int k = 0; k < dst.real.packed_per_tile; k++) { + if constexpr(std::is_same_v) { + asm volatile("" : "+f"(dst.real.tiles[i][j].data[k].x) :: "memory"); + asm volatile("" : "+f"(dst.real.tiles[i][j].data[k].y) :: "memory"); + asm volatile("" : "+f"(dst.imag.tiles[i][j].data[k].x) :: "memory"); + asm volatile("" : "+f"(dst.imag.tiles[i][j].data[k].y) :: "memory"); + } else if constexpr(std::is_same_v) { + asm volatile("" : "+r"(dst.real.tiles[i][j].data[k].x) :: "memory"); + asm volatile("" : "+r"(dst.real.tiles[i][j].data[k].y) :: "memory"); + asm volatile("" : "+r"(dst.imag.tiles[i][j].data[k].x) :: "memory"); + asm volatile("" : "+r"(dst.imag.tiles[i][j].data[k].y) :: "memory"); + } else { + asm volatile("" : "+r"(*(uint32_t*)&dst.real.tiles[i][j].data[k]) :: "memory"); + asm volatile("" : "+r"(*(uint32_t*)&dst.imag.tiles[i][j].data[k]) :: "memory"); + } + } + } + } + asm volatile ("wgmma.fence.sync.aligned;\n" ::: "memory"); + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); +} +template +__device__ static inline void mma_fence() { + KITTENS_CHECK_WARPGROUP + asm volatile ("wgmma.fence.sync.aligned;\n" ::: "memory"); + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); +} + + + + +template +__device__ static inline void mma_commit_group() { + KITTENS_CHECK_WARPGROUP + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); +} + + + + + + + + +template +__device__ static inline void mma_async_wait() { + KITTENS_CHECK_WARPGROUP + asm volatile ("wgmma.wait_group.sync.aligned %0;" : : "n"(N) : "memory"); +} + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +template +__device__ static inline void mma_AB(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M_DIV_4 = A::height; + static_assert(D::height == M_DIV_4); + constexpr int N = B::cols / kittens::TILE_COL_DIM; + static_assert(D::width == N); + constexpr int K = A::width; + static_assert(B::rows/kittens::TILE_ROW_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + static_assert(!std::is_same_v && !std::is_same_v, "Transpose is not supported for integer WGMMA; use mma_ABt / mm_ABt."); + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + using base = kittens::detail::wgmma::base*N, 0, 1>; + kittens::st_descriptor, 1> b_desc(b); + + if constexpr (fence) { mma_fence(d); } + + + #pragma unroll + for(int m = 0; m < M_DIV_4; m++) { + rt, TILE_COL_DIM*N, ducks::rt_layout::row> &d_ref = group<1>::subtile_inplace>(d, m); + base::rt_st( + d_ref, + a.tiles[m][0], + b_desc.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::rt_st( + d_ref, + a.tiles[m][k], + b_desc.chunk_descriptor(k), + 1 + ); + } + } + mma_commit_group(); +} +template +__device__ static inline void mm_AB(D &d, + const A &a, + const B &b) { + mma_AB(d, a, b); +} + +template +__device__ static inline void mma_AB(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M = A::rows / kittens::TILE_ROW_DIM; + static_assert(M == 4); + static_assert(D::height == 1); + constexpr int N = B::cols / kittens::TILE_COL_DIM; + constexpr int K = A::cols / kittens::TILE_COL_DIM; + static_assert(B::rows / kittens::TILE_ROW_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + static_assert(!std::is_same_v && !std::is_same_v, "Transpose is not supported for integer WGMMA; use mma_ABt / mm_ABt."); + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + using base = kittens::detail::wgmma::base*N, 0, 1>; + kittens::st_descriptor, 0> a_desc(a); + kittens::st_descriptor, 1> b_desc(b); + + if constexpr (fence) { mma_fence(d); } + + + base::st_st( + d, + a_desc.chunk_descriptor(0), + b_desc.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d, + a_desc.chunk_descriptor(k), + b_desc.chunk_descriptor(k), + 1 + ); + } + mma_commit_group(); +} +template +__device__ static inline void mm_AB(D &d, + const A &a, + const B &b) { + mma_AB(d, a, b); +} + + + + + + + + + + + + + + + + +template +__device__ static inline void mma_ABt(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M_DIV_4 = A::height; + static_assert(D::height == M_DIV_4); + constexpr int N = B::rows / kittens::TILE_ROW_DIM; + constexpr int K = A::width; + static_assert(B::cols/kittens::TILE_COL_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + using base = kittens::detail::wgmma::base*N, 0, 0>; + kittens::st_descriptor, 0> b_desc(b); + + if constexpr (fence) { mma_fence(d); } + + + #pragma unroll + for(int m = 0; m < M_DIV_4; m++) { + rt, TILE_COL_DIM*N, ducks::rt_layout::row> &d_ref = group<1>::subtile_inplace>(d, m); + base::rt_st( + d_ref, + a.tiles[m][0], + b_desc.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::rt_st( + d_ref, + a.tiles[m][k], + b_desc.chunk_descriptor(k), + 1 + ); + } + } + mma_commit_group(); +} +template +__device__ static inline void mm_ABt(D &d, + const A &a, + const B &b) { + mma_ABt(d, a, b); +} + + + + + + + + + + + + + + + + +template +__device__ static inline void mma_ABt(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M = A::rows / kittens::TILE_ROW_DIM; + static_assert(M == 4); + static_assert(D::height == 1); + constexpr int N = B::rows / kittens::TILE_ROW_DIM; + constexpr int K = A::cols / kittens::TILE_COL_DIM; + static_assert(B::cols/kittens::TILE_COL_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + using base = kittens::detail::wgmma::base*N, 0, 0>; + kittens::st_descriptor, 0> a_desc(a); + kittens::st_descriptor, 0> b_desc(b); + + if constexpr (fence) { mma_fence(d); } + + + base::st_st( + d, + a_desc.chunk_descriptor(0), + b_desc.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d, + a_desc.chunk_descriptor(k), + b_desc.chunk_descriptor(k), + 1 + ); + } + mma_commit_group(); +} +template +__device__ static inline void mm_ABt(D &d, + const A &a, + const B &b) { + mma_ABt(d, a, b); +} + + + + + + + + + + + + + + + + +template +__device__ static inline void mma_AtB(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M = A::cols / kittens::TILE_COL_DIM; + static_assert(M == 4); + static_assert(D::height == 1); + constexpr int N = B::cols / kittens::TILE_COL_DIM; + constexpr int K = A::rows / kittens::TILE_ROW_DIM; + static_assert(B::rows/kittens::TILE_ROW_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + static_assert(!std::is_same_v && !std::is_same_v, "Transpose is not supported for integer WGMMA; use mma_ABt / mm_ABt."); + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + using base = kittens::detail::wgmma::base*N, 1, 1>; + kittens::st_descriptor, 1> a_desc(a); + kittens::st_descriptor, 1> b_desc(b); + + if constexpr (fence) { mma_fence(d); } + + + base::st_st( + d, + a_desc.chunk_descriptor(0), + b_desc.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d, + a_desc.chunk_descriptor(k), + b_desc.chunk_descriptor(k), + 1 + ); + } + mma_commit_group(); +} +template +__device__ static inline void mm_AtB(D &d, + const A &a, + const B &b) { + mma_AtB(d, a, b); +} + + + + + + + + + + + + +template +__device__ static inline void mma_AtBt(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M = A::cols / kittens::TILE_COL_DIM; + static_assert(M == 4); + static_assert(D::height == 1); + constexpr int N = B::rows / kittens::TILE_ROW_DIM; + constexpr int K = A::rows / kittens::TILE_ROW_DIM; + static_assert(B::cols/kittens::TILE_COL_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + static_assert(!std::is_same_v && !std::is_same_v, "Transpose is not supported for integer WGMMA; use mma_ABt / mm_ABt."); + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + using base = kittens::detail::wgmma::base*N, 1, 0>; + kittens::st_descriptor, 1> a_desc(a); + kittens::st_descriptor, 0> b_desc(b); + + if constexpr (fence) { mma_fence(d); } + + + base::st_st( + d, + a_desc.chunk_descriptor(0), + b_desc.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d, + a_desc.chunk_descriptor(k), + b_desc.chunk_descriptor(k), + 1 + ); + } + mma_commit_group(); +} +template +__device__ static inline void mm_AtBt(D &d, + const A &a, + const B &b) { + mma_AtBt(d, a, b); +} + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +template +__device__ static inline void mma_AB(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M_DIV_4 = A::height; + static_assert(D::height == M_DIV_4); + constexpr int N = B::cols / kittens::TILE_COL_DIM; + constexpr int K = A::width; + static_assert(B::rows/kittens::TILE_ROW_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + static_assert(!std::is_same_v && !std::is_same_v, "Transpose is not supported for integer WGMMA; use mma_ABt / mm_ABt."); + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + using base = kittens::detail::wgmma::base*N, 0, 1>; + kittens::st_descriptor, 1> b_desc_real(b.real); + kittens::st_descriptor, 1> b_desc_imag(b.imag); + + if constexpr (fence) { mma_fence(d); } + + + #pragma unroll + for(int m = 0; m < M_DIV_4; m++) { + rt, TILE_COL_DIM*N, ducks::rt_layout::row> &d_ref = group<1>::subtile_inplace>(d.real, m); + base::rt_st( + d_ref, + a.real.tiles[m][0], + b_desc_real.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::rt_st( + d_ref, + a.real.tiles[m][k], + b_desc_real.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::rt_st<-1>( + d_ref, + a.imag.tiles[m][k], + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + } + #pragma unroll + for(int m = 0; m < M_DIV_4; m++) { + rt, TILE_COL_DIM*N, ducks::rt_layout::row> &d_ref = group<1>::subtile_inplace>(d.imag, m); + base::rt_st( + d_ref, + a.real.tiles[m][0], + b_desc_imag.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::rt_st( + d_ref, + a.real.tiles[m][k], + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::rt_st( + d_ref, + a.imag.tiles[m][k], + b_desc_real.chunk_descriptor(k), + 1 + ); + } + } + mma_commit_group(); +} +template +__device__ static inline void mm_AB(D &d, + const A &a, + const B &b) { + mma_AB(d, a, b); +} + +template +__device__ static inline void mma_AB(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M = A::rows / kittens::TILE_ROW_DIM; + static_assert(M == 4); + static_assert(D::height == 1); + constexpr int N = B::cols / kittens::TILE_COL_DIM; + constexpr int K = A::cols / kittens::TILE_COL_DIM; + static_assert(B::rows/kittens::TILE_ROW_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + static_assert(!std::is_same_v && !std::is_same_v, "Transpose is not supported for integer WGMMA; use mma_ABt / mm_ABt."); + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + using base = kittens::detail::wgmma::base*N, 0, 1>; + kittens::st_descriptor, 0> a_desc_real(a.real); + kittens::st_descriptor, 0> a_desc_imag(a.imag); + kittens::st_descriptor, 1> b_desc_real(b.real); + kittens::st_descriptor, 1> b_desc_imag(b.imag); + + if constexpr (fence) { mma_fence(d); } + + + base::st_st( + d.real, + a_desc_real.chunk_descriptor(0), + b_desc_real.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d.real, + a_desc_real.chunk_descriptor(k), + b_desc_real.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::st_st<-1>( + d.real, + a_desc_imag.chunk_descriptor(k), + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + base::st_st( + d.imag, + a_desc_real.chunk_descriptor(0), + b_desc_imag.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d.imag, + a_desc_real.chunk_descriptor(k), + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::st_st( + d.imag, + a_desc_imag.chunk_descriptor(k), + b_desc_real.chunk_descriptor(k), + 1 + ); + } + mma_commit_group(); +} +template +__device__ static inline void mm_AB(D &d, + const A &a, + const B &b) { + mma_AB(d, a, b); +} + + + + + + + + + + + + + + + + +template +__device__ static inline void mma_ABt(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M_DIV_4 = A::height; + static_assert(D::height == M_DIV_4); + constexpr int N = B::rows / kittens::TILE_ROW_DIM; + constexpr int K = A::width; + static_assert(B::cols/kittens::TILE_COL_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + using base = kittens::detail::wgmma::base*N, 0, 0>; + kittens::st_descriptor, 0> b_desc_real(b.real); + kittens::st_descriptor, 0> b_desc_imag(b.imag); + + if constexpr (fence) { mma_fence(d); } + + + #pragma unroll + for(int m = 0; m < M_DIV_4; m++) { + rt, TILE_ROW_DIM*N, ducks::rt_layout::row> &d_ref = group<1>::subtile_inplace>(d.real, m); + base::rt_st( + d_ref, + a.real.tiles[m][0], + b_desc_real.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::rt_st( + d_ref, + a.real.tiles[m][k], + b_desc_real.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::rt_st<-1>( + d_ref, + a.imag.tiles[m][k], + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + } + #pragma unroll + for(int m = 0; m < M_DIV_4; m++) { + rt, TILE_ROW_DIM*N, ducks::rt_layout::row> &d_ref = group<1>::subtile_inplace>(d.imag, m); + base::rt_st( + d_ref, + a.real.tiles[m][0], + b_desc_imag.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::rt_st( + d_ref, + a.real.tiles[m][k], + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::rt_st( + d_ref, + a.imag.tiles[m][k], + b_desc_real.chunk_descriptor(k), + 1 + ); + } + } + mma_commit_group(); +} +template +__device__ static inline void mm_ABt(D &d, + const A &a, + const B &b) { + mma_ABt(d, a, b); +} + + + + + + + + + + + + + + + + +template +__device__ static inline void mma_ABt(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M = A::rows / kittens::TILE_ROW_DIM; + static_assert(M == 4); + static_assert(D::height == 1); + constexpr int N = B::rows / kittens::TILE_ROW_DIM; + constexpr int K = A::cols / kittens::TILE_COL_DIM; + static_assert(B::cols/kittens::TILE_COL_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + using base = kittens::detail::wgmma::base*N, 0, 0>; + kittens::st_descriptor, 0> a_desc_real(a.real); + kittens::st_descriptor, 0> a_desc_imag(a.imag); + kittens::st_descriptor, 0> b_desc_real(b.real); + kittens::st_descriptor, 0> b_desc_imag(b.imag); + + if constexpr (fence) { mma_fence(d); } + + + base::st_st( + d.real, + a_desc_real.chunk_descriptor(0), + b_desc_real.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d.real, + a_desc_real.chunk_descriptor(k), + b_desc_real.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::st_st<-1>( + d.real, + a_desc_imag.chunk_descriptor(k), + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + base::st_st( + d.imag, + a_desc_real.chunk_descriptor(0), + b_desc_imag.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d.imag, + a_desc_real.chunk_descriptor(k), + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::st_st( + d.imag, + a_desc_imag.chunk_descriptor(k), + b_desc_real.chunk_descriptor(k), + 1 + ); + } + mma_commit_group(); +} +template +__device__ static inline void mm_ABt(D &d, + const A &a, + const B &b) { + mma_ABt(d, a, b); +} + + + + + + + + + + + + + + + + +template +__device__ static inline void mma_AtB(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M = A::cols / kittens::TILE_COL_DIM; + static_assert(M == 4); + static_assert(D::height == 1); + constexpr int N = B::cols / kittens::TILE_COL_DIM; + constexpr int K = A::rows / kittens::TILE_ROW_DIM; + static_assert(B::rows/kittens::TILE_ROW_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + static_assert(!std::is_same_v && !std::is_same_v, "Transpose is not supported for integer WGMMA; use mma_ABt / mm_ABt."); + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + using base = kittens::detail::wgmma::base*N, 1, 1>; + kittens::st_descriptor, 1> a_desc_real(a.real); + kittens::st_descriptor, 1> a_desc_imag(a.imag); + kittens::st_descriptor, 1> b_desc_real(b.real); + kittens::st_descriptor, 1> b_desc_imag(b.imag); + + if constexpr (fence) { mma_fence(d); } + + + base::st_st( + d.real, + a_desc_real.chunk_descriptor(0), + b_desc_real.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d.real, + a_desc_real.chunk_descriptor(k), + b_desc_real.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::st_st<-1>( + d.real, + a_desc_imag.chunk_descriptor(k), + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + base::st_st( + d.imag, + a_desc_real.chunk_descriptor(0), + b_desc_imag.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d.imag, + a_desc_real.chunk_descriptor(k), + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::st_st( + d.imag, + a_desc_imag.chunk_descriptor(k), + b_desc_real.chunk_descriptor(k), + 1 + ); + } + mma_commit_group(); +} +template +__device__ static inline void mm_AtB(D &d, + const A &a, + const B &b) { + mma_AtB(d, a, b); +} + + + + + + + + + + + + +template +__device__ static inline void mma_AtBt(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M = A::cols / kittens::TILE_COL_DIM; + static_assert(M == 4); + static_assert(D::height == 1); + constexpr int N = B::rows / kittens::TILE_ROW_DIM; + constexpr int K = A::rows / kittens::TILE_ROW_DIM; + static_assert(B::cols/kittens::TILE_COL_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + static_assert(!std::is_same_v && !std::is_same_v, "Transpose is not supported for integer WGMMA; use mma_ABt / mm_ABt."); + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + using base = kittens::detail::wgmma::base*N, 1, 0>; + kittens::st_descriptor, 1> a_desc_real(a.real); + kittens::st_descriptor, 1> a_desc_imag(a.imag); + kittens::st_descriptor, 0> b_desc_real(b.real); + kittens::st_descriptor, 0> b_desc_imag(b.imag); + + if constexpr (fence) { mma_fence(d); } + + + base::st_st( + d.real, + a_desc_real.chunk_descriptor(0), + b_desc_real.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d.real, + a_desc_real.chunk_descriptor(k), + b_desc_real.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::st_st<-1>( + d.real, + a_desc_imag.chunk_descriptor(k), + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + base::st_st( + d.imag, + a_desc_real.chunk_descriptor(0), + b_desc_imag.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d.imag, + a_desc_real.chunk_descriptor(k), + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::st_st( + d.imag, + a_desc_imag.chunk_descriptor(k), + b_desc_real.chunk_descriptor(k), + 1 + ); + } + mma_commit_group(); +} +template +__device__ static inline void mm_AtBt(D &d, + const A &a, + const B &b) { + mma_AtBt(d, a, b); +} + + + +template +__device__ static inline void mma(D &d, + const A &a, + const B &b) { + if constexpr(trans_A == transpose::T) { + if constexpr(trans_B == transpose::T) { + mma_AtBt(d, a, b); + } else { + mma_AtB(d, a, b); + } + } else { + if constexpr(trans_B == transpose::T) { + mma_ABt(d, a, b); + } else { + mma_AB(d, a, b); + } + } +} +template +__device__ static inline void mm(D &d, + const A &a, + const B &b) { + if constexpr(trans_A == transpose::T) { + if constexpr(trans_B == transpose::T) { + mm_AtBt(d, a, b); + } else { + mm_AtB(d, a, b); + } + } else { + if constexpr(trans_B == transpose::T) { + mm_ABt(d, a, b); + } else { + mm_AB(d, a, b); + } + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/register.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/register.cuh new file mode 100644 index 0000000..404f8af --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/register.cuh @@ -0,0 +1,7 @@ + + + + + +#include "tile/tile.cuh" +#include "vec/vec.cuh" diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/complex_conversions.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/complex_conversions.cuh new file mode 100644 index 0000000..c2bd5f1 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/complex_conversions.cuh @@ -0,0 +1,98 @@ + + + + + + + + + + + + + + + + + + + + +template +__device__ static inline void swap_layout(crt::type> &dst, const crt &src) { + swap_layout(dst.real, src.real); + swap_layout(dst.real, src.real); +} + + + + + + + + + + +template +__device__ static inline crt::type>& swap_layout_inplace(crt &tile) { + tile.real = swap_layout_inplace(tile.real); + tile.imag = swap_layout_inplace(tile.imag); + return tile; +} + + + + + + + + + + + + + + + + +template +__device__ static inline void transpose_sep(crt &dst, const crt &src) { + transpose_sep(dst.real, src.real); + transpose_sep(dst.imag, src.imag); +} + + + + + + + + + + +template +__device__ static inline crt& transpose_inplace(crt &tile) { + tile.real = transpose_inplace(tile.real); + tile.imag = transpose_inplace(tile.imag); + + return tile; +} + + + + + + + + + + + + + + +template +__device__ static inline void copy(crt &dst, const crt &src) { + copy(dst.real, src.real); + copy(dst.imag, src.imag); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/complex_maps.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/complex_maps.cuh new file mode 100644 index 0000000..55c3f36 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/complex_maps.cuh @@ -0,0 +1,137 @@ + + + + + + + + + + + +template +__device__ static inline void zero(T &dst) { + zero(dst.real); + zero(dst.imag); +} + + + + + + + +template +__device__ static inline void exp(T &dst, const T &src) { + using dtype = T::dtype; + dtype tmp; + + dtype rdst; + dtype idst; + + + exp(rdst, src.real); + copy(idst, rdst); + + cos(tmp, src.imag); + mul(rdst, rdst, tmp); + sin(tmp, src.imag); + mul(idst, idst, tmp); + + copy(dst.real, rdst); + copy(dst.imag, idst); +} + + + + + + + + +template +__device__ static inline void add(T &dst, const T &lhs, const T &rhs) { + add(dst.real, lhs.real, rhs.real); + add(dst.imag, lhs.imag, rhs.imag); +} + + + + + + + + + +template +__device__ static inline void sub(T &dst, const T &lhs, const T &rhs) { + sub(dst.real, lhs.real, rhs.real); + sub(dst.imag, lhs.imag, rhs.imag); +} + + + + + + + + +template +__device__ static inline void mul(T &dst, const T &lhs, const T &rhs) { + using dtype = T::component; + dtype tmp; + + dtype rdst; + dtype idst; + + + + mul(rdst, lhs.real, rhs.real); + mul(tmp, lhs.imag, rhs.imag); + sub(rdst, rdst, tmp); + + + mul(idst, lhs.imag, rhs.real); + mul(tmp, lhs.real, rhs.imag); + add(idst, idst, tmp); + + copy(dst.real, rdst); + copy(dst.imag, idst); +} + + + + + + + + +template +__device__ static inline void div(T &dst, const T &lhs, const T &rhs) { + using dtype = T::dtype; + dtype tmp; + dtype denom; + + dtype rdst; + dtype idst; + + + mul(tmp, rhs.real, rhs.real); + mul(denom, rhs.imag, rhs.imag); + add(denom, tmp, denom); + + mul(rdst, lhs.real, rhs.real); + mul(tmp, lhs.imag, rhs.imag); + add(rdst, rdst, tmp); + + mul(dst.imag, lhs.imag, rhs.real); + mul(tmp, lhs.real, rhs.imag); + sub(idst, idst, tmp); + + div(rdst, rdst, denom); + div(idst, idst, denom); + copy(dst.real, rdst); + copy(dst.imag, idst); +} + + diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/conversions.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/conversions.cuh new file mode 100644 index 0000000..44d2cfc --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/conversions.cuh @@ -0,0 +1,917 @@ + + + + + + + + + + + + + + + + + +__device__ static inline void swap_layout_8(bf16_2 &dst, const bf16_2 &src) { + KITTENS_CHECK_WARP + asm volatile ( + "movmatrix.sync.aligned.m8n8.trans.b16 %0, %1;\n" + : "+r"(*(uint32_t*)(&dst)) + : "r"(*(uint32_t*)(&src)) + ); +} + + + + + + + + + + + +template +__device__ static inline void swap_layout(rt_base::type> &dst, const rt_base &src) { + swap_layout_8(dst.data[0], src.data[0]); + + + typename rt_base::T2 data1_cache = src.data[1]; + swap_layout_8(dst.data[1], src.data[2]); + swap_layout_8(dst.data[2], data1_cache); + swap_layout_8(dst.data[3], src.data[3]); +} + + + + + + + + + + + + + +template +__device__ static inline void swap_layout(rt::type> &dst, const rt &src) { + #pragma unroll + for(int i = 0; i < dst.height; i++) { + #pragma unroll + for(int j = 0; j < dst.width; j++) { + swap_layout(dst.tiles[i][j], src.tiles[i][j]); + } + } +} + + + + + + + + + + + + +template +__device__ static inline rt_base::type>& swap_layout_inplace(const rt_base &src) { + rt_base::type> &dst = *(rt_base::type>*)(&src); + swap_layout(dst, src); + return dst; +} + + + + + + + + + + + + + +template +__device__ static inline rt::type>& swap_layout_inplace(rt &tile) { + #pragma unroll + for(int i = 0; i < tile.height; i++) { + #pragma unroll + for(int j = 0; j < tile.width; j++) { + swap_layout_inplace(tile.tiles[i][j]); + } + } + return *(rt::type>*)(&tile); +} + + + + + + + + + + + +template +__device__ static inline void transpose(rt_base &dst, const rt_base &src) { + swap_layout_8(dst.data[0], src.data[0]); + + + typename rt_base::T2 data1_cache = src.data[1]; + swap_layout_8(dst.data[1], src.data[2]); + swap_layout_8(dst.data[2], data1_cache); + swap_layout_8(dst.data[3], src.data[3]); +} + + + + + + + + + + + + + +template +__device__ static inline void transpose_sep(RT &dst, const rt &src) { + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + transpose(dst.tiles[i][j], src.tiles[j][i]); + } + } +} + + + + + + + + + +template +__device__ static inline rt_base& transpose_inplace(rt_base &src) { + transpose(src, src); + return src; +} + + + + + + + + + + +template +__device__ static inline rt& transpose_inplace(rt &tile) { + static_assert(_cols == _rows, "in-place register tile transpose is only allowed for square tiles."); + #pragma unroll + for(int i = 0; i < tile.height; i++) { + #pragma unroll + for(int j = 0; j < i; j++) { + rt_base tmp; + copy(tmp, tile.tiles[i][j]); + transpose(tile.tiles[i][j], tile.tiles[j][i]); + transpose(tile.tiles[j][i], tmp); + } + transpose_inplace(tile.tiles[i][i]); + } + return tile; +} + + + + + + + + + + + + +template +__device__ static inline void copy(rt_base &dst, const rt_base &src) { + using T2 = typename base_types::packing::packed_type; + using U2 = typename base_types::packing::packed_type; + #pragma unroll + for(int k = 0; k < dst.packed_per_thread; k++) { + dst.data[k] = base_types::convertor::convert(src.data[k]); + } +} +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + + + + +template +__device__ static inline void copy(rt &dst, const rt &src) { + + if constexpr ( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) + ) { + + int laneid = threadIdx.x % 32; + + #pragma unroll + for(int i = 0; i < dst.height; i++) { + #pragma unroll + for(int j = 0; j < dst.width; j++) { + #pragma unroll + for(int k = 0; k < dst.tiles[0][0].packed_per_thread; k++) { + + + using src_t = std::conditional_t, float2, std::conditional_t, bf16_2, half2>>; + src_t val1, val2; + + + if (laneid % 2 == 0) { + + val1 = src.tiles[i][2*j + k/2].data[(k%2)+0]; + val2 = src.tiles[i][2*j + k/2].data[(k%2)+2]; + } else { + + val1 = src.tiles[i][2*j + k/2].data[(k%2)+2]; + val2 = src.tiles[i][2*j + k/2].data[(k%2)+0]; + } + + + int row_mask = 4 * ( laneid / 4 ); + int row_offset = row_mask + ( (laneid-row_mask) / 2 ) + ( laneid % 2 ); + int src_offset = (laneid % 2 == 0 ) ? row_offset + 0 : ( row_offset + 1 ); + src_t val01 = packed_shfl_sync(MASK_ALL, val1, src_offset); + + int src_offset2 = (laneid % 4 < 2 ) ? src_offset + 1 : (src_offset - 1); + src_t val23 = packed_shfl_sync(MASK_ALL, val2, src_offset2); + + + float4 f4; + using fp8_4_t = std::conditional_t, fp8e4m3_4, fp8e5m2_4>; + fp8_4_t f4_fp8; + if ( laneid % 4 < 2 ) { + f4.x = val01.x; + f4.y = val01.y; + f4.z = val23.x; + f4.w = val23.y; + f4_fp8 = base_types::convertor::convert(f4); + dst.tiles[i][j].data[k] = f4_fp8; + } else { + f4.x = val23.x; + f4.y = val23.y; + f4.z = val01.x; + f4.w = val01.y; + f4_fp8 = base_types::convertor::convert(f4); + dst.tiles[i][j].data[k] = f4_fp8; + } + } + } + } + } + else if constexpr ( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) + ) { + + int laneid = threadIdx.x % 32; + + #pragma unroll + for(int i = 0; i < src.height; i++) { + #pragma unroll + for(int j = 0; j < src.width; j++) { + #pragma unroll + for(int k = 0; k < src.tiles[0][0].packed_per_thread; k++) { + int dst_j = 2*j + k/2; + + + using fp8_4_t = std::conditional_t, fp8e4m3_4, fp8e5m2_4>; + fp8_4_t val = src.tiles[i][j].data[k]; + float4 f4 = base_types::convertor::convert(val); + float2 f2_0, f2_1; + if ( laneid % 4 < 2 ) { + f2_0 = make_float2(f4.x, f4.y); + f2_1 = make_float2(f4.z, f4.w); + } + else { + f2_0 = make_float2(f4.z, f4.w); + f2_1 = make_float2(f4.x, f4.y); + } + + int row_offset = 4 * (laneid/4) + (laneid%2) * 2 + (laneid%4) / 2; + float2 f2_0_shfl = packed_shfl_sync(MASK_ALL, f2_0, row_offset); + float2 f2_1_shfl = packed_shfl_sync(MASK_ALL, f2_1, row_offset^2); + + + using dst_t = std::conditional_t, float2, std::conditional_t, bf16_2, half2>>; + if constexpr (!(std::is_same_v)) { + dst_t f2_0_shfl_t = base_types::convertor::convert(f2_0_shfl); + dst_t f2_1_shfl_t = base_types::convertor::convert(f2_1_shfl); + if (laneid % 2 == 0) { + dst.tiles[i][dst_j].data[(k%2)+0] = f2_0_shfl_t; + dst.tiles[i][dst_j].data[(k%2)+2] = f2_1_shfl_t; + } else { + dst.tiles[i][dst_j].data[(k%2)+0] = f2_1_shfl_t; + dst.tiles[i][dst_j].data[(k%2)+2] = f2_0_shfl_t; + } + } else { + if (laneid % 2 == 0) { + dst.tiles[i][dst_j].data[(k%2)+0] = f2_0_shfl; + dst.tiles[i][dst_j].data[(k%2)+2] = f2_1_shfl; + } else { + dst.tiles[i][dst_j].data[(k%2)+0] = f2_1_shfl; + dst.tiles[i][dst_j].data[(k%2)+2] = f2_0_shfl; + } + } + } + } + } + } + + else { + #pragma unroll + for(int i = 0; i < dst.height; i++) { + #pragma unroll + for(int j = 0; j < dst.width; j++) { + copy(dst.tiles[i][j], src.tiles[i][j]); + } + } + } +} +#else + + + + + + + + + + + +template +__device__ static inline void copy(rt &dst, const rt &src) { + #pragma unroll + for(int i = 0; i < dst.height; i++) { + #pragma unroll + for(int j = 0; j < dst.width; j++) { + copy(dst.tiles[i][j], src.tiles[i][j]); + } + } +} +#endif + + + + + + + + + + + + + + + +template +__device__ static inline void make_causal(RT &dst, const RT &src, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + const typename RT::dtype packed_val = base_types::packing::pack(val); + #if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert(!std::is_same_v && !std::is_same_v, "Unsupported type for make_causal"); + #endif + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + if(j < i) { + #pragma unroll + for(int k = 0; k < RT::packed_per_tile; k++) { + dst.tiles[i][j].data[k] = src.tiles[i][j].data[k]; + } + } + else if(j > i) { + #pragma unroll + for(int k = 0; k < RT::packed_per_tile; k++) { + dst.tiles[i][j].data[k] = packed_val; + } + } + else { + constexpr uint32_t MASK_X = 0xFF773311, MASK_Y = 0xF7733110; + dst.tiles[i][j].data[1] = src.tiles[i][j].data[1]; + dst.tiles[i][j].data[2] = packed_val; + if((MASK_X >> laneid()) & 1) { + dst.tiles[i][j].data[0].x = src.tiles[i][j].data[0].x; + dst.tiles[i][j].data[3].x = src.tiles[i][j].data[3].x; + } + else { + dst.tiles[i][j].data[0].x = val; + dst.tiles[i][j].data[3].x = val; + } + if((MASK_Y >> laneid()) & 1) { + dst.tiles[i][j].data[0].y = src.tiles[i][j].data[0].y; + dst.tiles[i][j].data[3].y = src.tiles[i][j].data[3].y; + } + else { + dst.tiles[i][j].data[0].y = val; + dst.tiles[i][j].data[3].y = val; + } + } + __syncwarp(); + } + } +} + + + + + + + + + + + + + +template +__device__ static inline void make_causal_t(RT &dst, const RT &src, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + const typename RT::dtype packed_val = base_types::packing::pack(val); + #if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert(!std::is_same_v && !std::is_same_v, "Unsupported type for make_causal"); + #endif + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + if(j > i) { + #pragma unroll + for(int k = 0; k < RT::packed_per_tile; k++) { + dst.tiles[i][j].data[k] = src.tiles[i][j].data[k]; + } + } + else if(j < i) { + #pragma unroll + for(int k = 0; k < RT::packed_per_tile; k++) { + dst.tiles[i][j].data[k] = packed_val; + } + } + else { + constexpr uint32_t MASK_X = 0x88CCEEF; + constexpr uint32_t MASK_Y = 0x88CCEEFF; + + dst.tiles[i][j].data[1] = packed_val; + dst.tiles[i][j].data[2] = src.tiles[i][j].data[2]; + + + if((MASK_X >> laneid()) & 1) { + dst.tiles[i][j].data[0].x = src.tiles[i][j].data[0].x; + dst.tiles[i][j].data[3].x = src.tiles[i][j].data[3].x; + } + + else { + dst.tiles[i][j].data[0].x = val; + dst.tiles[i][j].data[3].x = val; + } + + + if((MASK_Y >> laneid()) & 1) { + dst.tiles[i][j].data[0].y = src.tiles[i][j].data[0].y; + dst.tiles[i][j].data[3].y = src.tiles[i][j].data[3].y; + } + + else { + dst.tiles[i][j].data[0].y = val; + dst.tiles[i][j].data[3].y = val; + } + + } + __syncwarp(); + } + } +} + + + + + + + + + + + + +template +__device__ static inline void tril(RT &dst, const RT &src, const int diagonal, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + const typename RT::dtype packed_val = base_types::packing::pack(val); + + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int global_row_idx = (i * RT::tile_size_row) + ((k % 2) * 8) + (laneid() / 4); + const int global_col_idx_x = (j * RT::tile_size_col) + ((k / 2) * 8) + ((laneid() % 4) * 2); + const int global_col_idx_y = (j * RT::tile_size_col) + ((k / 2) * 8) + ((laneid() % 4) * 2) + 1; + + if (global_col_idx_x <= global_row_idx + diagonal) { dst.tiles[i][j].data[k].x = src.tiles[i][j].data[k].x; } + else { dst.tiles[i][j].data[k].x = val; } + + if (global_col_idx_y <= global_row_idx + diagonal) { dst.tiles[i][j].data[k].y = src.tiles[i][j].data[k].y; } + else { dst.tiles[i][j].data[k].y = val; } + } + } + __syncwarp(); + } +} +template +__device__ static inline void tril(RT &dst, const RT &src, const int diagonal, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int global_row_idx_x = (i * RT::tile_size_row) + ((k / 2) * 8) + ((laneid() % 4) * 2); + const int global_row_idx_y = (i * RT::tile_size_row) + ((k / 2) * 8) + ((laneid() % 4) * 2) + 1; + const int global_col_idx = (j * RT::tile_size_col) + ((k % 2) * 8) + (laneid() / 4); + + if (global_col_idx <= global_row_idx_x + diagonal) { dst.tiles[i][j].data[k].x = src.tiles[i][j].data[k].x; } + else { dst.tiles[i][j].data[k].x = val; } + if (global_col_idx <= global_row_idx_y + diagonal) { dst.tiles[i][j].data[k].y = src.tiles[i][j].data[k].y; } + else { dst.tiles[i][j].data[k].y = val; } + } + } + __syncwarp(); + } +} + + + + + + + + + + +template +__device__ static inline void triu(RT &dst, const RT &src, const int diagonal, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + const typename RT::dtype packed_val = base_types::packing::pack(val); + + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int global_row_idx = (i * RT::tile_size_row) + ((k % 2) * 8) + (laneid() / 4); + const int global_col_idx_x = (j * RT::tile_size_col) + ((k / 2) * 8) + ((laneid() % 4) * 2); + const int global_col_idx_y = (j * RT::tile_size_col) + ((k / 2) * 8) + ((laneid() % 4) * 2) + 1; + + if (global_col_idx_x >= global_row_idx + diagonal) { dst.tiles[i][j].data[k].x = src.tiles[i][j].data[k].x; } + else { dst.tiles[i][j].data[k].x = val; } + + if (global_col_idx_y >= global_row_idx + diagonal) { dst.tiles[i][j].data[k].y = src.tiles[i][j].data[k].y; } + else { dst.tiles[i][j].data[k].y = val; } + } + } + __syncwarp(); + } +} +template +__device__ static inline void triu(RT &dst, const RT &src, const int diagonal, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int global_row_idx_x = (i * RT::tile_size_row) + ((k / 2) * 8) + ((laneid() % 4) * 2); + const int global_row_idx_y = (i * RT::tile_size_row) + ((k / 2) * 8) + ((laneid() % 4) * 2) + 1; + const int global_col_idx = (j * RT::tile_size_col) + ((k % 2) * 8) + (laneid() / 4); + + if (global_col_idx >= global_row_idx_x + diagonal) { dst.tiles[i][j].data[k].x = src.tiles[i][j].data[k].x; } + else { dst.tiles[i][j].data[k].x = val; } + + if (global_col_idx >= global_row_idx_y + diagonal) { dst.tiles[i][j].data[k].y = src.tiles[i][j].data[k].y; } + else { dst.tiles[i][j].data[k].y = val; } + } + } + __syncwarp(); + } +} + + + + + + + + + + + + +template +__device__ static inline void right_fill(RT &dst, const RT &src, const int col_idx, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + if(col_idx >= RT::cols) return; + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int col_idx_x = (j * RT::tile_size_col) + ((k / 2) * 8) + ((laneid() % 4) * 2); + const int col_idx_y = (j * RT::tile_size_col) + ((k / 2) * 8) + ((laneid() % 4) * 2) + 1; + if (col_idx_x >= col_idx) { dst.tiles[i][j].data[k].x = val; } + else { dst.tiles[i][j].data[k].x = src.tiles[i][j].data[k].x; } + if (col_idx_y >= col_idx) { dst.tiles[i][j].data[k].y = val; } + else { dst.tiles[i][j].data[k].y = src.tiles[i][j].data[k].y; } + } + } + } +} +template +__device__ static inline void right_fill(RT &dst, const RT &src, const int col_idx, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + const typename RT::dtype packed_val = base_types::packing::pack(val); + + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int t_col_idx = (j * RT::tile_size_col) + ((k % 2) * 8) + (laneid() / 4); + if (t_col_idx >= col_idx) { dst.tiles[i][j].data[k] = packed_val; } + else { dst.tiles[i][j].data[k] = src.tiles[i][j].data[k]; } + } + } + __syncwarp(); + } +} + + + + + + + + + + +template +__device__ static inline void left_fill(RT &dst, const RT &src, const int col_idx, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + if(col_idx <= 0) return; + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int col_idx_x = (j * RT::tile_size_col) + ((k / 2) * 8) + ((laneid() % 4) * 2); + const int col_idx_y = (j * RT::tile_size_col) + ((k / 2) * 8) + ((laneid() % 4) * 2) + 1; + if (col_idx_x < col_idx) { dst.tiles[i][j].data[k].x = val; } + else { dst.tiles[i][j].data[k].x = src.tiles[i][j].data[k].x; } + if (col_idx_y < col_idx) { dst.tiles[i][j].data[k].y = val; } + else { dst.tiles[i][j].data[k].y = src.tiles[i][j].data[k].y; } + } + } + } +} +template +__device__ static inline void left_fill(RT &dst, const RT &src, const int col_idx, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + const typename RT::dtype packed_val = base_types::packing::pack(val); + + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int thread_col = (j * RT::tile_size_col) + ((k % 2) * 8) + ((laneid() / 4)); + if (thread_col < col_idx) { dst.tiles[i][j].data[k] = packed_val; } + else { dst.tiles[i][j].data[k] = src.tiles[i][j].data[k]; } + } + } + __syncwarp(); + } +} + + + + + + + + + + +template +__device__ static inline void upper_fill(RT &dst, const RT &src, const int row_idx, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + if(row_idx <= 0) return; + const typename RT::dtype packed_val = base_types::packing::pack(val); + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int thread_row = (i * RT::tile_size_row) + ((k % 2) * 8) + ((laneid() / 4)); + if (thread_row < row_idx) { dst.tiles[i][j].data[k] = packed_val; } + else { dst.tiles[i][j].data[k] = src.tiles[i][j].data[k]; } + } + } + } +} +template +__device__ static inline void upper_fill(RT &dst, const RT &src, const int row_idx, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int row_idx_x = (i * RT::tile_size_row) + ((k / 2) * 8) + ((laneid() % 4) * 2); + const int row_idx_y = (i * RT::tile_size_row) + ((k / 2) * 8) + ((laneid() % 4) * 2) + 1; + if (row_idx_x < row_idx) { dst.tiles[i][j].data[k].x = val; } + else { dst.tiles[i][j].data[k].x = src.tiles[i][j].data[k].x; } + if (row_idx_y < row_idx) { dst.tiles[i][j].data[k].y = val; } + else { dst.tiles[i][j].data[k].y = src.tiles[i][j].data[k].y; } + } + } + } +} + + + + + + + + + + +template +__device__ static inline void lower_fill(RT &dst, const RT &src, const int row_idx, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + if(row_idx >= RT::rows) return; + const typename RT::dtype packed_val = base_types::packing::pack(val); + + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int thread_row = (i * RT::tile_size_row) + ((k % 2) * 8) + ((laneid() / 4)); + if (thread_row >= row_idx) { dst.tiles[i][j].data[k] = packed_val; } + else { dst.tiles[i][j].data[k] = src.tiles[i][j].data[k]; } + } + } + } +} +template +__device__ static inline void lower_fill(RT &dst, const RT &src, const int row_idx, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int row_idx_x = (i * RT::tile_size_row) + ((k / 2) * 8) + ((laneid() % 4) * 2); + const int row_idx_y = (i * RT::tile_size_row) + ((k / 2) * 8) + ((laneid() % 4) * 2) + 1; + if (row_idx_x >= row_idx) { dst.tiles[i][j].data[k].x = val; } + else { dst.tiles[i][j].data[k].x = src.tiles[i][j].data[k].x; } + if (row_idx_y >= row_idx) { dst.tiles[i][j].data[k].y = val; } + else { dst.tiles[i][j].data[k].y = src.tiles[i][j].data[k].y; } + } + } + } +} + + + + + + + + + + + + + +template +__device__ static inline void upper_right_fill(RT &dst, const RT &src, const int row_idx, const int col_idx, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + if(col_idx >= RT::cols || row_idx < 0) return; + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int row_idx_xy = (i * RT::tile_size_row) + ((k % 2) * 8) + ((laneid() / 4)); + const int col_idx_x = (j * RT::tile_size_col) + ((k / 2) * 8) + ((laneid() % 4) * 2); + const int col_idx_y = (j * RT::tile_size_col) + ((k / 2) * 8) + ((laneid() % 4) * 2) + 1; + if (col_idx_x >= col_idx && row_idx_xy < row_idx) { dst.tiles[i][j].data[k].x = val; } + else { dst.tiles[i][j].data[k].x = src.tiles[i][j].data[k].x; } + if (col_idx_y >= col_idx && row_idx_xy < row_idx) { dst.tiles[i][j].data[k].y = val; } + else { dst.tiles[i][j].data[k].y = src.tiles[i][j].data[k].y; } + } + } + __syncwarp(); + } +} +template +__device__ static inline void upper_right_fill(RT &dst, const RT &src, const int row_idx, const int col_idx, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + const typename RT::dtype packed_val = base_types::packing::pack(val); + if(col_idx >= RT::cols || row_idx < 0) return; + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int row_idx_x = (i * RT::tile_size_row) + ((k / 2) * 8) + ((laneid() % 4) * 2); + const int row_idx_y = (i * RT::tile_size_row) + ((k / 2) * 8) + ((laneid() % 4) * 2) + 1; + const int col_idx_xy = (j * RT::tile_size_col) + ((k % 2) * 8) + (laneid() / 4); + if (row_idx_x < row_idx && col_idx_xy < col_idx) { dst.tiles[i][j].data[k].x = val; } + else { dst.tiles[i][j].data[k].x = src.tiles[i][j].data[k].x; } + if (row_idx_y < row_idx && col_idx_xy < col_idx) { dst.tiles[i][j].data[k].y = val; } + else { dst.tiles[i][j].data[k].y = src.tiles[i][j].data[k].y; } + } + } + __syncwarp(); + } +} + + + + + + + + + + + + + +template +__device__ static inline rt &subtile_inplace(RT & src, int idx) { + KITTENS_CHECK_WARP + using T = typename RT::T; + static_assert(RT::height % (subtile_rows / TILE_ROW_DIM) == 0, "subtile height should evenly divide tile height."); + return reinterpret_cast&>( + src.tiles[idx*(subtile_rows / TILE_ROW_DIM)] + ); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/maps.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/maps.cuh new file mode 100644 index 0000000..0142037 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/maps.cuh @@ -0,0 +1,890 @@ + + + + + + + + + + + + + + + +template +__device__ static inline void unary_map(T &dst, const T &src) { + #pragma unroll + for(int i = 0; i < T::height; i++) { + #pragma unroll + for(int j = 0; j < T::width; j++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile; k++) { + dst.tiles[i][j].data[k] = op::template op(src.tiles[i][j].data[k]); + } + } + } +} + + + + + + + + + + +template +__device__ static inline void bin_map(T &dst, const T &src, const typename T::dtype ¶m) { + #pragma unroll + for(int i = 0; i < T::height; i++) { + #pragma unroll + for(int j = 0; j < T::width; j++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile; k++) { + dst.tiles[i][j].data[k] = op::template op(src.tiles[i][j].data[k], param); + } + } + } +} + + + + + + + + + +template +__device__ static inline void bin_map(T &dst, const T &src, const typename base_types::packing::unpacked_type ¶m) { + + bin_map(dst, src, base_types::packing::pack(param)); +} + + + + + + + + + +template +__device__ static inline void bin_map(T &dst, const T &lhs, const T &rhs) { + #pragma unroll + for(int i = 0; i < T::height; i++) { + #pragma unroll + for(int j = 0; j < T::width; j++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile; k++) { + dst.tiles[i][j].data[k] = op::template op(lhs.tiles[i][j].data[k], rhs.tiles[i][j].data[k]); + } + } + } +} + +template +__device__ static inline void apply(RT &dst, const RT &src, Lambda &&lambda) { + int row_offset = 0; + if constexpr(GROUP_WARPS > 1) { + row_offset = warpid()*RT::height; + } + static_assert(sizeof(RT::T) != 1, "Cannot apply lambda to 8-bit types"); + if constexpr (ducks::rt::row_layout) { + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for(int k = 0; k < RT::packed_per_tile; k++) { + int row = row_offset + i*TILE_ROW_DIM + (k%2) * (TILE_ROW_DIM/2) + ::kittens::laneid()/4; + int col = j*TILE_COL_DIM + (k/2) * (TILE_COL_DIM/2) + (::kittens::laneid()%4)*2; + dst.tiles[i][j].data[k].x = lambda(row, col+0, src.tiles[i][j].data[k].x); + dst.tiles[i][j].data[k].y = lambda(row, col+1, src.tiles[i][j].data[k].y); + } + } + } + } + else { + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for(int k = 0; k < RT::packed_per_tile; k++) { + int row = row_offset + i*TILE_ROW_DIM + (k/2) * (TILE_ROW_DIM/2) + (::kittens::laneid()%4)*2; + int col = j*TILE_COL_DIM + (k%2) * (TILE_COL_DIM/2) + ::kittens::laneid()/4; + dst.tiles[i][j].data[k].x = lambda(row+0, col, src.tiles[i][j].data[k].x); + dst.tiles[i][j].data[k].y = lambda(row+1, col, src.tiles[i][j].data[k].y); + } + } + } + } +} +template +__device__ static inline RT apply(const RT &src, Lambda &&lambda) { + RT dst; + apply(dst, src, std::forward(lambda)); + return dst; +} + + + + + + + + + + + + + +template +__device__ static inline void row_map(T &dst, const T &src, const V &row_values) { + + static_assert(std::is_same_v::col_vec_layout>); + static_assert(std::is_same_v); + static_assert(V::outer_dim == T::height); + + using dtype = T::dtype; + + #pragma unroll + for(int i = 0; i < T::height; i++) { + dtype packed_top_row = base_types::packing::pack(row_values[i][0].x); + dtype packed_bottom_row = base_types::packing::pack(row_values[i][0].y); + #pragma unroll + for(int j = 0; j < T::width; j++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile; k+=2) { + dst.tiles[i][j].data[k+0] = op::template op(src.tiles[i][j].data[k+0], packed_top_row); + dst.tiles[i][j].data[k+1] = op::template op(src.tiles[i][j].data[k+1], packed_bottom_row); + } + } + } +} + + + + + + + + + + +template +__device__ static inline void row_map(T &dst, const T &src, const V &row_values) { + + static_assert(std::is_same_v); + static_assert(std::is_same_v::col_vec_layout>); + static_assert(V::outer_dim == T::height); + + using dtype = T::dtype; + + #pragma unroll + for(int i = 0; i < T::height; i++) { + #pragma unroll + for(int j = 0; j < T::width; j++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile/2; k++) { + dst.tiles[i][j].data[k+0] = op::template op(src.tiles[i][j].data[k+0], row_values[i][0]); + dst.tiles[i][j].data[k+2] = op::template op(src.tiles[i][j].data[k+2], row_values[i][1]); + } + } + } +} + + + + + + + + + + + + + + + +template +__device__ static inline void row_map(T &dst, const T &a, const T &b, const V &row_values) { + + static_assert(std::is_same_v::col_vec_layout>); + static_assert(std::is_same_v); + static_assert(V::outer_dim == T::height); + + using dtype = T::dtype; + + #pragma unroll + for(int i = 0; i < T::height; i++) { + dtype packed_top_row = base_types::packing::pack(row_values[i][0].x); + dtype packed_bottom_row = base_types::packing::pack(row_values[i][0].y); + #pragma unroll + for(int j = 0; j < T::width; j++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile; k+=2) { + dst.tiles[i][j].data[k+0] = op::template op(a.tiles[i][j].data[k+0], b.tiles[i][j].data[k+0], packed_top_row); + dst.tiles[i][j].data[k+1] = op::template op(a.tiles[i][j].data[k+1], b.tiles[i][j].data[k+1], packed_bottom_row); + } + } + } +} + + + + + + + + + + + +template +__device__ static inline void row_map(T &dst, const T &a, const T &b, const V &row_values) { + + static_assert(std::is_same_v::col_vec_layout>); + static_assert(std::is_same_v); + static_assert(V::outer_dim == T::height); + + using dtype = T::dtype; + + #pragma unroll + for(int i = 0; i < T::height; i++) { + #pragma unroll + for(int j = 0; j < T::width; j++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile/2; k++) { + dst.tiles[i][j].data[k+0] = op::template op(a.tiles[i][j].data[k+0], b.tiles[i][j].data[k+0], row_values[i][0]); + dst.tiles[i][j].data[k+2] = op::template op(a.tiles[i][j].data[k+2], b.tiles[i][j].data[k+2], row_values[i][1]); + } + } + } +} + + + + + + + + + + + + + +template +__device__ static inline void col_map(T &dst, const T &src, const V &col_values) { + KITTENS_CHECK_WARP + + static_assert(std::is_same_v::row_vec_layout>); + static_assert(std::is_same_v); + static_assert(V::outer_dim == T::width); + + using dtype = T::dtype; + + #pragma unroll + for(int j = 0; j < T::width; j++) { + #pragma unroll + for(int i = 0; i < T::height; i++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile/2; k++) { + dst.tiles[i][j].data[k+0] = op::template op(src.tiles[i][j].data[k+0], col_values[j][0]); + dst.tiles[i][j].data[k+2] = op::template op(src.tiles[i][j].data[k+2], col_values[j][1]); + } + } + } +} + + + + + + + + + + +template +__device__ static inline void col_map(T &dst, const T &src, const V &col_values) { + KITTENS_CHECK_WARP + + static_assert(std::is_same_v::row_vec_layout>); + static_assert(std::is_same_v); + static_assert(V::outer_dim == T::width); + + using dtype = T::dtype; + + #pragma unroll + for(int j = 0; j < T::width; j++) { + dtype packed_left_col = base_types::packing::pack(col_values[j][0].x); + dtype packed_right_col = base_types::packing::pack(col_values[j][0].y); + #pragma unroll + for(int i = 0; i < T::height; i++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile; k+=2) { + dst.tiles[i][j].data[k+0] = op::template op(src.tiles[i][j].data[k+0], packed_left_col); + dst.tiles[i][j].data[k+1] = op::template op(src.tiles[i][j].data[k+1], packed_right_col); + } + } + } +} + + + + + + + + + + + + + +template +__device__ static inline void col_map(T &dst, const T &a, const T &b, const V &col_values) { + KITTENS_CHECK_WARP + + static_assert(std::is_same_v::row_vec_layout>); + static_assert(std::is_same_v); + static_assert(V::outer_dim == T::width); + + using dtype = T::dtype; + + #pragma unroll + for(int j = 0; j < T::width; j++) { + #pragma unroll + for(int i = 0; i < T::height; i++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile/2; k++) { + dst.tiles[i][j].data[k+0] = op::template op(a.tiles[i][j].data[k+0], b.tiles[i][j].data[k+0], col_values[j][0]); + dst.tiles[i][j].data[k+2] = op::template op(a.tiles[i][j].data[k+2], b.tiles[i][j].data[k+2], col_values[j][1]); + } + } + } +} + + + + + + + + + + + +template +__device__ static inline void col_map(T &dst, const T &a, const T &b, const V &col_values) { + KITTENS_CHECK_WARP + + static_assert(std::is_same_v); + static_assert(std::is_same_v::row_vec_layout>); + static_assert(V::outer_dim == T::width); + + using dtype = T::dtype; + #pragma unroll + for(int j = 0; j < T::width; j++) { + dtype packed_left_col = base_types::packing::pack(col_values[j][0].x); + dtype packed_right_col = base_types::packing::pack(col_values[j][0].y); + #pragma unroll + for(int i = 0; i < T::height; i++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile; k+=2) { + dst.tiles[i][j].data[k+0] = op::template op(a.tiles[i][j].data[k+0], b.tiles[i][j].data[k+0], packed_left_col); + dst.tiles[i][j].data[k+1] = op::template op(a.tiles[i][j].data[k+1], b.tiles[i][j].data[k+1], packed_right_col); + } + } + } +} + + + + + + + + + + + + + +template +__device__ static inline void zero(T &dst) { + unary_map(dst, dst); +} + + + + + + +template +__device__ static inline void one(T &dst) { + unary_map(dst, dst); +} + + + + + + +template +__device__ static inline void pos_infty(T &dst) { + unary_map(dst, dst); +} + + + + + + +template +__device__ static inline void neg_infty(T &dst) { + unary_map(dst, dst); +} + + + + + + + + +template +__device__ static inline void exp(T &dst, const T &src) { + unary_map(dst, src); +} +template +__device__ static inline T exp(const T &src) { + T dst; + exp(dst, src); + return dst; +} + + + + + + + + +template +__device__ static inline void exp2(T &dst, const T &src) { + unary_map(dst, src); +} +template +__device__ static inline T exp2(const T &src) { + T dst; + exp2(dst, src); + return dst; +} + + + + + + + + +template +__device__ static inline void log(T &dst, const T &src) { + unary_map(dst, src); +} +template +__device__ static inline T log(const T &src) { + T dst; + log(dst, src); + return dst; +} + + + + + + + + +template +__device__ static inline void log2(T &dst, const T &src) { + unary_map(dst, src); +} +template +__device__ static inline T log2(const T &src) { + T dst; + log2(dst, src); + return dst; +} + + + + + + + + +template +__device__ static inline void abs(T &dst, const T &src) { + unary_map(dst, src); +} +template +__device__ static inline T abs(const T &src) { + T dst; + abs(dst, src); + return dst; +} + + + + + + + + +template +__device__ static inline void relu(T &dst, const T &src) { + unary_map(dst, src); +} +template +__device__ static inline T relu(const T &src) { + T dst; + relu(dst, src); + return dst; +} + + + + + + + + +template +__device__ static inline void neg(T &dst, const T &src) { + unary_map(dst, src); +} +template +__device__ static inline T neg(const T &src) { + T dst; + neg(dst, src); + return dst; +} + + + + + + + + +template +__device__ static inline void sqrt(T &dst, const T &src) { + unary_map(dst, src); +} +template +__device__ static inline T sqrt(const T &src) { + T dst; + sqrt(dst, src); + return dst; +} + + + + + + + + +template +__device__ static inline void rsqrt(T &dst, const T &src) { + unary_map(dst, src); +} +template +__device__ static inline T rsqrt(const T &src) { + T dst; + rsqrt(dst, src); + return dst; +} + + + + + + + + + +template +__device__ static inline void copy(T &dst, const U &src) { + bin_map(dst, src); +} + + + + + + + + + + +template +__device__ static inline void max(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} +template +__device__ static inline T max(const T &lhs, const U &rhs) { + T dst; + max(dst, lhs, rhs); + return dst; +} + + + + + + + + + + +template +__device__ static inline void min(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} +template +__device__ static inline T min(const T &lhs, const U &rhs) { + T dst; + min(dst, lhs, rhs); + return dst; +} + + + + + + + + + + +template +__device__ static inline void add(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + + + + + + + + + + +template +__device__ static inline void sub(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void mul(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + + + + + + + + + + +template +__device__ static inline void div(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + + + + + + + + + + +template +__device__ static inline void add_row(T &dst, const T &src, const V &row_values) { + row_map(dst, src, row_values); +} + + + + + + + + + + +template +__device__ static inline void sub_row(T &dst, const T &src, const V &row_values) { + row_map(dst, src, row_values); +} + + + + + + + + + + +template +__device__ static inline void mul_row(T &dst, const T &src, const V &row_values) { + row_map(dst, src, row_values); +} + + + + + + + + + + +template +__device__ static inline void div_row(T &dst, const T &src, const V &row_values) { + row_map(dst, src, row_values); +} + + + + + + + + + +template +__device__ static inline void broadcast_row(T &dst, const V &row_values) { + row_map(dst, dst, row_values); +} +template +__device__ static inline T broadcast_row(const V &row_values) { + T dst; + broadcast_row(dst, row_values); + return dst; +} + + + + + + + + + + + + +template +__device__ static inline void add_col(T &dst, const T &src, const V &col_values) { + col_map(dst, src, col_values); +} + + + + + + + + + + +template +__device__ static inline void sub_col(T &dst, const T &src, const V &col_values) { + col_map(dst, src, col_values); +} + + + + + + + + + + +template +__device__ static inline void mul_col(T &dst, const T &src, const V &col_values) { + col_map(dst, src, col_values); +} + + + + + + + + + + +template +__device__ static inline void div_col(T &dst, const T &src, const V &col_values) { + col_map(dst, src, col_values); +} + + + + + + + + + +template +__device__ static inline void broadcast_col(T &dst, const V &col_values) { + col_map(dst, dst, col_values); +} +template +__device__ static inline T broadcast_col(const V &col_values) { + T dst; + broadcast_col(dst, col_values); + return dst; +} + + +template +__device__ static inline void tril(RT &dst, const RT &src, int diagonal=0, const typename base_types::packing::unpacked_type &val=0) { + apply(dst, src, [val, diagonal]__device__(int row, int col, auto &src_val) { + return col <= row + diagonal ? src_val : val; + }); +} +template +__device__ static inline void triu(RT &dst, const RT &src, int diagonal=0, const typename base_types::packing::unpacked_type &val=0) { + apply(dst, src, [val, diagonal]__device__(int row, int col, auto &src_val) { + return col >= row + diagonal ? src_val : val; + }); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/reductions.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/reductions.cuh new file mode 100644 index 0000000..16322b1 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/reductions.cuh @@ -0,0 +1,540 @@ + + + + + + + + + + + + + + + + + + + +template +__device__ static inline void row_reduce(V &row_accum, const T &src, const V &src_accum) { + + static_assert(std::is_same_v::col_vec_layout>); + static_assert(std::is_same_v); + static_assert(V::outer_dim == T::height); + + using dtype = V::dtype; + + #pragma unroll + for(int i = 0; i < T::height; i++) { + dtype accum_top_row = op::template op(src.tiles[i][0].data[0], src.tiles[i][0].data[2]); + dtype accum_bottom_row = op::template op(src.tiles[i][0].data[1], src.tiles[i][0].data[3]); + #pragma unroll + for(int j = 1; j < T::width; j++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile; k+=2) { + accum_top_row = op::template op(accum_top_row, src.tiles[i][j].data[k+0]); + accum_bottom_row = op::template op(accum_bottom_row, src.tiles[i][j].data[k+1]); + } + } + dtype accum_packed; + accum_packed.x = op::template op::unpacked_type>(accum_top_row.x, accum_top_row.y); + accum_packed.y = op::template op::unpacked_type>(accum_bottom_row.x, accum_bottom_row.y); + + + + accum_packed = op::template op(accum_packed, packed_shfl_xor_sync(MASK_ALL, accum_packed, 2)); + accum_packed = op::template op(accum_packed, packed_shfl_xor_sync(MASK_ALL, accum_packed, 1)); + + if(reset) { + row_accum[i][0] = accum_packed; + } + else { + row_accum[i][0] = op::template op(src_accum[i][0], accum_packed); + } + } +} + + + + + + + + + + + + + + +template +__device__ static inline void row_reduce(V &row_accum, const T &src, const V &src_accum) { + + static_assert(std::is_same_v::col_vec_layout>); + static_assert(std::is_same_v); + static_assert(V::outer_dim == T::height); + + using dtype = V::dtype; + + #pragma unroll + for(int i = 0; i < T::height; i++) { + dtype accum_top_rows = op::template op(src.tiles[i][0].data[0], src.tiles[i][0].data[1]); + dtype accum_bottom_rows = op::template op(src.tiles[i][0].data[2], src.tiles[i][0].data[3]); + #pragma unroll + for(int j = 1; j < T::width; j++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile/2; k++) { + accum_top_rows = op::template op(accum_top_rows, src.tiles[i][j].data[k+0]); + accum_bottom_rows = op::template op(accum_bottom_rows, src.tiles[i][j].data[k+2]); + } + } + + + + accum_top_rows = op::template op(accum_top_rows, packed_shfl_xor_sync(MASK_ALL, accum_top_rows, 16)); + accum_top_rows = op::template op(accum_top_rows, packed_shfl_xor_sync(MASK_ALL, accum_top_rows, 8)); + accum_top_rows = op::template op(accum_top_rows, packed_shfl_xor_sync(MASK_ALL, accum_top_rows, 4)); + + accum_bottom_rows = op::template op(accum_bottom_rows, packed_shfl_xor_sync(MASK_ALL, accum_bottom_rows, 16)); + accum_bottom_rows = op::template op(accum_bottom_rows, packed_shfl_xor_sync(MASK_ALL, accum_bottom_rows, 8)); + accum_bottom_rows = op::template op(accum_bottom_rows, packed_shfl_xor_sync(MASK_ALL, accum_bottom_rows, 4)); + + if(reset) { + row_accum[i][0] = accum_top_rows; + row_accum[i][1] = accum_bottom_rows; + } + else { + row_accum[i][0] = op::template op(src_accum[i][0], accum_top_rows); + row_accum[i][1] = op::template op(src_accum[i][1], accum_bottom_rows); + } + } +} + + + + + + + + + + + + + + + + +template +__device__ static inline void col_reduce(V &col_accum, const T &src, const V &src_accum) { + + KITTENS_CHECK_WARP + static_assert(std::is_same_v::row_vec_layout>); + static_assert(std::is_same_v); + static_assert(V::outer_dim == T::width); + + using dtype = V::dtype; + + #pragma unroll + for(int j = 0; j < T::width; j++) { + dtype accum_left_cols = op::template op(src.tiles[0][j].data[0], src.tiles[0][j].data[1]); + dtype accum_right_cols = op::template op(src.tiles[0][j].data[2], src.tiles[0][j].data[3]); + #pragma unroll + for(int i = 1; i < T::height; i++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile/2; k++) { + accum_left_cols = op::template op(accum_left_cols, src.tiles[i][j].data[k+0]); + accum_right_cols = op::template op(accum_right_cols, src.tiles[i][j].data[k+2]); + } + } + + + + accum_left_cols = op::template op(accum_left_cols, packed_shfl_xor_sync(MASK_ALL, accum_left_cols, 16)); + accum_left_cols = op::template op(accum_left_cols, packed_shfl_xor_sync(MASK_ALL, accum_left_cols, 8)); + accum_left_cols = op::template op(accum_left_cols, packed_shfl_xor_sync(MASK_ALL, accum_left_cols, 4)); + + accum_right_cols = op::template op(accum_right_cols, packed_shfl_xor_sync(MASK_ALL, accum_right_cols, 16)); + accum_right_cols = op::template op(accum_right_cols, packed_shfl_xor_sync(MASK_ALL, accum_right_cols, 8)); + accum_right_cols = op::template op(accum_right_cols, packed_shfl_xor_sync(MASK_ALL, accum_right_cols, 4)); + + if(reset) { + col_accum[j][0] = accum_left_cols; + col_accum[j][1] = accum_right_cols; + } + else { + col_accum[j][0] = op::template op(src_accum[j][0], accum_left_cols); + col_accum[j][1] = op::template op(src_accum[j][1], accum_right_cols); + } + } +} + + + + + + + + + + + + + + +template +__device__ static inline void col_reduce(V &col_accum, const T &src, const V &src_accum) { + + KITTENS_CHECK_WARP + static_assert(std::is_same_v::row_vec_layout>); + static_assert(std::is_same_v); + static_assert(V::outer_dim == T::width); + + using dtype = V::dtype; + #pragma unroll + for(int j = 0; j < T::width; j++) { + dtype accum_left_col = op::template op(src.tiles[0][j].data[0], src.tiles[0][j].data[2]); + dtype accum_right_col = op::template op(src.tiles[0][j].data[1], src.tiles[0][j].data[3]); + #pragma unroll + for(int i = 1; i < T::height; i++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile; k+=2) { + accum_left_col = op::template op(accum_left_col, src.tiles[i][j].data[k+0]); + accum_right_col = op::template op(accum_right_col, src.tiles[i][j].data[k+1]); + } + } + dtype accum_packed; + accum_packed.x = op::template op::unpacked_type>(accum_left_col.x, accum_left_col.y); + accum_packed.y = op::template op::unpacked_type>(accum_right_col.x, accum_right_col.y); + + + + accum_packed = op::template op(accum_packed, packed_shfl_xor_sync(MASK_ALL, accum_packed, 2)); + accum_packed = op::template op(accum_packed, packed_shfl_xor_sync(MASK_ALL, accum_packed, 1)); + + if(reset) { + col_accum[j][0] = accum_packed; + } + else { + col_accum[j][0] = op::template op(src_accum[j][0], accum_packed); + } + } +} + + + + + + + + + + + + + +template +__device__ static inline void row_max(V &row_accum, const T &src) { + row_reduce(row_accum, src, row_accum); +} + + + + + + + + +template +__device__ static inline void row_min(V &row_accum, const T &src) { + row_reduce(row_accum, src, row_accum); +} + + + + + + + + +template +__device__ static inline void row_sum(V &row_accum, const T &src) { + row_reduce(row_accum, src, row_accum); +} + + + + + + + + +template +__device__ static inline void row_prod(V &row_accum, const T &src) { + row_reduce(row_accum, src, row_accum); +} + + + + + + + + + + +template +__device__ static inline void row_max(V &row_accum, const T &src, const V &src_accum) { + row_reduce(row_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void row_min(V &row_accum, const T &src, const V &src_accum) { + row_reduce(row_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void row_sum(V &row_accum, const T &src, const V &src_accum) { + row_reduce(row_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void row_prod(V &row_accum, const T &src, const V &src_accum) { + row_reduce(row_accum, src, src_accum); +} + + + + + + + + + + + +template +__device__ static inline void col_max(V &col_accum, const T &src) { + col_reduce(col_accum, src, col_accum); +} + + + + + + + + +template +__device__ static inline void col_min(V &col_accum, const T &src) { + col_reduce(col_accum, src, col_accum); +} + + + + + + + + +template +__device__ static inline void col_sum(V &col_accum, const T &src) { + col_reduce(col_accum, src, col_accum); +} + + + + + + + + +template +__device__ static inline void col_prod(V &col_accum, const T &src) { + col_reduce(col_accum, src, col_accum); +} + + + + + + + + + + +template +__device__ static inline void col_max(V &col_accum, const T &src, const V &src_accum) { + col_reduce(col_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void col_min(V &col_accum, const T &src, const V &src_accum) { + col_reduce(col_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void col_sum(V &col_accum, const T &src, const V &src_accum) { + col_reduce(col_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void col_prod(V &col_accum, const T &src, const V &src_accum) { + col_reduce(col_accum, src, src_accum); +} + + + +template +__device__ static inline void max(RV &dst, const T &src, const RV &src_accum) { + if constexpr (ax == axis::COL) row_max(dst, src, src_accum); + else col_max(dst, src, src_accum); +} +template +__device__ static inline auto max(const T &src, const RV &src_accum) { + RV dst; + if constexpr (ax == axis::COL) row_max(dst, src, src_accum); + else col_max(dst, src, src_accum); + return dst; +} +template +__device__ static inline void max(RV &dst, const T &src) { + if constexpr (ax == axis::COL) row_max(dst, src); + else col_max(dst, src); +} +template +__device__ static inline auto max(const T &src) { + using RV = std::conditional_t; + RV dst; + if constexpr (ax == axis::COL) row_max(dst, src); + else col_max(dst, src); + return dst; +} + +template +__device__ static inline void min(RV &dst, const T &src, const RV &src_accum) { + if constexpr (ax == axis::COL) row_min(dst, src, src_accum); + else col_min(dst, src, src_accum); +} +template +__device__ static inline auto min(const T &src, const RV &src_accum) { + RV dst; + if constexpr (ax == axis::COL) row_min(dst, src, src_accum); + else col_min(dst, src, src_accum); + return dst; +} +template +__device__ static inline void min(RV &dst, const T &src) { + if constexpr (ax == axis::COL) row_min(dst, src); + else col_min(dst, src); +} +template +__device__ static inline auto min(const T &src) { + using RV = std::conditional_t; + RV dst; + if constexpr (ax == axis::COL) row_min(dst, src); + else col_min(dst, src); + return dst; +} + +template +__device__ static inline void sum(RV &dst, const T &src, const RV &src_accum) { + if constexpr (ax == axis::COL) row_sum(dst, src, src_accum); + else col_sum(dst, src, src_accum); +} +template +__device__ static inline auto sum(const T &src, const RV &src_accum) { + RV dst; + if constexpr (ax == axis::COL) row_sum(dst, src, src_accum); + else col_sum(dst, src, src_accum); + return dst; +} +template +__device__ static inline void sum(RV &dst, const T &src) { + if constexpr (ax == axis::COL) row_sum(dst, src); + else col_sum(dst, src); +} +template +__device__ static inline auto sum(const T &src) { + using RV = std::conditional_t; + RV dst; + if constexpr (ax == axis::COL) row_sum(dst, src); + else col_sum(dst, src); + return dst; +} + +template +__device__ static inline void prod(RV &dst, const T &src, const RV &src_accum) { + if constexpr (ax == axis::COL) row_prod(dst, src, src_accum); + else col_prod(dst, src, src_accum); +} +template +__device__ static inline auto prod(const T &src, const RV &src_accum) { + RV dst; + if constexpr (ax == axis::COL) row_prod(dst, src, src_accum); + else col_prod(dst, src, src_accum); + return dst; +} +template +__device__ static inline void prod(RV &dst, const T &src) { + if constexpr (ax == axis::COL) row_prod(dst, src); + else col_prod(dst, src); +} +template +__device__ static inline auto prod(const T &src) { + using RV = std::conditional_t; + RV dst; + if constexpr (ax == axis::COL) row_prod(dst, src); + else col_prod(dst, src); + return dst; +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/tile.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/tile.cuh new file mode 100644 index 0000000..5d736b6 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/tile.cuh @@ -0,0 +1,46 @@ + + + + + +#include "conversions.cuh" +#include "maps.cuh" +#include "reductions.cuh" + +template +__device__ static inline bool hasnan(const RT &src) { + KITTENS_CHECK_WARP + bool nan_detected = false; + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for(int k = 0; k < RT::packed_per_tile; k++) { + if constexpr (std::is_same_v) { + if(isnan(src.tiles[i][j].data[k].x) || isnan(src.tiles[i][j].data[k].y)) { + nan_detected = true; + } + } + else if constexpr (std::is_same_v) { + if(isnan(__bfloat162float(src.tiles[i][j].data[k].x)) || isnan(__bfloat162float(src.tiles[i][j].data[k].y))) { + nan_detected = true; + } + } + else if constexpr (std::is_same_v) { + if(isnan(__half2float(src.tiles[i][j].data[k].x)) || isnan(__half2float(src.tiles[i][j].data[k].y))) { + nan_detected = true; + } + } + else { + static_assert(sizeof(typename RT::T) == 999, "Unsupported dtype"); + } + } + } + } + + return (__ballot_sync(0xffffffff, nan_detected) != 0); +} + +#include "complex_conversions.cuh" +#include "complex_maps.cuh" diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/vec/conversions.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/vec/conversions.cuh new file mode 100644 index 0000000..ab37d25 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/vec/conversions.cuh @@ -0,0 +1,153 @@ + + + + + +struct vec_conversion_detail { + + + +__device__ static inline int row_from_indices_dim2(int laneid, int inner_dim, int x_or_y) { + return 8*inner_dim + (laneid%4)*2 + x_or_y; +} +__device__ static inline int row_from_indices_dim1(int laneid, int x_or_y) { + return 8*x_or_y + (laneid/4); +} +__device__ static inline int canonical_src_lane_dim2(int row) { + return (row/2)%4 + 4*(row%2); +} +__device__ static inline int canonical_src_lane_dim1(int row) { + return (row*4)%32; +} + +}; + + + + + + + + + +template +__device__ static inline void copy(RV1 &dst, const RV2 &src) { + KITTENS_CHECK_WARP + static_assert(RV1::length == RV2::length, "Register vectors must be the same length."); + using D1 = RV1::dtype; + using D2 = RV2::dtype; + if constexpr (std::is_same_v) { + #pragma unroll + for(int i = 0; i < RV1::outer_dim; i++) { + #pragma unroll + for(int j = 0; j < RV1::inner_dim; j++) { + dst[i][j] = base_types::convertor::convert(src[i][j]); + } + } + } + else { + int laneid = ::kittens::laneid(); + if constexpr (std::is_same_v && std::is_same_v) { + #pragma unroll + for(int i = 0; i < RV1::outer_dim; i++) { + dst[i][0].x = packed_shfl_sync( + kittens::MASK_ALL, + laneid < 4 ? src[i][0].x : src[i][0].y, + vec_conversion_detail::canonical_src_lane_dim2(vec_conversion_detail::row_from_indices_dim1(laneid, 0)) + ); + dst[i][0].y = packed_shfl_sync( + kittens::MASK_ALL, + laneid < 4 ? src[i][1].x : src[i][1].y, + vec_conversion_detail::canonical_src_lane_dim2(vec_conversion_detail::row_from_indices_dim1(laneid, 1)) + ); + } + } + else if constexpr (std::is_same_v && std::is_same_v) { + #pragma unroll + for(int i = 0; i < RV1::outer_dim; i++) { + dst[i][0].x = packed_shfl_sync( + kittens::MASK_ALL, + src[i][0].x, + vec_conversion_detail::canonical_src_lane_dim1(vec_conversion_detail::row_from_indices_dim2(laneid, 0, 0)) + ); + dst[i][0].y = packed_shfl_sync( + kittens::MASK_ALL, + src[i][0].x, + vec_conversion_detail::canonical_src_lane_dim1(vec_conversion_detail::row_from_indices_dim2(laneid, 0, 1)) + ); + dst[i][1].x = packed_shfl_sync( + kittens::MASK_ALL, + src[i][0].y, + vec_conversion_detail::canonical_src_lane_dim1(vec_conversion_detail::row_from_indices_dim2(laneid, 1, 0)) + ); + dst[i][1].y = packed_shfl_sync( + kittens::MASK_ALL, + src[i][0].y, + vec_conversion_detail::canonical_src_lane_dim1(vec_conversion_detail::row_from_indices_dim2(laneid, 1, 1)) + ); + } + } + else if constexpr (std::is_same_v && std::is_same_v) { + #pragma unroll + for(int i = 0; i < RV1::outer_dim; i++) { + dst[i][0].x = packed_shfl_sync( + kittens::MASK_ALL, src[i/2][0], + 16*(i%2) + 0 + (laneid/4) + ); + dst[i][0].y = packed_shfl_sync( + kittens::MASK_ALL, src[i/2][0], + 16*(i%2) + 8 + (laneid/4) + ); + } + } + else if constexpr (std::is_same_v && std::is_same_v) { + int lane_replication = laneid%4; + #pragma unroll + for(int i = 0; i < RV1::outer_dim; i++) { + D1 tmp = 0; + if(RV1::length%32==0 || i < RV1::outer_dim-1 || lane_replication<2) { + tmp = lane_replication%2 ? src[2*i + (lane_replication>=2)][0].y : src[2*i + (lane_replication>=2)][0].x; + } + dst[i][0] = packed_shfl_sync( + kittens::MASK_ALL, tmp, + (laneid%8)*4 + (laneid/8) + ); + } + } + else if constexpr (std::is_same_v && std::is_same_v) { + #pragma unroll + for(int i = 0; i < RV1::outer_dim; i++) { + dst[i][0].x = packed_shfl_sync( + kittens::MASK_ALL, src[i/2][0], + 16*(i%2) + 0 + 2*(laneid%4) + 0 + ); + dst[i][0].y = packed_shfl_sync( + kittens::MASK_ALL, src[i/2][0], + 16*(i%2) + 0 + 2*(laneid%4) + 1 + ); + dst[i][1].x = packed_shfl_sync( + kittens::MASK_ALL, src[i/2][0], + 16*(i%2) + 8 + 2*(laneid%4) + 0 + ); + dst[i][1].y = packed_shfl_sync( + kittens::MASK_ALL, src[i/2][0], + 16*(i%2) + 8 + 2*(laneid%4) + 1 + ); + } + } + else if constexpr (std::is_same_v && std::is_same_v) { + int lane_replication = laneid/8; + #pragma unroll + for(int i = 0; i < RV1::outer_dim; i++) { + D1 tmp = 0; + if(RV1::length%32==0 || i < RV1::outer_dim-1 || laneid<16) { + tmp = (laneid%8)<4 ? src[2*i + (lane_replication>=2)][lane_replication%2].x : src[2*i + (lane_replication>=2)][lane_replication%2].y; + } + dst[i][0] = packed_shfl_sync( + kittens::MASK_ALL, tmp, + 4*(laneid%2) + (laneid%8)/2 + (laneid&0b11000) + ); + } + } + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/vec/maps.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/vec/maps.cuh new file mode 100644 index 0000000..91fefb2 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/vec/maps.cuh @@ -0,0 +1,374 @@ + + + + + + + + + + + + + + + +template +__device__ static inline void unary_op(T &dst, const T &src) { + #pragma unroll + for(int i = 0; i < T::outer_dim; i++) { + #pragma unroll + for(int j = 0; j < T::inner_dim; j++) { + dst[i][j] = op::template op(src[i][j]); + } + } +} + + + + + + + + + +template +__device__ static inline void bin_op(T &dst, const T &lhs, const T &rhs) { + #pragma unroll + for(int i = 0; i < T::outer_dim; i++) { + #pragma unroll + for(int j = 0; j < T::inner_dim; j++) { + dst[i][j] = op::template op(lhs[i][j], rhs[i][j]); + } + } +} + + + + + + + + + +template +__device__ static inline void bin_op(T &dst, const T &src, const typename T::dtype ¶m) { + #pragma unroll + for(int i = 0; i < T::outer_dim; i++) { + #pragma unroll + for(int j = 0; j < T::inner_dim; j++) { + dst[i][j] = op::template op(src[i][j], param); + } + } +} + + + + + + + + + +template +__device__ static inline void bin_op(T &dst, const T &src, const typename base_types::packing::unpacked_type ¶m) { + bin_op(dst, src, base_types::packing::pack(param)); +} + + +template +__device__ static inline void apply(RV &dst, const RV &src, Lambda &&lambda) { + int group_offset = 0; + if constexpr(GROUP_WARPS > 1) { + group_offset = warpid()*RV::length; + } + static_assert(sizeof(RV::T) != 1, "Cannot apply lambda to 8-bit types"); + if constexpr (ducks::rv::ortho_layout) { + #pragma unroll + for(int i = 0; i < RV::outer_dim; i++) { + int base_idx = group_offset + i*16 + ::kittens::laneid()/4; + dst[i][0].x = lambda(base_idx+0, src[i][0].x); + dst[i][0].y = lambda(base_idx+8, src[i][0].y); + } + } + else if constexpr (ducks::rv::align_layout) { + #pragma unroll + for(int i = 0; i < RV::outer_dim; i++) { + int base_idx = group_offset + i*16 + 2*(::kittens::laneid()%4); + dst[i][0].x = lambda(base_idx+0, src[i][0].x); + dst[i][0].y = lambda(base_idx+1, src[i][0].y); + dst[i][1].x = lambda(base_idx+8, src[i][1].x); + dst[i][1].y = lambda(base_idx+9, src[i][1].y); + } + } + else { + #pragma unroll + for(int i = 0; i < RV::outer_dim; i++) { + int base_idx = group_offset + i*32 + ::kittens::laneid(); + if (i < RV::outer_dim-1 || RV::length%32 == 0 || ::kittens::laneid()<16) { + dst[i][0] = lambda(base_idx, src[i][0]); + } + } + } +} +template +__device__ static inline RV apply(const RV &src, Lambda &&lambda) { + RV dst; + apply(dst, src, std::forward(lambda)); + return dst; +} + + + + + + + + + + + +template +__device__ static inline void zero(T &dst) { + unary_op(dst, dst); +} + + + + + + +template +__device__ static inline void one(T &dst) { + unary_op(dst, dst); +} + + + + + + +template +__device__ static inline void pos_infty(T &dst) { + unary_op(dst, dst); +} + + + + + + +template +__device__ static inline void neg_infty(T &dst) { + unary_op(dst, dst); +} + + + + + + + + + + + +template +__device__ static inline void copy(T &dst, const U &src) { + bin_op(dst, dst, src); +} + + + + + + + +template +__device__ static inline void exp(T &dst, const T &src) { + unary_op(dst, src); +} +template +__device__ static inline T exp(const T &src) { + T dst; + exp(dst, src); + return dst; +} + + + + + + + +template +__device__ static inline void exp2(T &dst, const T &src) { + unary_op(dst, src); +} +template +__device__ static inline T exp2(const T &src) { + T dst; + exp2(dst, src); + return dst; +} + + + + + + + +template +__device__ static inline void log(T &dst, const T &src) { + unary_op(dst, src); +} +template +__device__ static inline T log(const T &src) { + T dst; + log(dst, src); + return dst; +} + + + + + + + +template +__device__ static inline void log2(T &dst, const T &src) { + unary_op(dst, src); +} +template +__device__ static inline T log2(const T &src) { + T dst; + log2(dst, src); + return dst; +} + + + + + + + +template +__device__ static inline void abs(T &dst, const T &src) { + unary_op(dst, src); +} +template +__device__ static inline T abs(const T &src) { + T dst; + abs(dst, src); + return dst; +} + + + + + + + +template +__device__ static inline void relu(T &dst, const T &src) { + unary_op(dst, src); +} +template +__device__ static inline T relu(const T &src) { + T dst; + relu(dst, src); + return dst; +} + + + + + + + + + + + + +template +__device__ static inline void max(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} +template +__device__ static inline T max(const T &lhs, const U &rhs) { + T dst; + max(dst, lhs, rhs); + return dst; +} + + + + + + + + + +template +__device__ static inline void min(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} +template +__device__ static inline T min(const T &lhs, const U &rhs) { + T dst; + min(dst, lhs, rhs); + return dst; +} + + + + + + + + + +template +__device__ static inline void add(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void sub(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void mul(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void div(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/vec/reductions.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/vec/reductions.cuh new file mode 100644 index 0000000..62808d0 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/vec/reductions.cuh @@ -0,0 +1,233 @@ + + + + + + + + + + + + + + + + + + + + + +template +__device__ static inline void reduce( + typename base_types::packing::unpacked_type &dst_accum, + const RV &src, + const typename base_types::packing::unpacked_type &src_accum) { + KITTENS_CHECK_WARP + using T = base_types::packing::unpacked_type; + int laneid = kittens::laneid(); + if constexpr (std::is_same_v) { + T accum = op::template op(src[0][0].x, src[0][0].y); + #pragma unroll + for(int i = 1; i < RV::outer_dim; i++) { + accum = op::template op(accum, src[i][0].x); + accum = op::template op(accum, src[i][0].y); + } + + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 16)); + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 8)); + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 4)); + + if constexpr (!reset) accum = op::template op(accum, src_accum); + + dst_accum = packed_shfl_sync(kittens::MASK_ALL, accum, 0); + } + else if constexpr (std::is_same_v) { + T accum = op::template op(src[0][0].x, src[0][0].y); + accum = op::template op(accum, src[0][1].x); + accum = op::template op(accum, src[0][1].y); + #pragma unroll + for(int i = 1; i < RV::outer_dim; i++) { + + accum = op::template op(accum, src[i][0].x); + accum = op::template op(accum, src[i][0].y); + accum = op::template op(accum, src[i][1].x); + accum = op::template op(accum, src[i][1].y); + } + + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 2)); + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 1)); + + if constexpr (!reset) accum = op::template op(accum, src_accum); + + dst_accum = packed_shfl_sync(kittens::MASK_ALL, accum, 0); + } + else if constexpr (std::is_same_v) { + T accum = src[0][0]; + #pragma unroll + for(int i = 1; i < RV::outer_dim; i++) { + if (i < RV::outer_dim-1 || i*kittens::TILE_ROW_DIM*2 + laneid < RV::length) { + accum = op::template op(accum, src[i][0]); + } + } + if(RV::length > 16) accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 16)); + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 8)); + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 4)); + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 2)); + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 1)); + if constexpr (!reset) accum = op::template op(accum, src_accum); + dst_accum = packed_shfl_sync(kittens::MASK_ALL, accum, 0); + } +} + + + + + + + + + +template +__device__ static inline void max(typename base_types::packing::unpacked_type &max_val, const RV &src) { + reduce(max_val, src, max_val); +} +template +__device__ static inline typename base_types::packing::unpacked_type max(const RV &src) { + typename base_types::packing::unpacked_type max_val; + reduce(max_val, src, max_val); + return max_val; +} + + + + + + + + +template +__device__ static inline void min(typename base_types::packing::unpacked_type &min_val, const RV &src) { + reduce(min_val, src, min_val); +} +template +__device__ static inline typename base_types::packing::unpacked_type min(const RV &src) { + typename base_types::packing::unpacked_type min_val; + reduce(min_val, src, min_val); + return min_val; +} + + + + + + + + +template +__device__ static inline void sum(typename base_types::packing::unpacked_type &sum_val, const RV &src) { + reduce(sum_val, src, sum_val); +} +template +__device__ static inline typename base_types::packing::unpacked_type sum(const RV &src) { + typename base_types::packing::unpacked_type sum_val; + reduce(sum_val, src, sum_val); + return sum_val; +} + + + + + + + + +template +__device__ static inline void prod(typename base_types::packing::unpacked_type &prod_val, const RV &src) { + reduce(prod_val, src, prod_val); +} +template +__device__ static inline typename base_types::packing::unpacked_type prod(const RV &src) { + typename base_types::packing::unpacked_type prod_val; + reduce(prod_val, src, prod_val); + return prod_val; +} + + + + + + + + + + + +template +__device__ static inline void max(typename base_types::packing::unpacked_type &max_val, const RV &src, const typename base_types::packing::unpacked_type &src_accum) { + reduce(max_val, src, src_accum); +} +template +__device__ static inline typename base_types::packing::unpacked_type max(const RV &src, const typename base_types::packing::unpacked_type &src_accum) { + typename base_types::packing::unpacked_type max_val; + reduce(max_val, src, src_accum); + return max_val; +} + + + + + + + + + +template +__device__ static inline void min(typename base_types::packing::unpacked_type &min_val, const RV &src, const typename base_types::packing::unpacked_type &src_accum) { + reduce(min_val, src, src_accum); +} +template +__device__ static inline typename base_types::packing::unpacked_type min(const RV &src, const typename base_types::packing::unpacked_type &src_accum) { + typename base_types::packing::unpacked_type min_val; + reduce(min_val, src, src_accum); + return min_val; +} + + + + + + + + + +template +__device__ static inline void sum(typename base_types::packing::unpacked_type &sum_val, const RV &src, const typename base_types::packing::unpacked_type &src_accum) { + reduce(sum_val, src, src_accum); +} +template +__device__ static inline typename base_types::packing::unpacked_type sum(const RV &src, const typename base_types::packing::unpacked_type &src_accum) { + typename base_types::packing::unpacked_type sum_val; + reduce(sum_val, src, src_accum); + return sum_val; +} + + + + + + + + + +template +__device__ static inline void prod(typename base_types::packing::unpacked_type &prod_val, const RV &src, const typename base_types::packing::unpacked_type &src_accum) { + reduce(prod_val, src, src_accum); +} +template +__device__ static inline typename base_types::packing::unpacked_type prod(const RV &src, const typename base_types::packing::unpacked_type &src_accum) { + typename base_types::packing::unpacked_type prod_val; + reduce(prod_val, src, src_accum); + return prod_val; +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/vec/vec.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/vec/vec.cuh new file mode 100644 index 0000000..fa67f5d --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/vec/vec.cuh @@ -0,0 +1,59 @@ + + + + + +#include "conversions.cuh" +#include "maps.cuh" +#include "reductions.cuh" + +template +__device__ static inline bool hasnan(const RV &src) { + KITTENS_CHECK_WARP + bool nan_detected = false; + #pragma unroll + for(int i = 0; i < RV::outer_dim; i++) { + #pragma unroll + for(int j = 0; j < RV::inner_dim; j++) { + if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { + if(isnan(src[i][j].x) || isnan(src[i][j].y)) { + nan_detected = true; + } + } + else if constexpr (std::is_same_v) { + if(isnan(__bfloat162float(src[i][j].x)) || isnan(__bfloat162float(src[i][j].y))) { + nan_detected = true; + } + } + else if constexpr (std::is_same_v) { + if(isnan(__half2float(src[i][j].x)) || isnan(__half2float(src[i][j].y))) { + nan_detected = true; + } + } + } + else if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { + if(isnan(src[i][j])) { + nan_detected = true; + } + } + else if constexpr (std::is_same_v) { + if(isnan(__bfloat162float(src[i][j]))) { + nan_detected = true; + } + } + else if constexpr (std::is_same_v) { + if(isnan(__half2float(src[i][j]))) { + nan_detected = true; + } + } + } + else { + static_assert(sizeof(typename RV::dtype) == 999, "Unsupported dtype"); + } + } + } + + return (__ballot_sync(0xffffffff, nan_detected) != 0); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/shared.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/shared.cuh new file mode 100644 index 0000000..404f8af --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/shared.cuh @@ -0,0 +1,7 @@ + + + + + +#include "tile/tile.cuh" +#include "vec/vec.cuh" diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/tile/conversions.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/tile/conversions.cuh new file mode 100644 index 0000000..45810c1 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/tile/conversions.cuh @@ -0,0 +1,16 @@ + + + + + + + +template +__device__ static inline void copy(ST1 &dst, const ST2 &src) { + static_assert(ST1::rows == ST2::rows && ST1::cols == ST2::cols, "Tiles must have the same rows and cols"); + #pragma unroll + for(int i = laneid(); i < ST1::num_elements; i+=GROUP_THREADS) { + int row = i/ST1::cols, col = i%ST1::cols; + dst[{row, col}] = base_types::convertor::convert(src[{row, col}]); + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/tile/maps.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/tile/maps.cuh new file mode 100644 index 0000000..c5c1b36 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/tile/maps.cuh @@ -0,0 +1,487 @@ + + + + + + + + + + + + + + + + +template +__device__ static inline void unary_map(T &dst, const T &src) { + #pragma unroll + for(int i = laneid(); i < T::num_elements; i += GROUP_THREADS) { + dst.data[i] = op::template op(src.data[i]); + } +} + + + + + + + + + + + + + +template +__device__ static inline void bin_map(T &dst, const T &src, const typename T::dtype ¶m) { + #pragma unroll + for(int i = laneid(); i < T::num_elements; i += GROUP_THREADS) { + dst.data[i] = op::template op(src.data[i], param); + } +} + + + + + + + + + + + + + +template +__device__ static inline void bin_map(T &dst, const T &lhs, const T &rhs) { + #pragma unroll + for(int i = laneid(); i < T::num_elements; i += GROUP_THREADS) { + dst.data[i] = op::template op(lhs.data[i], rhs.data[i]); + } +} + + + + + + + + + + + + + + + +template +__device__ static inline void row_map(T &dst, const T &src, const V &vec) { + static_assert(std::is_same::value, "Tile and vector must have the same data type"); + static_assert(V::length == T::rows, "Vector length must match the number of rows in the tile"); + #pragma unroll + for(int i = laneid(); i < T::num_elements; i += GROUP_THREADS) { + int row = i/T::cols, col = i%T::cols; + dst[{row, col}] = op::template op(src[{row, col}], vec[row]); + } +} + + + + + + + + + + + + + + + +template +__device__ static inline void col_map(T &dst, const T &src, const V &vec) { + static_assert(std::is_same::value, "Tile and vector must have the same data type"); + static_assert(V::length == T::cols, "Vector length must match the number of columns in the tile"); + #pragma unroll + for(int i = laneid(); i < T::num_elements; i += GROUP_THREADS) { + int row = i/T::cols, col = i%T::cols; + dst[{row, col}] = op::template op(src[{row, col}], vec[col]); + } +} + + + + + + + + + + + + + + +template +__device__ static inline void zero(T &dst) { + unary_map(dst, dst); +} + + + + + + +template +__device__ static inline void one(T &dst) { + unary_map(dst, dst); +} + + + + + + +template +__device__ static inline void pos_infty(T &dst) { + unary_map(dst, dst); +} + + + + + + +template +__device__ static inline void neg_infty(T &dst) { + unary_map(dst, dst); +} + + + + + + + + + +template +__device__ static inline void exp(T &dst, const T &src) { + unary_map(dst, src); +} + + + + + + + +template +__device__ static inline void exp2(T &dst, const T &src) { + unary_map(dst, src); +} + + + + + + + +template +__device__ static inline void log(T &dst, const T &src) { + unary_map(dst, src); +} + + + + + + + +template +__device__ static inline void log2(T &dst, const T &src) { + unary_map(dst, src); +} + + + + + + + +template +__device__ static inline void abs(T &dst, const T &src) { + unary_map(dst, src); +} + + + + + + + +template +__device__ static inline void relu(T &dst, const T &src) { + unary_map(dst, src); +} + + + + + + + + +template +__device__ static inline void copy(T &dst, const U &src) { + bin_map(dst, src); +} + + + + + + + + + + + +template +__device__ static inline void max(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void min(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void add(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void sub(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void mul(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void div(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + + + + + + + + + + + + +template +__device__ static inline void add_row(T &dst, const T &src, const V &row_values) { + row_map(dst, src, row_values); +} + + + + + + + + + +template +__device__ static inline void sub_row(T &dst, const T &src, const V &row_values) { + row_map(dst, src, row_values); +} + + + + + + + + + +template +__device__ static inline void mul_row(T &dst, const T &src, const V &row_values) { + row_map(dst, src, row_values); +} + + + + + + + + + +template +__device__ static inline void div_row(T &dst, const T &src, const V &row_values) { + row_map(dst, src, row_values); +} + + + + + + + + +template +__device__ static inline void broadcast_row(T &dst, const V &row_values) { + row_map(dst, dst, row_values); +} + + + + + + + + + + + + +template +__device__ static inline void add_col(T &dst, const T &src, const V &col_values) { + col_map(dst, src, col_values); +} + + + + + + + + + +template +__device__ static inline void sub_col(T &dst, const T &src, const V &col_values) { + col_map(dst, src, col_values); +} + + + + + + + + + +template +__device__ static inline void mul_col(T &dst, const T &src, const V &col_values) { + col_map(dst, src, col_values); +} + + + + + + + + + +template +__device__ static inline void div_col(T &dst, const T &src, const V &col_values) { + col_map(dst, src, col_values); +} + + + + + + + + +template +__device__ static inline void broadcast_col(T &dst, const V &col_values) { + col_map(dst, dst, col_values); +} + + + +template +__device__ static inline void add(T &dst, const T &src, const V &col_values) { + if constexpr (axis == axis::COL) add_col(dst, src, col_values); + else add_row(dst, src, col_values); +} + +template +__device__ static inline void sub(T &dst, const T &src, const V &col_values) { + if constexpr (axis == axis::COL) sub_col(dst, src, col_values); + else sub_row(dst, src, col_values); +} + +template +__device__ static inline void mul(T &dst, const T &src, const V &col_values) { + if constexpr (axis == axis::COL) mul_col(dst, src, col_values); + else mul_row(dst, src, col_values); +} + +template +__device__ static inline void div(T &dst, const T &src, const V &col_values) { + if constexpr (axis == axis::COL) div_col(dst, src, col_values); + else div_row(dst, src, col_values); +} + +template +__device__ static inline void broadcast(T &dst, const V &col_values) { + if constexpr (axis == axis::COL) broadcast_col(dst, col_values); + else broadcast_row(dst, col_values); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/tile/reductions.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/tile/reductions.cuh new file mode 100644 index 0000000..b4926e1 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/tile/reductions.cuh @@ -0,0 +1,372 @@ + + + + + + + + + + + + + + + + +template +__device__ static inline void row_reduce(V &row_accum, const T &src, const V &src_accum) { + using dtype = typename V::dtype; + for (int row = laneid(); row < T::rows; row += GROUP_THREADS) { + dtype accum = src[{row, 0}]; + #pragma unroll + for (int col = 1; col < T::cols; col++) { + accum = op::template op(accum, src[{row, col}]); + } + if (reset) { + row_accum[row] = accum; + } else { + row_accum[row] = op::template op(src_accum[row], accum); + } + } +} + + + + + + + + + + + + +template +__device__ static inline void col_reduce(V &col_accum, const T &src, const V &src_accum) { + using dtype = typename V::dtype; + for (int col = laneid(); col < T::cols; col += GROUP_THREADS) { + dtype accum = src[{0, col}]; + #pragma unroll + for (int row = 1; row < T::rows; row++) { + accum = op::template op(accum, src[{row, col}]); + } + if (reset) { + col_accum[col] = accum; + } else { + col_accum[col] = op::template op(src_accum[col], accum); + } + } +} + + + + + + + + + + + +template +__device__ static inline void row_max(V &row_accum, const T &src) { + row_reduce(row_accum, src, row_accum); +} + + + + + + + + +template +__device__ static inline void row_min(V &row_accum, const T &src) { + row_reduce(row_accum, src, row_accum); +} + + + + + + + + +template +__device__ static inline void row_sum(V &row_accum, const T &src) { + row_reduce(row_accum, src, row_accum); +} + + + + + + + + +template +__device__ static inline void row_prod(V &row_accum, const T &src) { + row_reduce(row_accum, src, row_accum); +} + + + + + + + + + + +template +__device__ static inline void row_max(V &row_accum, const T &src, const V &src_accum) { + row_reduce(row_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void row_min(V &row_accum, const T &src, const V &src_accum) { + row_reduce(row_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void row_sum(V &row_accum, const T &src, const V &src_accum) { + row_reduce(row_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void row_prod(V &row_accum, const T &src, const V &src_accum) { + row_reduce(row_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void col_max(V &col_accum, const T &src) { + col_reduce(col_accum, src, col_accum); +} + + + + + + + + +template +__device__ static inline void col_min(V &col_accum, const T &src) { + col_reduce(col_accum, src, col_accum); +} + + + + + + + + +template +__device__ static inline void col_sum(V &col_accum, const T &src) { + col_reduce(col_accum, src, col_accum); +} + + + + + + + + +template +__device__ static inline void col_prod(V &col_accum, const T &src) { + col_reduce(col_accum, src, col_accum); +} + + + + + + + + + + +template +__device__ static inline void col_max(V &col_accum, const T &src, const V &src_accum) { + col_reduce(col_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void col_min(V &col_accum, const T &src, const V &src_accum) { + col_reduce(col_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void col_sum(V &col_accum, const T &src, const V &src_accum) { + col_reduce(col_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void col_prod(V &col_accum, const T &src, const V &src_accum) { + col_reduce(col_accum, src, src_accum); +} + + + +template +__device__ static inline void max(V &dst, const T &src, const V &src_accum) { + if constexpr (ax == axis::COL) row_max(dst, src, src_accum); + else col_max(dst, src, src_accum); +} +template +__device__ static inline auto max(const T &src, const V &src_accum) { + V dst; + if constexpr (ax == axis::COL) row_max(dst, src, src_accum); + else col_max(dst, src, src_accum); + return dst; +} +template +__device__ static inline void max(V &dst, const T &src) { + if constexpr (ax == axis::COL) row_max(dst, src); + else col_max(dst, src); +} +template +__device__ static inline auto max(const T &src) { + using V = std::conditional_t; + V dst; + if constexpr (ax == axis::COL) row_max(dst, src); + else col_max(dst, src); + return dst; +} + +template +__device__ static inline void min(V &dst, const T &src, const V &src_accum) { + if constexpr (ax == axis::COL) row_min(dst, src, src_accum); + else col_min(dst, src, src_accum); +} +template +__device__ static inline auto min(const T &src, const V &src_accum) { + V dst; + if constexpr (ax == axis::COL) row_min(dst, src, src_accum); + else col_min(dst, src, src_accum); + return dst; +} +template +__device__ static inline void min(V &dst, const T &src) { + if constexpr (ax == axis::COL) row_min(dst, src); + else col_min(dst, src); +} +template +__device__ static inline auto min(const T &src) { + using V = std::conditional_t; + V dst; + if constexpr (ax == axis::COL) row_min(dst, src); + else col_min(dst, src); + return dst; +} + +template +__device__ static inline void sum(V &dst, const T &src, const V &src_accum) { + if constexpr (ax == axis::COL) row_sum(dst, src, src_accum); + else col_sum(dst, src, src_accum); +} +template +__device__ static inline auto sum(const T &src, const V &src_accum) { + V dst; + if constexpr (ax == axis::COL) row_sum(dst, src, src_accum); + else col_sum(dst, src, src_accum); + return dst; +} +template +__device__ static inline void sum(V &dst, const T &src) { + if constexpr (ax == axis::COL) row_sum(dst, src); + else col_sum(dst, src); +} +template +__device__ static inline auto sum(const T &src) { + using V = std::conditional_t; + V dst; + if constexpr (ax == axis::COL) row_sum(dst, src); + else col_sum(dst, src); + return dst; +} + +template +__device__ static inline void prod(V &dst, const T &src, const V &src_accum) { + if constexpr (ax == axis::COL) row_prod(dst, src, src_accum); + else col_prod(dst, src, src_accum); +} +template +__device__ static inline auto prod(const T &src, const V &src_accum) { + V dst; + if constexpr (ax == axis::COL) row_prod(dst, src, src_accum); + else col_prod(dst, src, src_accum); + return dst; +} +template +__device__ static inline void prod(V &dst, const T &src) { + if constexpr (ax == axis::COL) row_prod(dst, src); + else col_prod(dst, src); +} +template +__device__ static inline auto prod(const T &src) { + using V = std::conditional_t; + V dst; + if constexpr (ax == axis::COL) row_prod(dst, src); + else col_prod(dst, src); + return dst; +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/tile/tile.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/tile/tile.cuh new file mode 100644 index 0000000..cf76bda --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/tile/tile.cuh @@ -0,0 +1,37 @@ + + + + + +#include "conversions.cuh" +#include "maps.cuh" +#include "reductions.cuh" + +template +__device__ static inline bool hasnan(const ST &src) { + KITTENS_CHECK_WARP + bool nan_detected = false; + #pragma unroll + for(int i = laneid(); i < ST::num_elements; i+=GROUP_THREADS) { + if constexpr (std::is_same_v) { + if(isnan(src[i])) { + nan_detected = true; + } + } + else if constexpr (std::is_same_v) { + if(isnan(__bfloat162float(src[i]))) { + nan_detected = true; + } + } + else if constexpr (std::is_same_v) { + if(isnan(__half2float(src[i]))) { + nan_detected = true; + } + } + else { + static_assert(sizeof(typename ST::T) == 999, "Unsupported dtype"); + } + } + + return (__ballot_sync(0xffffffff, nan_detected) != 0); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/vec/conversions.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/vec/conversions.cuh new file mode 100644 index 0000000..5a0fa4a --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/vec/conversions.cuh @@ -0,0 +1,27 @@ + + + + + + + + + + + + + + + + + + + +template +__device__ static inline void copy(SV1 &dst, const SV2 &src) { + static_assert(SV1::length == SV2::length, "Source and destination vectors must have the same length."); + #pragma unroll + for(int i = laneid(); i < SV1::length; i+=GROUP_THREADS) { + dst[i] = base_types::convertor::convert(src[i]); + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/vec/maps.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/vec/maps.cuh new file mode 100644 index 0000000..e491c3e --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/vec/maps.cuh @@ -0,0 +1,259 @@ + + + + + + + + + + + + + +template +__device__ static inline void unary_op(T &dst, const T &src) { + #pragma unroll + for(auto cur = laneid(); cur < T::length; cur+=GROUP_THREADS) { + dst[cur] = op::template op(src[cur]); + } +} + + + + + + + + + +template +__device__ static inline void bin_op(T &dst, const T &lhs, const T &rhs) { + #pragma unroll + for(auto cur = laneid(); cur < T::length; cur+=GROUP_THREADS) { + dst[cur] = op::template op(lhs[cur], rhs[cur]); + } +} + + + + + + + + + +template +__device__ static inline void bin_op(T &dst, const T &src, const typename T::dtype ¶m) { + #pragma unroll + for(auto cur = laneid(); cur < T::length; cur+=GROUP_THREADS) { + dst[cur] = op::template op(src[cur], param); + } +} + + + + + + + + + + + +template +__device__ static inline void zero(T &dst) { + unary_op(dst, dst); +} + + + + + + +template +__device__ static inline void one(T &dst) { + unary_op(dst, dst); +} + + + + + + +template +__device__ static inline void pos_infty(T &dst) { + unary_op(dst, dst); +} + + + + + + +template +__device__ static inline void neg_infty(T &dst) { + unary_op(dst, dst); +} + + + + + + + + + + + +template +__device__ static inline void copy(T &dst, const U &src) { + bin_op(dst, dst, src); +} + + + + + + + +template +__device__ static inline void exp(T &dst, const T &src) { + unary_op(dst, src); +} + + + + + + + +template +__device__ static inline void exp2(T &dst, const T &src) { + unary_op(dst, src); +} + + + + + + + +template +__device__ static inline void log(T &dst, const T &src) { + unary_op(dst, src); +} + + + + + + + +template +__device__ static inline void log2(T &dst, const T &src) { + unary_op(dst, src); +} + + + + + + + +template +__device__ static inline void abs(T &dst, const T &src) { + unary_op(dst, src); +} + + + + + + + +template +__device__ static inline void relu(T &dst, const T &src) { + unary_op(dst, src); +} + + + + + + + + + + + + +template +__device__ static inline void max(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void min(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void add(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void sub(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void mul(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void div(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/vec/reductions.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/vec/reductions.cuh new file mode 100644 index 0000000..89212da --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/vec/reductions.cuh @@ -0,0 +1,207 @@ + + + + + + + + + + + + + + + + + + + + + + +template +__device__ static inline void reduce(typename SV::dtype &dst_accum, const SV &src, const typename SV::dtype &src_accum) { + if constexpr (GROUP_WARPS == 1) { + using T = SV::dtype; + int lane = laneid(); + T accum; + if(lane < SV::length) accum = src[lane]; + __syncwarp(); + for(int i = lane+kittens::WARP_THREADS; i < SV::length; i+=kittens::WARP_THREADS) { + accum = op::template op(accum, src[i]); + } + __syncwarp(); + + if constexpr (SV::length > 16) { + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 16)); + __syncwarp(); + } + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 8)); + __syncwarp(); + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 4)); + __syncwarp(); + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 2)); + __syncwarp(); + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 1)); + __syncwarp(); + if constexpr (!reset) accum = op::template op(accum, src_accum); + + dst_accum = packed_shfl_sync(kittens::MASK_ALL, accum, 0); + } + else { + ::kittens::group<1>::reduce(dst_accum, src, src_accum); + } +} + + + + + + + + + + +template +__device__ static inline void max(typename SV::dtype &max_val, const SV &src) { + reduce(max_val, src, max_val); +} +template +__device__ static inline typename SV::dtype max(const SV &src) { + typename SV::dtype max_val; + reduce(max_val, src, max_val); + return max_val; +} + + + + + + + + +template +__device__ static inline void min(typename SV::dtype &min_val, const SV &src) { + reduce(min_val, src, min_val); +} +template +__device__ static inline typename SV::dtype min(const SV &src) { + typename SV::dtype min_val; + reduce(min_val, src, min_val); + return min_val; +} + + + + + + + + +template +__device__ static inline void sum(typename SV::dtype &sum_val, const SV &src) { + reduce(sum_val, src, sum_val); +} +template +__device__ static inline typename SV::dtype sum(const SV &src) { + typename SV::dtype sum_val; + reduce(sum_val, src, sum_val); + return sum_val; +} + + + + + + + + +template +__device__ static inline void prod(typename SV::dtype &prod_val, const SV &src) { + reduce(prod_val, src, prod_val); +} +template +__device__ static inline typename SV::dtype prod(const SV &src) { + typename SV::dtype prod_val; + reduce(prod_val, src, prod_val); + return prod_val; +} + + + + + + + + + + + +template +__device__ static inline void max(typename SV::dtype &max_val, const SV &src, const typename SV::dtype &src_accum) { + reduce(max_val, src, src_accum); +} +template +__device__ static inline typename SV::dtype max(const SV &src, const typename SV::dtype &src_accum) { + typename SV::dtype max_val; + reduce(max_val, src, src_accum); + return max_val; +} + + + + + + + + + +template +__device__ static inline void min(typename SV::dtype &min_val, const SV &src, const typename SV::dtype &src_accum) { + reduce(min_val, src, src_accum); +} +template +__device__ static inline typename SV::dtype min(const SV &src, const typename SV::dtype &src_accum) { + typename SV::dtype min_val; + reduce(min_val, src, src_accum); + return min_val; +} + + + + + + + + + +template +__device__ static inline void sum(typename SV::dtype &sum_val, const SV &src, const typename SV::dtype &src_accum) { + reduce(sum_val, src, src_accum); +} +template +__device__ static inline typename SV::dtype sum(const SV &src, const typename SV::dtype &src_accum) { + typename SV::dtype sum_val; + reduce(sum_val, src, src_accum); + return sum_val; +} + + + + + + + + + +template +__device__ static inline void prod(typename SV::dtype &prod_val, const SV &src, const typename SV::dtype &src_accum) { + reduce(prod_val, src, src_accum); +} +template +__device__ static inline typename SV::dtype prod(const SV &src, const typename SV::dtype &src_accum) { + typename SV::dtype prod_val; + reduce(prod_val, src, src_accum); + return prod_val; +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/vec/vec.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/vec/vec.cuh new file mode 100644 index 0000000..5354d34 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/vec/vec.cuh @@ -0,0 +1,37 @@ + + + + + +#include "conversions.cuh" +#include "maps.cuh" +#include "reductions.cuh" + +template +__device__ static inline bool hasnan(const SV &src) { + KITTENS_CHECK_WARP + bool nan_detected = false; + #pragma unroll + for(int i = laneid(); i < SV::length; i+=GROUP_THREADS) { + if constexpr (std::is_same_v) { + if(isnan(src[i])) { + nan_detected = true; + } + } + else if constexpr (std::is_same_v) { + if(isnan(__bfloat162float(src[i]))) { + nan_detected = true; + } + } + else if constexpr (std::is_same_v) { + if(isnan(__half2float(src[i]))) { + nan_detected = true; + } + } + else { + static_assert(sizeof(typename SV::T) == 999, "Unsupported dtype"); + } + } + + return (__ballot_sync(0xffffffff, nan_detected) != 0); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/util/sync.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/util/sync.cuh new file mode 100644 index 0000000..e500aba --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/util/sync.cuh @@ -0,0 +1,224 @@ + + + + + +template __device__ static inline void load_async_wait(int bar_id) { + asm volatile("cp.async.wait_group %0;\n" : : "n"(N) : "memory"); + sync(bar_id); +} +template __device__ static inline void load_async_wait() { + KITTENS_CHECK_WARP + asm volatile("cp.async.wait_group %0;\n" : : "n"(N) : "memory"); + __syncwarp(); +} + +__device__ static inline void arrive(barrier bar) { + asm volatile("bar.arrive %0, %1;\n" :: "r"(bar.barrier_id), "n"(GROUP_WARPS*WARP_THREADS) : "memory"); +} +__device__ static inline void arrive_and_wait(barrier bar) { + asm volatile("bar.sync %0, %1;\n" :: "r"(bar.barrier_id), "n"(GROUP_WARPS*WARP_THREADS) : "memory"); +} + + + + + + + + + + + + + +__device__ static inline void init_semaphore(semaphore& bar, int thread_count, int transaction_count=0) { + if (laneid() == 0) { + void const* const ptr = &bar; + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + + asm volatile ( + "mbarrier.init.shared::cta.b64 [%0], %1;\n" + :: "r"(bar_ptr), "r"(thread_count+transaction_count) + ); + } +} + + + + + +__device__ static inline void invalidate_semaphore(semaphore& bar) { + if (laneid() == 0) { + void const* const ptr = &bar; + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + asm volatile ( + "mbarrier.inval.shared::cta.b64 [%0];\n" + :: "r"(bar_ptr) + ); + } +} + + + + + + + + +__device__ static inline void arrive(semaphore& sem) { + if(laneid() == 0) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(&sem)); + asm volatile ( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];\n" + : + : "r"(mbar_ptr) + : "memory" + ); + } +} +template __device__ static inline void arrive(barrier bar) { + asm volatile("bar.arrive %0, %1;\n" :: "r"(bar.barrier_id), "n"(num_warps*WARP_THREADS) : "memory"); +} + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + +__device__ static inline void arrive(semaphore& sem, uint32_t count) { + if(laneid() == 0) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(&sem)); + asm volatile ( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0], %1;\n" + : + : "r"(mbar_ptr), "r"(count) + : "memory" + ); + } +} +#endif + + + + + + + +__device__ static inline void wait(semaphore& sem, int kPhaseBit) { + void const* const ptr = &sem; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +#else + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.test_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "nanosleep.u32 5;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +#endif +} + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +__device__ static inline bool try_wait(semaphore &sem, int kPhaseBit) { + void const* const ptr = &sem; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + uint32_t success; + + asm volatile( + "{\n" + ".reg .pred P1; \n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2; \n" + "selp.b32 %0, 1, 0, P1; \n" + "}\n" + : "=r"(success) + : "r"(mbar_ptr), "r"(kPhaseBit) + : "memory" + ); + + return static_cast(success); +} +#endif + + + + + + + +__device__ static inline int test_wait(semaphore& sem, int kPhaseBit) { + void const* const ptr = &sem; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + int result; + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "mbarrier.test_wait.parity.shared::cta.b64 P1, [%1], %2;\n" + "selp.u32 %0,1,0,P1;" + "}\n" + : "=r"(result) + : "r"(mbar_ptr), "r"(kPhaseBit) + ); + return result; +} + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + +__device__ static inline bool elect_leader() { + if constexpr (GROUP_WARPS == 1) { + uint32_t elected = 0; + asm volatile( + "{.reg .pred P;\n" + " elect.sync _|P, %1;\n" + " selp.u32 %0, 1, 0, P;}\n" + : "+r"(elected) + : "r"(0xFFFFFFFF) + ); + return static_cast(elected); + } else { + return (warpid() == 0 && ::kittens::group<1>::elect_leader()); + } +} + +__device__ static inline bool elect_leader_from_active() { + if constexpr (GROUP_WARPS == 1) { + uint32_t elected = 0; + asm volatile( + "{.reg .pred P;\n" + " .reg .u32 mask;\n" + " activemask.b32 mask;\n" + " elect.sync _|P, mask;\n" + " selp.u32 %0, 1, 0, P;}\n" + : "+r"(elected) + ); + return static_cast(elected); + } else { + return (warpid() == 0 && ::kittens::group<1>::elect_leader_from_active()); + } +} + +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/util/tma.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/util/tma.cuh new file mode 100644 index 0000000..015b31a --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/util/tma.cuh @@ -0,0 +1,90 @@ + + + + + + + + + + + + + + + + + +__device__ static inline void expect_bytes(semaphore& bar, uint32_t bytes) { + if(laneid() == 0) { + ::kittens::tma::expect_bytes(bar, bytes); + } +} + + + + + +template +__device__ static inline void expect(semaphore& bar, const T& _1, const args&... _2) { + expect_bytes(bar, size_bytes); +} + + + + + + +__device__ static inline void store_commit_group() { + asm volatile("cp.async.bulk.commit_group;"); +} + + + + + +template +__device__ static inline void store_async_wait() { + asm volatile ( + "cp.async.bulk.wait_group %0;" + : + : "n"(N) + : "memory" + ); +} + + + + + +template +__device__ static inline void store_async_read_wait() { + asm volatile ( + "cp.async.bulk.wait_group.read %0;" + : + : "n"(N) + : "memory" + ); +} + + + +__device__ static inline void load_async(void *dst, void *src, uint32_t size_bytes, semaphore& bar) { + if(laneid() == 0) { + ::kittens::tma::load_async(dst, src, size_bytes, bar); + } +} +template +__device__ static inline void load_async(T &dst, T &src, uint32_t size_bytes, semaphore& bar) { + load_async(reinterpret_cast(&dst), reinterpret_cast(&src), size_bytes, bar); +} + +__device__ static inline void store_async(void *dst, void *src, uint32_t size_bytes) { + if(laneid() == 0) { + ::kittens::tma::store_async(dst, src, size_bytes); + } +} +template +__device__ static inline void store_async(T &dst, T &src, uint32_t size_bytes) { + store_async(reinterpret_cast(&dst), reinterpret_cast(&src), size_bytes); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/util/tma_cluster.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/util/tma_cluster.cuh new file mode 100644 index 0000000..7404629 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/util/tma_cluster.cuh @@ -0,0 +1,118 @@ + + + + + + + + + + + +__device__ static inline void wait(semaphore& bar, int kPhaseBit) { + void const* const ptr = &bar; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +} + +__device__ static inline bool try_wait(semaphore &bar, int kPhaseBit) { + void const* const ptr = &bar; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + uint32_t success; + + asm volatile( + "{\n" + ".reg .pred P1; \n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [%1], %2; \n" + "selp.b32 %0, 1, 0, P1; \n" + "}\n" + : "=r"(success) + : "r"(mbar_ptr), "r"(kPhaseBit) + : "memory" + ); + + return static_cast(success); +} + + + + + + + + + + + + + + + +__device__ static inline void expect_bytes(semaphore& bar, uint32_t bytes) { + if(laneid() == 0) { + ::kittens::tma::cluster::expect_bytes(bar, bytes); + } +} +__device__ static inline void expect_bytes(semaphore& bar, uint32_t bytes, int dst_cta) { + if(laneid() == 0) { + ::kittens::tma::cluster::expect_bytes(bar, bytes, dst_cta); + } +} + + + + + + + + + + + + + + + +template +__device__ static inline void expect(semaphore& bar, const T& _1, const args&... _2) { + expect_bytes(bar, size_bytes); +} + + + + + + + + + + +__device__ static inline void arrive(semaphore& bar, int dst_cta, uint32_t count=1) { + if(laneid() == 0) { + ::kittens::tma::cluster::arrive(bar, dst_cta, count); + } +} + + + +__device__ static inline void load_async(void *dst, void *src, uint32_t size_bytes, semaphore& bar, uint16_t cta_mask) { + if(laneid() == 0) { + ::kittens::tma::cluster::load_async(dst, src, size_bytes, bar, cta_mask); + } +} +template +__device__ static inline void load_async(T &dst, T &src, uint32_t size_bytes, semaphore& bar, uint16_t cta_mask) { + load_async(reinterpret_cast(&dst), reinterpret_cast(&src), size_bytes, bar, cta_mask); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/util/util.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/util/util.cuh new file mode 100644 index 0000000..295407b --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/util/util.cuh @@ -0,0 +1,94 @@ + + + + + +#include "sync.cuh" + + + + +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + +struct clc { + + + + + + + + +__device__ static inline void schedule(kittens::clc::handle &h, kittens::semaphore &sem) { + if (laneid() == 0) { + asm volatile("{clusterlaunchcontrol.try_cancel.async.shared::cta.mbarrier::complete_tx::bytes.multicast::cluster::all.b128 [%0], [%1];}" + :: "r"(static_cast(__cvta_generic_to_shared(&h.internal_value))), "r"(static_cast(__cvta_generic_to_shared(&sem))) + : "memory" + ); + } +} + + + + + +__device__ static inline kittens::clc::result query(kittens::clc::handle &h) { + kittens::clc::result r; + asm volatile( + "{\n" + ".reg .pred SUCCESS;\n" + ".reg .b128 CLC_HANDLE;\n" + "ld.shared.b128 CLC_HANDLE, [%4];\n" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 SUCCESS, CLC_HANDLE;\n" + "selp.u32 %0, 1, 0, SUCCESS;\n" + "@!SUCCESS bra.uni DONE;\n" + "clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%1, %2, %3, _}, CLC_HANDLE;\n" + "fence.proxy.async.shared::cta;\n" + "DONE:\n" + "}" + : "=r"(r.success), "=r"(r.x), "=r"(r.y), "=r"(r.z) + : "r"(static_cast(__cvta_generic_to_shared(&h.internal_value))) + : "memory" + ); + return r; +} + +}; + +#endif + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + + +struct pdl { + + + + + + + + +__device__ static inline void arrive() { + if (laneid() == 0) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +} + + + + +__device__ static inline void wait() { + asm volatile("griddepcontrol.wait;" ::: "memory"); +} + +}; + +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/ops.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/ops.cuh new file mode 100644 index 0000000..0c58623 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/ops.cuh @@ -0,0 +1,261 @@ + + + + + +#pragma once + +#include "thread/thread.cuh" +#include "group/group.cuh" + +namespace kittens { + + + + + +template +__device__ static inline T operator+(const T &lhs, const U &rhs) { + T dst; + warp::add(dst, lhs, rhs); + return dst; +} +template +__device__ static inline void operator+=(T &lhs, const U &rhs) { + warp::add(lhs, lhs, rhs); +} +template +__device__ static inline T operator-(const T &lhs, const U &rhs) { + T dst; + warp::sub(dst, lhs, rhs); + return dst; +} +template +__device__ static inline void operator-=(T &lhs, const U &rhs) { + warp::sub(lhs, lhs, rhs); +} +template +__device__ static inline T operator*(const T &lhs, const U &rhs) { + T dst; + warp::mul(dst, lhs, rhs); + return dst; +} +template +__device__ static inline void operator*=(T &lhs, const U &rhs) { + warp::mul(lhs, lhs, rhs); +} +template +__device__ static inline T operator/(const T &lhs, const U &rhs) { + T dst; + warp::div(dst, lhs, rhs); + return dst; +} +template +__device__ static inline void operator/=(T &lhs, const U &rhs) { + warp::div(lhs, lhs, rhs); +} +template +__device__ static inline T operator+(const T &src, const V &row_values) { + T dst; + warp::add_row(dst, src, row_values); + return dst; +} +template +__device__ static inline T operator+(const T &src, const V &row_values) { + T dst; + warp::add_row(dst, src, row_values); + return dst; +} +template +__device__ static inline void operator+=(T &lhs, const V &row_values) { + warp::add_row(lhs, lhs, row_values); +} +template +__device__ static inline void operator+=(T &lhs, const V &row_values) { + warp::add_row(lhs, lhs, row_values); +} +template +__device__ static inline T operator-(const T &src, const V &row_values) { + T dst; + warp::sub_row(dst, src, row_values); + return dst; +} +template +__device__ static inline T operator-(const T &src, const V &row_values) { + T dst; + warp::sub_row(dst, src, row_values); + return dst; +} +template +__device__ static inline void operator-=(T &lhs, const V &row_values) { + warp::sub_row(lhs, lhs, row_values); +} +template +__device__ static inline void operator-=(T &lhs, const V &row_values) { + warp::sub_row(lhs, lhs, row_values); +} +template +__device__ static inline T operator*(const T &src, const V &row_values) { + T dst; + warp::mul_row(dst, src, row_values); + return dst; +} +template +__device__ static inline T operator*(const T &src, const V &row_values) { + T dst; + warp::mul_row(dst, src, row_values); + return dst; +} +template +__device__ static inline void operator*=(T &lhs, const V &row_values) { + warp::mul_row(lhs, lhs, row_values); +} +template +__device__ static inline void operator*=(T &lhs, const V &row_values) { + warp::mul_row(lhs, lhs, row_values); +} +template +__device__ static inline T operator/(const T &src, const V &row_values) { + T dst; + warp::div_row(dst, src, row_values); + return dst; +} +template +__device__ static inline T operator/(const T &src, const V &row_values) { + T dst; + warp::div_row(dst, src, row_values); + return dst; +} +template +__device__ static inline void operator/=(T &lhs, const V &row_values) { + warp::div_row(lhs, lhs, row_values); +} +template +__device__ static inline void operator/=(T &lhs, const V &row_values) { + warp::div_row(lhs, lhs, row_values); +} +template +__device__ static inline T operator+(const T &src, const V &col_values) { + T dst; + warp::add_col(dst, src, col_values); + return dst; +} +template +__device__ static inline T operator+(const T &src, const V &col_values) { + T dst; + warp::add_col(dst, src, col_values); + return dst; +} +template +__device__ static inline void operator+=(T &lhs, const V &col_values) { + warp::add_col(lhs, lhs, col_values); +} +template +__device__ static inline void operator+=(T &lhs, const V &col_values) { + warp::add_col(lhs, lhs, col_values); +} +template +__device__ static inline T operator-(const T &src, const V &col_values) { + T dst; + warp::sub_col(dst, src, col_values); + return dst; +} +template +__device__ static inline T operator-(const T &src, const V &col_values) { + T dst; + warp::sub_col(dst, src, col_values); + return dst; +} +template +__device__ static inline void operator-=(T &lhs, const V &col_values) { + warp::sub_col(lhs, lhs, col_values); +} +template +__device__ static inline void operator-=(T &lhs, const V &col_values) { + warp::sub_col(lhs, lhs, col_values); +} +template +__device__ static inline T operator*(const T &src, const V &col_values) { + T dst; + warp::mul_col(dst, src, col_values); + return dst; +} +template +__device__ static inline T operator*(const T &src, const V &col_values) { + T dst; + warp::mul_col(dst, src, col_values); + return dst; +} +template +__device__ static inline void operator*=(T &lhs, const V &col_values) { + warp::mul_col(lhs, lhs, col_values); +} +template +__device__ static inline void operator*=(T &lhs, const V &col_values) { + warp::mul_col(lhs, lhs, col_values); +} +template +__device__ static inline T operator/(const T &src, const V &col_values) { + T dst; + warp::div_col(dst, src, col_values); + return dst; +} +template +__device__ static inline T operator/(const T &src, const V &col_values) { + T dst; + warp::div_col(dst, src, col_values); + return dst; +} +template +__device__ static inline void operator/=(T &lhs, const V &col_values) { + warp::div_col(lhs, lhs, col_values); +} +template +__device__ static inline void operator/=(T &lhs, const V &col_values) { + warp::div_col(lhs, lhs, col_values); +} + + + +template +__device__ static inline T operator+(const T &lhs, const U &rhs) { + T dst; + warp::add(dst, lhs, rhs); + return dst; +} +template +__device__ static inline void operator+=(T &lhs, const U &rhs) { + warp::add(lhs, lhs, rhs); +} +template +__device__ static inline T operator-(const T &lhs, const U &rhs) { + T dst; + warp::sub(dst, lhs, rhs); + return dst; +} +template +__device__ static inline void operator-=(T &lhs, const U &rhs) { + warp::sub(lhs, lhs, rhs); +} +template +__device__ static inline T operator*(const T &lhs, const U &rhs) { + T dst; + warp::mul(dst, lhs, rhs); + return dst; +} +template +__device__ static inline void operator*=(T &lhs, const U &rhs) { + warp::mul(lhs, lhs, rhs); +} +template +__device__ static inline T operator/(const T &lhs, const U &rhs) { + T dst; + warp::div(dst, lhs, rhs); + return dst; +} +template +__device__ static inline void operator/=(T &lhs, const U &rhs) { + warp::div(lhs, lhs, rhs); +} + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/memory.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/memory.cuh new file mode 100644 index 0000000..3b0cb5f --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/memory.cuh @@ -0,0 +1,9 @@ + + + + + +#pragma once + +#include "tile/tile.cuh" +#include "vec/vec.cuh" diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/tile/shared_to_tensor.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/tile/shared_to_tensor.cuh new file mode 100644 index 0000000..72eed49 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/tile/shared_to_tensor.cuh @@ -0,0 +1,39 @@ + + + + + +#pragma once + +#include "../../mma/tcgen05.cuh" + +template +__device__ inline static void load_mxnv_scale_async(TT &dst, const ST &src) { + static_assert(ncta == 1 || ncta == 2, "ncta must be 1 or 2"); + static_assert(std::is_same_v || std::is_same_v, "Scale TT must be fp8e8m0 or fp8e4m3"); + static_assert(std::is_same_v, "Scale TT and ST must have the same type"); + static_assert(TT::rows == 128 && TT::cols == 16, "Tensor memory scale tile must always be 128x16"); + static_assert(ST::rows == 32 && ST::cols == 16, "Shared memory scale tile must always be 32x16"); + static_assert(!ST::swizzle, "Shared memory scale tile must not be TMA-swizzled"); + + uint64_t st_desc = kittens::detail::matrix_descriptor_raw(&src.data[0], 128, 128, 0); + + if constexpr (ncta == 1) { + asm volatile("{tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;}" :: "r"(dst.addr), "l"(st_desc)); + } else { + asm volatile("{tcgen05.cp.cta_group::2.32x128b.warpx4 [%0], %1;}" :: "r"(dst.addr), "l"(st_desc)); + } +} +template +__device__ inline static void load_mxnv_scale_async(TT &dst, const ST &src, kittens::semaphore &sem) { + load_mxnv_scale_async(dst, src); + kittens::detail::tcgen05::commit(sem); +} +template +__device__ inline static void load_mxnv_scale_async2(TT &dst, const ST &src) { + load_mxnv_scale_async<2>(dst, src); +} +template +__device__ inline static void load_mxnv_scale_async2(TT &dst, const ST &src, kittens::semaphore &sem) { + load_mxnv_scale_async<2>(dst, src, sem); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/tile/tile.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/tile/tile.cuh new file mode 100644 index 0000000..9b41b67 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/tile/tile.cuh @@ -0,0 +1,13 @@ + + + + + +#pragma once + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +#include "tma.cuh" +#endif +#if defined(KITTENS_SM10X) +#include "shared_to_tensor.cuh" +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/tile/tma.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/tile/tma.cuh new file mode 100644 index 0000000..0a063ad --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/tile/tma.cuh @@ -0,0 +1,880 @@ + + + + + +#pragma once + +#include "../../../../common/common.cuh" +#include "../../../../types/types.cuh" +#include "../../util/util.cuh" + +namespace kittens { +namespace tma { + +namespace detail { +template __device__ inline int4 tma_coords(const coord &unit_coord) { + static_assert(ST::swizzle, "tma_coords only should be called for swizzled tiles"); + constexpr int swizzle_elements = ST::swizzle_bytes / sizeof(typename ST::dtype); + if constexpr (axis == 2) return {unit_coord.r, unit_coord.c / swizzle_elements, unit_coord.d, unit_coord.b}; + else if constexpr (axis == 1) return {unit_coord.d, unit_coord.c / swizzle_elements, unit_coord.r, unit_coord.b}; + else if constexpr (axis == 0) return {unit_coord.b, unit_coord.c / swizzle_elements, unit_coord.r, unit_coord.d}; +} +} + + + + + + + + + + + +template> +__device__ static inline void prefetch(ST &dst, const GL &src, const COORD &idx) { + uint64_t tma_ptr = reinterpret_cast(src.template get_tma()); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.prefetch.tensor.5d.L2.global.tile" + " [%0, {%1, %2, %3, %4, %5}];" + : + : "l"(tma_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.prefetch.tensor.5d.L2.global.tile.L2::cache_hint" + " [%0, {%1, %2, %3, %4, %5}], %6;" + : + : "l"(tma_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.prefetch.tensor.4d.L2.global.tile" + " [%0, {%1, %2, %3, %4}];" + : + : "l"(tma_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.prefetch.tensor.4d.L2.global.tile.L2::cache_hint" + " [%0, {%1, %2, %3, %4}], %5;" + : + : "l"(tma_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } + } +} +template> +__device__ static inline void prefetch(ST &dst, const GL &src, const COORD &idx) { + prefetch(dst, src, idx); +} + + + + + + + + + + + + + +template> +__device__ static inline void store_async(const GL &dst, const ST &src, const COORD &idx) { + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.5d.global.shared::cta.tile.bulk_group" + " [%0, {%2, %3, %4, %5, %6}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.tensor.5d.global.shared::cta.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5, %6}], [%1], %7;" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } + } + store_commit_group(); +} +template> +__device__ static inline void store_async(const GL &dst, const ST &src, const COORD &idx) { + store_async(dst, src, idx); +} +template> +__device__ static inline void store_async(const PGL &dst, const ST &src, const COORD &idx) { + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.5d.global.shared::cta.tile.bulk_group" + " [%0, {%2, %3, %4, %5, %6}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.tensor.5d.global.shared::cta.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5, %6}], [%1], %7;" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } + } + store_commit_group(); +} +template> +__device__ static inline void store_async(const PGL &dst, const ST &src, const COORD &idx) { + store_async(dst, src, idx); +} + + + + + + + + + + + + + +template> +__device__ static inline void store_add_async(const GL &dst, const ST &src, const COORD &idx) { + + static_assert(!(std::is_same_v || + std::is_same_v), + "TMA does not support async add reductions for fp8 types."); + + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.add.tile.bulk_group" + " [%0, {%2, %3, %4, %5, %6}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.add.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5, %6}], [%1], %7;" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.add.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.add.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } + } + store_commit_group(); +} +template> +__device__ static inline void store_add_async(const GL &dst, const ST &src, const COORD &idx) { + store_add_async(dst, src, idx); +} +template> +__device__ static inline void store_add_async(const PGL &dst, const ST &src, const COORD &idx) { + + static_assert(!(std::is_same_v || + std::is_same_v), + "TMA does not support async add reductions for fp8 types."); + + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.add.tile.bulk_group" + " [%0, {%2, %3, %4, %5, %6}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.add.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5, %6}], [%1], %7;" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.add.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.add.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } + } + store_commit_group(); +} +template> +__device__ static inline void store_add_async(const PGL &dst, const ST &src, const COORD &idx) { + store_add_async(dst, src, idx); +} + + + + + + + + + + + +template> +__device__ static inline void store_min_async(const GL &dst, const ST &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + + static_assert(!(std::is_same_v || + std::is_same_v), + "TMA does not support async min reductions for fp8 types."); + + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.min.tile.bulk_group" + " [%0, {%2, %3, %4, %5, %6}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.min.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5, %6}], [%1], %7;" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.min.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.min.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } + } + store_commit_group(); +} +template> +__device__ static inline void store_min_async(const GL &dst, const ST &src, const COORD &idx) { + store_min_async(dst, src, idx); +} +template> +__device__ static inline void store_min_async(const PGL &dst, const ST &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + + static_assert(!(std::is_same_v || + std::is_same_v), + "TMA does not support async min reductions for fp8 types."); + + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.min.tile.bulk_group" + " [%0, {%2, %3, %4, %5, %6}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.min.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5, %6}], [%1], %7;" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.min.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.min.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } + } + store_commit_group(); +} +template> +__device__ static inline void store_min_async(const PGL &dst, const ST &src, const COORD &idx) { + store_min_async(dst, src, idx); +} + + + + + + + + + + + +template> +__device__ static inline void store_max_async(const GL &dst, const ST &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + + static_assert(!(std::is_same_v || + std::is_same_v), + "TMA does not support async max reductions for fp8 types."); + + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.max.tile.bulk_group" + " [%0, {%2, %3, %4, %5, %6}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.max.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5, %6}], [%1], %7;" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.max.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.max.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } + } + store_commit_group(); +} +template> +__device__ static inline void store_max_async(const GL &dst, const ST &src, const COORD &idx) { + store_max_async(dst, src, idx); +} +template> +__device__ static inline void store_max_async(const PGL &dst, const ST &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + + static_assert(!(std::is_same_v || + std::is_same_v), + "TMA does not support async max reductions for fp8 types."); + + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.max.tile.bulk_group" + " [%0, {%2, %3, %4, %5, %6}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.max.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5, %6}], [%1], %7;" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.max.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.max.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } + } + store_commit_group(); +} +template> +__device__ static inline void store_max_async(const PGL &dst, const ST &src, const COORD &idx) { + store_max_async(dst, src, idx); +} + + + + + + + + + + + + +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar) { + uint64_t tma_ptr = reinterpret_cast(src.template get_tma()); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(&bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst)); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + + if constexpr (policy == cache_policy::NORMAL) { + asm volatile( + "cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5, %6, %7}], [%2];" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w) + : "memory" + ); + } + else { + asm volatile( + "cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + + if constexpr (policy == cache_policy::NORMAL) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5, %6}], [%2];" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b) + : "memory" + ); + } + else { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%3, %4, %5, %6}], [%2], %7;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } + } +} +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar) { + load_async(dst, src, idx, bar); +} + +namespace cluster { + + + + + + + + + + + + + +#ifdef KITTENS_SM10X +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) +#else +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask) +#endif +{ + uint64_t tma_ptr = reinterpret_cast(src.template get_tma()); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(&bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst)); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + +#ifdef KITTENS_SM10X + if(dst_mbar_cta != -1) { + uint32_t neighbor_mbar_ptr; + asm volatile ( + "mapa.shared::cluster.u32 %0, %1, %2;\n" + : "=r"(neighbor_mbar_ptr) + : "r"(mbar_ptr), "r"(dst_mbar_cta) + ); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.cta_group::2.multicast::cluster" + " [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(neighbor_mbar_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "h"(cluster_mask) + : "memory" + ); + } else { + asm volatile ( + "cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.cta_group::2.multicast::cluster.L2::cache_hint" + " [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8, %9;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(neighbor_mbar_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "h"(cluster_mask), "l"(make_cache_policy()) + : "memory" + ); + } + } else +#endif + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "h"(cluster_mask) + : "memory" + ); + } else { + asm volatile ( + "cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster.L2::cache_hint" + " [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8, %9;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "h"(cluster_mask), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + +#ifdef KITTENS_SM10X + if(dst_mbar_cta != -1) { + uint32_t neighbor_mbar_ptr; + asm volatile ( + "mapa.shared::cluster.u32 %0, %1, %2;\n" + : "=r"(neighbor_mbar_ptr) + : "r"(mbar_ptr), "r"(dst_mbar_cta) + ); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.cta_group::2.multicast::cluster" + " [%0], [%1, {%3, %4, %5, %6}], [%2], %7;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(neighbor_mbar_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "h"(cluster_mask) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.cta_group::2.multicast::cluster.L2::cache_hint" + " [%0], [%1, {%3, %4, %5, %6}], [%2], %7, %8;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(neighbor_mbar_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "h"(cluster_mask), "l"(make_cache_policy()) + : "memory" + ); + } + } else +#endif + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%3, %4, %5, %6}], [%2], %7;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "h"(cluster_mask) + : "memory" + ); + } else { + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster.L2::cache_hint" + " [%0], [%1, {%3, %4, %5, %6}], [%2], %7, %8;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "h"(cluster_mask), "l"(make_cache_policy()) + : "memory" + ); + } + } +} +#ifdef KITTENS_SM10X +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) { + load_async(dst, src, idx, bar, cluster_mask, dst_mbar_cta); +} +#else +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask) { + load_async(dst, src, idx, bar, cluster_mask); +} +#endif + +} +} + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/vec/tma.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/vec/tma.cuh new file mode 100644 index 0000000..f288ef4 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/vec/tma.cuh @@ -0,0 +1,418 @@ + + + + + +#pragma once + +#include "../../../../common/common.cuh" +#include "../../../../types/types.cuh" +#include "../../util/util.cuh" + + +#define __KITTENS_TMA_DEFINE_DEFAULT_LOAD_CACHE_VEC__(function_name) \ +template> \ +__device__ static inline void function_name(SV &dst, const GL &src, const COORD &idx) { \ + function_name(dst, src, idx); \ +} +#define __KITTENS_TMA_DEFINE_PGL_DEFAULT_LOAD_CACHE_VEC__(function_name) \ +template> \ +__device__ static inline void function_name(SV &dst, const PGL &src, const COORD &idx) { \ + function_name(dst, src, idx); \ +} +#define __KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(function_name) \ +template> \ +__device__ static inline void function_name(const GL &dst, const SV &src, const COORD &idx) { \ + function_name(dst, src, idx); \ +} +#define __KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(function_name) \ +template> \ +__device__ static inline void function_name(const PGL &dst, const SV &src, const COORD &idx) { \ + function_name(dst, src, idx); \ +} +#define __KITTENS_TMA_DEFINE_SEMAPHORE_CACHE_VEC__(function_name) \ +template> \ +__device__ static inline void function_name(SV &dst, const GL &src, const COORD &idx, semaphore& bar) { \ + function_name(dst, src, idx, bar); \ +} +#define __KITTENS_TMA_DEFINE_PGL_SEMAPHORE_CACHE_VEC__(function_name) \ +template> \ +__device__ static inline void function_name(SV &dst, const PGL &src, const COORD &idx, semaphore& bar) { \ + function_name(dst, src, idx, bar); \ +} +#define __KITTENS_TMA_DEFINE_CLUSTER_SEMAPHORE_CACHE_VEC__(function_name) \ +template> \ +__device__ static inline void function_name(SV &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) { \ + function_name(dst, src, idx, bar, cluster_mask, dst_mbar_cta); \ +} +#define __KITTENS_TMA_DEFINE_PGL_CLUSTER_SEMAPHORE_CACHE_VEC__(function_name) \ +template> \ +__device__ static inline void function_name(SV &dst, const PGL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) { \ + function_name(dst, src, idx, bar, cluster_mask, dst_mbar_cta); \ +} + + +namespace kittens { + +namespace detail { +namespace tma { + +template __device__ static inline void vec_prefetch_tma_internal(uint64_t tma_ptr, coord<> tma_coord) { + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.prefetch.tensor.4d.L2.global.tile" + " [%0, {%1, %2, %3, %4}];" + : + : "l"(tma_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.prefetch.tensor.4d.L2.global.tile.L2::cache_hint" + " [%0, {%1, %2, %3, %4}], %5;" + : + : "l"(tma_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } +} + +template __device__ static inline void vec_store_async_tma_internal(uint64_t tma_ptr, uint32_t src_i_ptr, coord<> tma_coord) { + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } +} + +template __device__ static inline void vec_store_add_async_tma_internal(uint64_t tma_ptr, uint32_t src_i_ptr, coord<> tma_coord) { + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.add.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.add.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } +} + +template __device__ static inline void vec_store_min_async_tma_internal(uint64_t tma_ptr, uint32_t src_i_ptr, coord<> tma_coord) { + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.min.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.min.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } +} + +template __device__ static inline void vec_store_max_async_tma_internal(uint64_t tma_ptr, uint32_t src_i_ptr, coord<> tma_coord) { + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.max.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.max.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } +} + +template __device__ static inline void vec_load_async_tma_internal(uint64_t tma_ptr, uint32_t dst_i_ptr, uint32_t mbar_ptr, coord<> tma_coord) { + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5, %6}], [%2];" + : + : "r"(dst_i_ptr), "l"(tma_ptr), "r"(mbar_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%3, %4, %5, %6}], [%2], %7;" + : + : "r"(dst_i_ptr), "l"(tma_ptr), "r"(mbar_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } +} + +namespace cluster { +template __device__ static inline void vec_load_async_tma_internal(uint64_t tma_ptr, uint32_t dst_i_ptr, uint32_t mbar_ptr, coord<> tma_coord, uint16_t cluster_mask, int dst_mbar_cta=-1) { +#ifdef KITTENS_SM10X + if(dst_mbar_cta != -1) { + uint32_t neighbor_mbar_ptr; + asm volatile ( + "mapa.shared::cluster.u32 %0, %1, %2;\n" + : "=r"(neighbor_mbar_ptr) + : "r"(mbar_ptr), "r"(dst_mbar_cta) + ); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.cta_group::2.multicast::cluster" + " [%0], [%1, {%3, %4, %5, %6}], [%2], %7;" + : + : "r"(dst_i_ptr), "l"(tma_ptr), "r"(neighbor_mbar_ptr), + "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "h"(cluster_mask) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.cta_group::2.multicast::cluster.L2::cache_hint" + " [%0], [%1, {%3, %4, %5, %6}], [%2], %7, %8;" + : + : "r"(dst_i_ptr), "l"(tma_ptr), "r"(neighbor_mbar_ptr), + "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "h"(cluster_mask), "l"(make_cache_policy()) + : "memory" + ); + } + } else +#endif + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%3, %4, %5, %6}], [%2], %7;" + : + : "r"(dst_i_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "h"(cluster_mask) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster.L2::cache_hint" + " [%0], [%1, {%3, %4, %5, %6}], [%2], %7, %8;" + : + : "r"(dst_i_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "h"(cluster_mask), "l"(make_cache_policy()) + : "memory" + ); + } +} +} + +} +} + +namespace tma { + +template> +__device__ static inline void prefetch(SV &dst, const GL &src, const COORD &idx) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(src.template get_tma()); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + ::kittens::detail::tma::vec_prefetch_tma_internal(tma_ptr, tma_coord); + } +} +__KITTENS_TMA_DEFINE_DEFAULT_LOAD_CACHE_VEC__(prefetch) + +template> +__device__ static inline void store_async(const GL &dst, const SV &src, const COORD &idx) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + ::kittens::tma::store_commit_group(); +} +__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_async) + +template> +__device__ static inline void store_async(const PGL &dst, const SV &src, const COORD &idx) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + ::kittens::tma::store_commit_group(); +} +__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_async) + +template> +__device__ static inline void store_add_async(const GL &dst, const SV &src, const COORD &idx) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_add_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + ::kittens::tma::store_commit_group(); +} +__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_add_async) + +template> +__device__ static inline void store_add_async(const PGL &dst, const SV &src, const COORD &idx) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_add_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + ::kittens::tma::store_commit_group(); +} +__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_add_async) + +template> +__device__ static inline void store_min_async(const GL &dst, const SV &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_min_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + ::kittens::tma::store_commit_group(); +} +__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_min_async) + +template> +__device__ static inline void store_min_async(const PGL &dst, const SV &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_min_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + ::kittens::tma::store_commit_group(); +} +__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_min_async) + +template> +__device__ static inline void store_max_async(const GL &dst, const SV &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_max_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + ::kittens::tma::store_commit_group(); +} +__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_max_async) + +template> +__device__ static inline void store_max_async(const PGL &dst, const SV &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_max_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + ::kittens::tma::store_commit_group(); +} +__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_max_async) + +template> +__device__ static inline void load_async(SV &dst, const GL &src, const COORD &idx, semaphore& bar) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(src.template get_tma()); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(&bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst)); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t dst_i_ptr = dst_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_load_async_tma_internal(tma_ptr, dst_i_ptr, mbar_ptr, tma_coord); + } +} +__KITTENS_TMA_DEFINE_SEMAPHORE_CACHE_VEC__(load_async) + +namespace cluster { +template> +__device__ static inline void load_async(SV &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(src.template get_tma()); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(&bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst)); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t dst_i_ptr = dst_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::cluster::vec_load_async_tma_internal(tma_ptr, dst_i_ptr, mbar_ptr, tma_coord, cluster_mask, dst_mbar_cta); + } +} +__KITTENS_TMA_DEFINE_CLUSTER_SEMAPHORE_CACHE_VEC__(load_async) +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/vec/vec.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/vec/vec.cuh new file mode 100644 index 0000000..8730f36 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/vec/vec.cuh @@ -0,0 +1,10 @@ + + + + + +#pragma once + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +#include "tma.cuh" +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/mma/mma.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/mma/mma.cuh new file mode 100644 index 0000000..9bf9e9c --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/mma/mma.cuh @@ -0,0 +1,10 @@ + + + + + +#pragma once + +#if defined(KITTENS_SM10X) +#include "tcgen05.cuh" +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/mma/tcgen05.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/mma/tcgen05.cuh new file mode 100644 index 0000000..6fc4f05 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/mma/tcgen05.cuh @@ -0,0 +1,985 @@ + + + + + +#pragma once + +#include "../../../common/common.cuh" +#include "../../../types/types.cuh" +#include "../util/sync.cuh" + +namespace kittens { +namespace detail { +namespace tcgen05 { + + +template +__device__ static inline constexpr uint32_t instruction_descriptor() { + uint32_t desc = 0; + if constexpr (std::is_same_v || std::is_same_v) { + + static_assert(std::is_same_v || std::is_same_v); + desc |= 0b00 << 0; + desc |= 0b0 << 2; + desc |= 0b0 << 3; + if constexpr (std::is_same_v) { + desc |= 0b01 << 4; + } + else { + desc |= 0b00 << 4; + } + desc |= 0b0 << 6; + if constexpr (std::is_same_v) { + desc |= 0b000 << 7; + desc |= 0b000 << 10; + } else if constexpr (std::is_same_v) { + desc |= 0b001 << 7; + desc |= 0b001 << 10; + } + if constexpr (neg) { + desc |= 0b1 << 13; + } + else { + desc |= 0b0 << 13; + } + desc |= 0b0 << 14; + if constexpr (trans_a) { + desc |= 0b1 << 15; + } + else { + desc |= 0b0 << 15; + } + if constexpr (trans_b) { + desc |= 0b1 << 16; + } + else { + desc |= 0b0 << 16; + } + desc |= (N >> 3) << 17; + desc |= 0b0 << 23; + desc |= (M >> 4) << 24; + desc |= 0b0 << 29; + desc |= 0b00 << 30; + } else if constexpr (std::is_same_v || std::is_same_v) { + static_assert(std::is_same_v, "Integer MMA must accumulate to int32."); + desc |= 0b00 << 0; + desc |= 0b0 << 2; + desc |= 0b1 << 3; + desc |= 0b10 << 4; + desc |= 0b0 << 6; + if constexpr (std::is_same_v) { + desc |= 0b000 << 7; + desc |= 0b000 << 10; + } else if constexpr (std::is_same_v) { + desc |= 0b001 << 7; + desc |= 0b001 << 10; + } else { + static_assert(sizeof(AB) == 999, "Invalid AB type."); + } + static_assert(!neg, "Negation is not supported for integer matrix multiplies."); + desc |= 0b0 << 13; + desc |= 0b0 << 14; + if constexpr (trans_a) { + desc |= 0b1 << 15; + } + else { + desc |= 0b0 << 15; + } + if constexpr (trans_b) { + desc |= 0b1 << 16; + } + else { + desc |= 0b0 << 16; + } + desc |= (N >> 3) << 17; + desc |= 0b0 << 23; + desc |= (M >> 4) << 24; + desc |= 0b0 << 29; + desc |= 0b00 << 30; + } else if constexpr (std::is_same_v || std::is_same_v || std::is_same_v) { + static_assert(std::is_same_v || std::is_same_v); + desc |= 0b00 << 0; + desc |= 0b0 << 2; + desc |= 0b0 << 3; + if constexpr (std::is_same_v) { + desc |= 0b01 << 4; + } + else { + desc |= 0b00 << 4; + } + desc |= 0b0 << 6; + if constexpr (std::is_same_v) { + desc |= 0b000 << 7; + desc |= 0b000 << 10; + } else if constexpr (std::is_same_v) { + desc |= 0b001 << 7; + desc |= 0b001 << 10; + } else if constexpr (std::is_same_v) { + desc |= 0b101 << 7; + desc |= 0b101 << 10; + } + if constexpr (neg) { + desc |= 0b1 << 13; + } + else { + desc |= 0b0 << 13; + } + desc |= 0b0 << 14; + if constexpr (trans_a) { + desc |= 0b1 << 15; + } + else { + desc |= 0b0 << 15; + } + if constexpr (trans_b) { + desc |= 0b1 << 16; + } + else { + desc |= 0b0 << 16; + } + desc |= (N >> 3) << 17; + desc |= 0b0 << 23; + desc |= (M >> 4) << 24; + desc |= 0b0 << 29; + desc |= 0b00 << 30; + } else { + static_assert(sizeof(AB) == 999, "Invalid AB type size; not implemented yet."); + } + return desc; +}; + +template +__device__ static inline constexpr uint32_t instruction_descriptor() { + + static_assert(std::is_same_v || std::is_same_v, "AB must be fp8e4m3 for f4e2m1"); + static_assert(std::is_same_v || std::is_same_v, "SAB must be either fp8e4m3 or fp8e8m0"); +#if !defined(KITTENS_SM103) + static_assert(mma_k_bytes == 32, "32B block-scaled MMA is only supported on non-SM103."); +#endif + static_assert(mma_k_bytes == 32 || (std::is_same_v && mma_k_bytes == 48), "48B (K96) block-scaled MMA is only supported for packed NVFP4 (E4M3 or E8M0 scales)."); + constexpr int scale_type = std::is_same_v ? 0 : std::is_same_v ? 1 : -1; + + uint32_t desc = 0; + desc |= 0b00 << 0; + desc |= 0b0 << 2; + desc |= 0b0 << 3; + desc |= scale_factor_id << 4; + desc |= 0b0 << 6; + if constexpr (std::is_same_v) { + desc |= (0b000 << 7); + desc |= (0b000 << 10); + } else if constexpr (std::is_same_v) { + desc |= 0b001 << 7; + desc |= 0b01 << 10; + desc |= 0b0 << 12; + } else { + static_assert(sizeof(AB) == 999, "Invalid AB type."); + } + if constexpr (neg) { + desc |= 0b1 << 13; + } + else { + desc |= 0b0 << 13; + } + desc |= 0b0 << 14; + desc |= 0b0 << 15; + desc |= 0b0 << 16; + desc |= (N >> 3) << 17; + desc |= scale_type << 23; + desc |= 0b000 << 24; + desc |= (M >> 7) << 27; + desc |= scale_factor_id << 29; + desc |= (mma_k_bytes == 48 ? 1u : 0u) << 31; + + return desc; +} + +template +__device__ static inline void tt_st(uint32_t d_tt_addr, uint32_t a_tt_addr, uint64_t b_desc, uint32_t idesc) { + if constexpr (std::is_same_v || std::is_same_v) { + if constexpr (ncta == 1) { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::1.kind::i8 [%0], [%1], %2, %3, p;}\n" + :: "r"(d_tt_addr), "r"(a_tt_addr), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + else { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::2.kind::i8 [%0], [%1], %2, %3, p;}\n" + :: "r"(d_tt_addr), "r"(a_tt_addr), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + } else if constexpr (std::is_same_v || std::is_same_v || std::is_same_v) { + if constexpr (ncta == 1) { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [%1], %2, %3, p;}\n" + :: "r"(d_tt_addr), "r"(a_tt_addr), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + else { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], [%1], %2, %3, p;}\n" + :: "r"(d_tt_addr), "r"(a_tt_addr), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + } else { + if constexpr (ncta == 1) { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::1.kind::f16 [%0], [%1], %2, %3, p;}\n" + :: "r"(d_tt_addr), "r"(a_tt_addr), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + else { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::2.kind::f16 [%0], [%1], %2, %3, p;}\n" + :: "r"(d_tt_addr), "r"(a_tt_addr), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + } +} + +template +__device__ static inline void st_st(uint32_t d_tt_addr, uint64_t a_desc, uint64_t b_desc, uint32_t idesc) { + if constexpr (std::is_same_v || std::is_same_v) { + if constexpr (ncta == 1) { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::1.kind::i8 [%0], %1, %2, %3, p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + else { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::2.kind::i8 [%0], %1, %2, %3, p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + } else if constexpr (std::is_same_v || std::is_same_v || std::is_same_v) { + if constexpr (ncta == 1) { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + else { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], %1, %2, %3, p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + } else { + if constexpr (ncta == 1) { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + else { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + } +} + +template +__device__ static inline void st_st(uint32_t d_tt_addr, uint64_t a_desc, uint64_t b_desc, uint32_t sa_tt_addr, uint32_t sb_tt_addr, uint32_t idesc) { + static_assert(std::is_same_v || std::is_same_v, "AB must be fp8e4m3 for f4e2m1"); + if constexpr (ncta == 1) { + if constexpr (std::is_same_v) { + asm volatile( + "{.reg .pred p;\n\t" \ + "setp.eq.u32 p, 1, %6;\n\t" \ + "tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [%0], %1, %2, %3, [%4], [%5], p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "r"(sa_tt_addr), "r"(sb_tt_addr), "n"(acc) + ); + } else if constexpr (std::is_same_v) { + if constexpr (block_size == 32) { + asm volatile( + "{.reg .pred p;\n\t" \ + "setp.eq.u32 p, 1, %6;\n\t" \ + "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::2X [%0], %1, %2, %3, [%4], [%5], p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "r"(sa_tt_addr), "r"(sb_tt_addr), "n"(acc) + ); + } + else { + asm volatile( + "{.reg .pred p;\n\t" \ + "setp.eq.u32 p, 1, %6;\n\t" \ + "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::4X [%0], %1, %2, %3, [%4], [%5], p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "r"(sa_tt_addr), "r"(sb_tt_addr), "n"(acc) + ); + } + } else { + static_assert(sizeof(AB) == 999, "Invalid AB type."); + } + } + else { + if constexpr (std::is_same_v) { + asm volatile( + "{.reg .pred p;\n\t" \ + "setp.eq.u32 p, 1, %6;\n\t" \ + "tcgen05.mma.cta_group::2.kind::mxf8f6f4.block_scale.scale_vec::1X [%0], %1, %2, %3, [%4], [%5], p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "r"(sa_tt_addr), "r"(sb_tt_addr), "n"(acc) + ); + } else if constexpr (std::is_same_v) { + if constexpr (block_size == 32) { + asm volatile( + "{.reg .pred p;\n\t" \ + "setp.eq.u32 p, 1, %6;\n\t" \ + "tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::2X [%0], %1, %2, %3, [%4], [%5], p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "r"(sa_tt_addr), "r"(sb_tt_addr), "n"(acc) + ); + } + else { + asm volatile( + "{.reg .pred p;\n\t" \ + "setp.eq.u32 p, 1, %6;\n\t" \ + "tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X [%0], %1, %2, %3, [%4], [%5], p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "r"(sa_tt_addr), "r"(sb_tt_addr), "n"(acc) + ); + } + } else { + static_assert(sizeof(AB) == 999, "Invalid AB type."); + } + } +} + +template +__device__ static inline void commit(kittens::semaphore &sem, uint16_t dst_cta_mask = 0b11) { + if constexpr (ncta == 1) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one.b64 [%0];\n" + :: "l"(__cvta_generic_to_shared(&sem))); + } + else { + asm volatile( + "tcgen05.commit.cta_group::2.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [%0], %1;\n" + :: "l"(__cvta_generic_to_shared(&sem)), "h"(dst_cta_mask)); + } +} + +} +} + + +template constexpr int reduction_dimension = sizeof(T_AB) == 2 ? 16 : sizeof(T_AB) == 4 ? 8 : 32; + + +template +__device__ static inline void mma(D &d, const A &a, const B &b) { + constexpr int trans_b = 1 - n_trans_b; + + + constexpr int M = (trans_a ? A::cols : A::rows) * ncta; + static_assert(M == D::rows*ncta && ((ncta == 1 && (M == 64 || M == 128)) || (ncta == 2 && (M == 128 || M == 256)))); + + constexpr int N = (trans_b ? B::cols : B::rows) * ncta; + static_assert(N == D::cols); + + constexpr int K = trans_a ? A::rows : A::cols; + static_assert((trans_b ? B::rows : B::cols) == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; static_assert(std::is_same_v); + using T_D = D::T; + + constexpr int red_dim = reduction_dimension; + static_assert(K%red_dim == 0, "K dimension must be divisible by red_dim."); + + static_assert( + + (std::is_same_v && ( + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v + )) || + + (std::is_same_v && ( + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v + )) || + + (std::is_same_v && ( + std::is_same_v || + std::is_same_v + )), + "Currently unsupported type combination for matrix multiply." + ); + + static_assert(trans_a == transpose::N, "A from TMEM can't be transposed in TS (tensor-shared) MMA."); + + uint32_t idesc = detail::tcgen05::instruction_descriptor(); + kittens::st_descriptor, trans_b> b_desc(b); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + + detail::tcgen05::template tt_st( + d.addr, + a.template chunk_addr(0), + b_desc.chunk_descriptor(0), + idesc + ); + + #pragma unroll + for(int i = 1; i < K/red_dim; i++) { + detail::tcgen05::template tt_st( + d.addr, + a.template chunk_addr(i), + b_desc.chunk_descriptor(i), + idesc + ); + } +} +template +__device__ static inline void mma(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b); + detail::tcgen05::commit(sem); +} + + +template +__device__ static inline void mma(D &d, const A &a, const B &b) { + constexpr int trans_b = 1 - n_trans_b; + + + constexpr int M = (trans_a ? A::cols : A::rows) * ncta; + static_assert(M == D::rows*ncta && ((ncta == 1 && (M == 64 || M == 128)) || (ncta == 2 && (M == 128 || M == 256)))); + + constexpr int N = (trans_b ? B::cols : B::rows) * ncta; + static_assert(N == D::cols); + + + constexpr int K = (trans_a ? A::rows : A::cols); + static_assert((trans_b ? B::rows : B::cols) == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; static_assert(std::is_same_v); + using T_D = D::T; + + constexpr int red_dim = reduction_dimension; + static_assert(K%red_dim == 0, "K dimension must be divisible by red_dim."); + + static_assert( + + (std::is_same_v && ( + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v + )) || + + (std::is_same_v && ( + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v + )) || + + (std::is_same_v && ( + std::is_same_v || + std::is_same_v + )), + "Currently unsupported type combination for matrix multiply." + ); + + uint32_t idesc = detail::tcgen05::instruction_descriptor(); + kittens::st_descriptor, trans_a> a_desc(a); + kittens::st_descriptor, trans_b> b_desc(b); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + + detail::tcgen05::template st_st( + d.addr, + a_desc.chunk_descriptor(0), + b_desc.chunk_descriptor(0), + idesc + ); + #pragma unroll + for(int i = 1; i < K/red_dim; i++) { + detail::tcgen05::template st_st( + d.addr, + a_desc.chunk_descriptor(i), + b_desc.chunk_descriptor(i), + idesc + ); + } +} +template +__device__ static inline void mma(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b); + detail::tcgen05::commit(sem); +} + + +template +__device__ static inline void mma(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + + + static_assert(std::is_same_v); + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "A and B must be fp8e4m3 or fp4e2m1_2" + ); + static_assert( + (std::is_same_v && ( + std::is_same_v && std::is_same_v + )) || + (std::is_same_v && ( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) + )), + "SAB must be fp8e8m0 for fp8e4m3 element type, or fp8e4m3 / fp8e8m0 for fp4e2m1_2 element type"); + + static_assert(std::is_same_v, "Only float32 accumulator is supported for microscaling formats"); + using T_AB = A::T; + using T_SAB = SA::T; + using T_D = D::T; + + constexpr int block_size = std::is_same_v ? 16 : 32; + constexpr int trans_b = 1 - n_trans_b; + static_assert(trans_a == transpose::N && trans_b == transpose::N, "Only ABt supported for microscaling formats currently"); +#if !defined(KITTENS_SM103) + static_assert(mma_k_bytes == 32, "32B NVFP4 tcgen05 MMA is only supported on non-SM103."); +#endif + static_assert(mma_k_bytes == 32 || mma_k_bytes == 48, "Microscaling MMA K dimension must be 32 or 48 bytes."); + if constexpr (mma_k_bytes == 48) { + static_assert(std::is_same_v, "K96 microscaling MMA is only supported for packed FP4."); + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "K96 expects matching E8M0 (MXFP4, block32) or E4M3 (NVFP4, block16) scale factors."); + } + + + constexpr int M = (trans_a ? A::cols : A::rows) * ncta; + constexpr int N = (trans_b ? B::cols : B::rows) * ncta; + constexpr int K = std::is_same_v ? (trans_a ? A::rows : A::cols) * 2 : (trans_a ? A::rows : A::cols); + constexpr bool is_k96 = std::is_same_v && mma_k_bytes == 48; + constexpr int red_dim = std::is_same_v ? mma_k_bytes * 2 : mma_k_bytes; + static_assert(K % red_dim == 0, "K dimension must be divisible by red_dim."); + + + static_assert(M == D::rows*ncta && ((ncta == 1 && M == 128) || (ncta == 2 && (M == 128 || M == 256)))); + + + static_assert(N == D::cols && ((ncta == 1 && N%8 == 0) || (ncta == 2 && N%16 == 0))); + + + kittens::st_descriptor, trans_a> a_desc(a); + kittens::st_descriptor, trans_b> b_desc(b); + + + + + + + constexpr uint32_t idescs[4] = { + detail::tcgen05::instruction_descriptor(), + detail::tcgen05::instruction_descriptor(), + detail::tcgen05::instruction_descriptor(), + detail::tcgen05::instruction_descriptor() + }; + + + constexpr int N_offset = N / 32; + constexpr int M_offset = M / 32 / ncta; + + uint64_t a_desc0, b_desc0; + if constexpr (is_k96) { + a_desc0 = a_desc.chunk_descriptor_k96(0); + b_desc0 = b_desc.chunk_descriptor_k96(0); + } else { + a_desc0 = a_desc.chunk_descriptor(0); + b_desc0 = b_desc.chunk_descriptor(0); + } + detail::tcgen05::template st_st( + d.addr, + a_desc0, + b_desc0, + sa.addr, + sb.addr, + idescs[0] + ); + + if constexpr (is_k96) { + constexpr int sf_atoms = ((96 / block_size) + 3) / 4; + static_assert((K / red_dim) * M_offset * sf_atoms <= SA::cols / 4, "A scale tensor tile is too narrow for the K96 MMA count."); + static_assert((K / red_dim) * N_offset * sf_atoms <= SB::cols / 4, "B scale tensor tile is too narrow for the K96 MMA count."); + #pragma unroll + for (int i = 1; i < K / red_dim; i++) { + detail::tcgen05::template st_st( + d.addr, + a_desc.chunk_descriptor_k96(i), + b_desc.chunk_descriptor_k96(i), + sa.addr + i * M_offset * sf_atoms, + sb.addr + i * N_offset * sf_atoms, + idescs[0] + ); + } + } else if constexpr (std::is_same_v) { + #pragma unroll + for (int i = 1; i < K / red_dim; i++) { + detail::tcgen05::template st_st( + d.addr, + a_desc.chunk_descriptor(i), + b_desc.chunk_descriptor(i), + sa.addr + (i >> 2) * M_offset, + sb.addr + (i >> 2) * N_offset, + idescs[i % 4] + ); + } + } else if constexpr (std::is_same_v && block_size == 16) { + #pragma unroll + for (int i = 1; i < K / red_dim; i++) { + detail::tcgen05::template st_st( + d.addr, + a_desc.chunk_descriptor(i), + b_desc.chunk_descriptor(i), + sa.addr + i * M_offset, + sb.addr + i * N_offset, + idescs[0] + ); + } + } else if constexpr (std::is_same_v && block_size == 32) { + #pragma unroll + for (int i = 1; i < K / red_dim; i++) { + detail::tcgen05::template st_st( + d.addr, + a_desc.chunk_descriptor(i), + b_desc.chunk_descriptor(i), + sa.addr + (i >> 1) * M_offset, + sb.addr + (i >> 1) * N_offset, + (i & 1) ? idescs[2] : idescs[0] + ); + } + } else { + static_assert(sizeof(T_AB) == 999, "Should not reach here."); + } +} +template +__device__ static inline void mma(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma(d, a, b, sa, sb); + detail::tcgen05::commit(sem); +} + + +template +__device__ static inline void mma2(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mma2(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mma2(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + static_assert(!trans_a && trans_b, "Only ABt supported for microscaling formats currently"); + mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mma2(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + static_assert(!trans_a && trans_b, "Only ABt supported for microscaling formats currently"); + mma(d, a, b, sa, sb); +} +template +__device__ static inline void mm(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + static_assert(!trans_a && trans_b, "Only ABt supported for microscaling formats currently"); + mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mm(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + static_assert(!trans_a && trans_b, "Only ABt supported for microscaling formats currently"); + mma(d, a, b, sa, sb); +} +template +__device__ static inline void mm2(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mm2(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mm2(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + static_assert(!trans_a && trans_b, "Only ABt supported for microscaling formats currently"); + mma2(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mm2(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + static_assert(!trans_a && trans_b, "Only ABt supported for microscaling formats currently"); + mma2(d, a, b, sa, sb); +} + + +template +__device__ static inline void mma_AB(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mma_AB(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mma2_AB(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mma2_AB(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mma_ABt(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mma_ABt(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mma2_ABt(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mma2_ABt(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mma_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mma_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma(d, a, b, sa, sb); +} +template +__device__ static inline void mma2_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mma2_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma(d, a, b, sa, sb); +} +template +__device__ static inline void mma_AtB(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mma_AtB(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mma2_AtB(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mma2_AtB(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mma_AtBt(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mma_AtBt(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mma2_AtBt(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mma2_AtBt(D &d, const A &a, const B &b) { + mma2(d, a, b); +} + +template +__device__ static inline void mm_AB(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm_AB(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm2_AB(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mm2_AB(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mm_ABt(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm_ABt(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm2_ABt(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mm2_ABt(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mm_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mm_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma(d, a, b, sa, sb); +} +template +__device__ static inline void mm2_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mm2_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma(d, a, b, sa, sb); +} +template +__device__ static inline void mm_AtB(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm_AtB(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm2_AtB(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mm2_AtB(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mm_AtBt(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm_AtBt(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm2_AtBt(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mm2_AtBt(D &d, const A &a, const B &b) { + mma2(d, a, b); +} + +template +__device__ static inline void mma_chunk(D &d, const A &a, const B &b, const SA &sa, const SB &sb, int ab_chunk_idx, int sf_chunk_idx, bool init) { + static_assert(mma_k_bytes == 48, "Chunked block-scaled MMA currently only supports the packed-FP4 48B (K96) form."); + static_assert(std::is_same_v && std::is_same_v, "A and B must be fp4e2m1_2"); + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "SAB must be fp8e4m3 or fp8e8m0"); + static_assert(std::is_same_v, "Only float32 accumulator is supported for microscaling formats"); + using T_AB = A::T; + using T_SAB = SA::T; + using T_D = D::T; + + constexpr int block_size = std::is_same_v ? 16 : 32; + constexpr int trans_b = 1 - n_trans_b; + static_assert(trans_a == transpose::N && trans_b == transpose::N, "Only ABt supported for microscaling formats currently"); + + + constexpr int M = (trans_a ? A::cols : A::rows) * ncta; + constexpr int N = (trans_b ? B::cols : B::rows) * ncta; + + + kittens::st_descriptor, trans_a> a_desc(a); + kittens::st_descriptor, trans_b> b_desc(b); + + + constexpr uint32_t idescs[4] = { + detail::tcgen05::instruction_descriptor(), + detail::tcgen05::instruction_descriptor(), + detail::tcgen05::instruction_descriptor(), + detail::tcgen05::instruction_descriptor() + }; + + + constexpr int N_offset = N / 32; + constexpr int M_offset = M / 32 / ncta; + + + + uint32_t sa_off, sb_off, idesc; + if constexpr (tight_sf) { + + + const int start_block = sf_chunk_idx * (96 / block_size); + sa_off = (start_block / 4) * M_offset; + sb_off = (start_block / 4) * N_offset; + idesc = idescs[start_block % 4]; + } else { + + constexpr int sf_atoms = ((96 / block_size) + 3) / 4; + sa_off = sf_chunk_idx * M_offset * sf_atoms; + sb_off = sf_chunk_idx * N_offset * sf_atoms; + idesc = idescs[0]; + } + + const uint64_t ad = a_desc.chunk_descriptor_k96(ab_chunk_idx); + const uint64_t bd = b_desc.chunk_descriptor_k96(ab_chunk_idx); + if (init) + detail::tcgen05::template st_st(d.addr, ad, bd, sa.addr + sa_off, sb.addr + sb_off, idesc); + else + detail::tcgen05::template st_st(d.addr, ad, bd, sa.addr + sa_off, sb.addr + sb_off, idesc); +} +template +__device__ static inline void mma2_ABt_chunk(D &d, const A &a, const B &b, const SA &sa, const SB &sb, int ab_chunk_idx, int sf_chunk_idx, bool init) { + mma_chunk(d, a, b, sa, sb, ab_chunk_idx, sf_chunk_idx, init); +} + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/thread.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/thread.cuh new file mode 100644 index 0000000..aa24901 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/thread.cuh @@ -0,0 +1,11 @@ + + + + + +#pragma once + + +#include "memory/memory.cuh" +#include "mma/mma.cuh" +#include "util/util.cuh" diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/util/sync.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/util/sync.cuh new file mode 100644 index 0000000..875d82a --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/util/sync.cuh @@ -0,0 +1,331 @@ + + + + + +#pragma once + +#include "../../../common/common.cuh" +#include "../../../types/types.cuh" + +namespace kittens { + +struct semaphore { +private: + uint64_t value; +}; +template struct barrier { + int barrier_id; + __device__ __forceinline__ barrier(int _id) : barrier_id(_id) {} + __device__ __forceinline__ barrier operator[](int i) { + return barrier(barrier_id + i); + } +}; + + + + + + + + + + + + + +__device__ static inline void init_semaphore(semaphore& bar, int thread_count, int transaction_count=0) { + void const* const ptr = &bar; + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + + asm volatile ( + "mbarrier.init.shared::cta.b64 [%0], %1;\n" + :: "r"(bar_ptr), "r"(thread_count+transaction_count) + ); +} + + + + + +__device__ static inline void invalidate_semaphore(semaphore& bar) { + void const* const ptr = &bar; + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + asm volatile ( + "mbarrier.inval.shared::cta.b64 [%0];\n" + :: "r"(bar_ptr) + ); +} + + + + + + + + +__device__ static inline void arrive(semaphore& sem) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(&sem)); + asm volatile ( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];\n" + : + : "r"(mbar_ptr) + : "memory" + ); +} +template __device__ static inline void arrive(barrier bar) { + asm volatile("bar.arrive %0, %1;\n" :: "r"(bar.barrier_id), "n"(num_warps*WARP_THREADS) : "memory"); +} + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + +__device__ static inline void arrive(semaphore& sem, uint32_t count) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(&sem)); + asm volatile ( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0], %1;\n" + : + : "r"(mbar_ptr), "r"(count) + : "memory" + ); +} +#endif + + + + + + + +__device__ static inline void wait(semaphore& sem, int kPhaseBit) { + void const* const ptr = &sem; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +#else + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.test_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "nanosleep.u32 5;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +#endif +} + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +__device__ static inline bool try_wait(semaphore &sem, int kPhaseBit) { + void const* const ptr = &sem; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + uint32_t success; + + asm volatile( + "{\n" + ".reg .pred P1; \n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2; \n" + "selp.b32 %0, 1, 0, P1; \n" + "}\n" + : "=r"(success) + : "r"(mbar_ptr), "r"(kPhaseBit) + : "memory" + ); + + return static_cast(success); +} +#endif + +__device__ static inline void careful_wait(semaphore& sem, int kPhaseBit) { + void const* const ptr = &sem; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + asm volatile ( + "{\n" + ".reg .b64 start_clock, current_clock;\n" + "mov.b64 start_clock, %clock64;\n" + ".reg .pred P_CLOCK;\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "mov.b64 current_clock, %clock64;\n" + "sub.u64 current_clock, current_clock, start_clock;\n" + "setp.ge.u64 P_CLOCK, current_clock, 1000000;\n" + "@P_CLOCK trap;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +#else + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.test_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "nanosleep.u32 5;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +#endif +} + + + + + + + +__device__ static inline int test_wait(semaphore& sem, int kPhaseBit) { + void const* const ptr = &sem; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + int result; + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "mbarrier.test_wait.parity.shared::cta.b64 P1, [%1], %2;\n" + "selp.u32 %0,1,0,P1;" + "}\n" + : "=r"(result) + : "r"(mbar_ptr), "r"(kPhaseBit) + ); + return result; +} + +__device__ static inline void arrive_and_wait(semaphore& sem, int kPhaseBit) { + arrive(sem); + wait(sem, kPhaseBit); +} +template __device__ static inline void arrive_and_wait(barrier bar) { + asm volatile("bar.sync %0, %1;\n" :: "r"(bar.barrier_id), "n"(num_warps*WARP_THREADS) : "memory"); +} + +template __device__ static inline void load_async_wait() { + if constexpr (N == 0) { + asm volatile("cp.async.wait_all;\n" ::); + } else { + asm volatile("cp.async.wait_group %0;\n" :: "n"(N)); + } + __syncwarp(); +} + + +namespace detail { +template struct size_info { + static constexpr uint32_t bytes = sizeof(std::remove_reference_t); +}; +template struct size_info { + static constexpr uint32_t elements = ST::num_elements; + static constexpr uint32_t bytes = ST::num_elements * sizeof(typename ST::dtype); +}; +template struct size_info { + static constexpr uint32_t elements = SV::length; + static constexpr uint32_t bytes = SV::length * sizeof(typename SV::dtype); +}; +} +template inline constexpr uint32_t size_bytes = 0; +template inline constexpr uint32_t size_bytes = detail::size_info::bytes + size_bytes; + + + +#if defined(KITTENS_SM10X) + +__device__ static inline void tensor_before_thread_sync() { + asm volatile("tcgen05.fence::before_thread_sync;\n"); +} +__device__ static inline void tensor_after_thread_sync() { + asm volatile("tcgen05.fence::after_thread_sync;\n"); +} + +__device__ inline static void tensor_load_wait() { + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} +__device__ inline static void tensor_store_wait() { + asm volatile("tcgen05.wait::st.sync.aligned;"); +} + +template +__device__ static inline void tensor_commit(kittens::semaphore &sem, uint16_t dst_cta_mask = 0b11) { + if constexpr (ncta == 1) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one.b64 [%0];\n" + :: "l"(__cvta_generic_to_shared(&sem))); + } + else { + asm volatile( + "tcgen05.commit.cta_group::2.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [%0], %1;\n" + :: "l"(__cvta_generic_to_shared(&sem)), "h"(dst_cta_mask)); + } +} + +#endif + + + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + +template +__device__ static inline void signal( + const barrier_t &barrier, const coord &idx, const int dst_dev_idx, const int val +) { + asm volatile("{red.release.sys.global.add.s32 [%0], %1;}" :: "l"(&barrier[dst_dev_idx][idx]), "r"(val) : "memory"); +} + +template +__device__ static inline void signal_all( + const barrier_t &barrier, const coord &idx, const int val +) { + asm volatile("{multimem.red.release.sys.global.add.s32 [%0], %1;}" :: "l"(barrier.mc_ptr_at(idx)), "r"(val) : "memory"); +} + +template +__device__ static inline void wait( + const barrier_t &barrier, const coord &idx, const int dev_idx, const int expected +) { + int val; + do { + asm volatile("{ld.relaxed.sys.global.s32 %0, [%1];}" : "=r"(val) : "l"(&barrier[dev_idx][idx]) : "memory"); + } while (val != expected); +} + +template +__device__ static inline void barrier_all( + const barrier_t &barrier, const coord &idx, const int dev_idx +) { + signal_all(barrier, idx, 1); + wait(barrier, idx, dev_idx, NUM_DEVICES); + asm volatile("{red.release.sys.global.add.s32 [%0], %1;}" :: "l"(&barrier[dev_idx][idx]), "r"(-NUM_DEVICES) : "memory"); +} + +#endif + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/util/tma.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/util/tma.cuh new file mode 100644 index 0000000..886aa21 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/util/tma.cuh @@ -0,0 +1,280 @@ + + + + + +#pragma once + +#include "../../../common/common.cuh" +#include "../../../types/types.cuh" +#include "sync.cuh" + +namespace kittens { + + + +namespace tma { + + + + + + + + + + + + + +__device__ static inline void expect_bytes(semaphore& bar, uint32_t bytes) { + void const* const ptr = &bar; + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + + asm volatile ("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" + :: "r"(bar_ptr), "r"(bytes)); +} + + + + + +template +__device__ static inline void expect(semaphore& bar, const T& _1, const args&... _2) { + expect_bytes(bar, size_bytes); +} + + + + + + +__device__ static inline void store_commit_group() { + asm volatile("cp.async.bulk.commit_group;"); +} + + + + + +template +__device__ static inline void store_async_wait() { + asm volatile ( + "cp.async.bulk.wait_group %0;" + : + : "n"(N) + : "memory" + ); +} + + + + + +template +__device__ static inline void store_async_read_wait() { + asm volatile ( + "cp.async.bulk.wait_group.read %0;" + : + : "n"(N) + : "memory" + ); +} + + + +__device__ static inline void load_async(void *dst, void *src, uint32_t size_bytes, semaphore& bar) { + asm volatile ( + "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n" + :: "r"(static_cast(__cvta_generic_to_shared(dst))), + "l"(src), "r"(size_bytes), + "r"(static_cast(__cvta_generic_to_shared(&bar))) + : "memory" + ); +} +template +__device__ static inline void load_async(T &dst, T &src, uint32_t size_bytes, semaphore& bar) { + load_async(reinterpret_cast(&dst), reinterpret_cast(&src), size_bytes, bar); +} + +__device__ static inline void store_async(void *dst, void *src, uint32_t size_bytes) { + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + asm volatile ( + "cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;\n" + :: "l"(dst), "r"(static_cast(__cvta_generic_to_shared(src))), "r"(size_bytes) + : "memory" + ); + store_commit_group(); +} +template +__device__ static inline void store_async(T &dst, T &src, uint32_t size_bytes) { + store_async(reinterpret_cast(&dst), reinterpret_cast(&src), size_bytes); +} + + + +namespace cluster { + + + + + + + +__device__ static inline void wait(semaphore& bar, int kPhaseBit) { + void const* const ptr = &bar; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +} + +__device__ static inline bool try_wait(semaphore &bar, int kPhaseBit) { + void const* const ptr = &bar; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + uint32_t success; + + asm volatile( + "{\n" + ".reg .pred P1; \n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [%1], %2; \n" + "selp.b32 %0, 1, 0, P1; \n" + "}\n" + : "=r"(success) + : "r"(mbar_ptr), "r"(kPhaseBit) + : "memory" + ); + + return static_cast(success); +} + +__device__ static inline void careful_wait(semaphore& bar, int kPhaseBit) { + void const* const ptr = &bar; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + + asm volatile ( + "{\n" + ".reg .b64 start_clock, current_clock;\n" + "mov.b64 start_clock, %clock64;\n" + ".reg .pred P_CLOCK;\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "mov.b64 current_clock, %clock64;\n" + "sub.u64 current_clock, current_clock, start_clock;\n" + "setp.ge.u64 P_CLOCK, current_clock, 1000000;\n" + "@P_CLOCK trap;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +} + + + + + + + + + + + + + + + +__device__ static inline void expect_bytes(semaphore& bar, uint32_t bytes) { + uint32_t mbar_addr = static_cast(__cvta_generic_to_shared(&bar)); + asm volatile ("mbarrier.arrive.expect_tx.shared::cluster.b64 _, [%0], %1;\n" + :: "r"(mbar_addr), "r"(bytes)); +} +__device__ static inline void expect_bytes(semaphore& bar, uint32_t bytes, int dst_cta) { + uint32_t mbar_addr = static_cast(__cvta_generic_to_shared(&bar)); + uint32_t neighbor_mbar_addr; + asm volatile ( + "mapa.shared::cluster.u32 %0, %1, %2;\n" + : "=r"(neighbor_mbar_addr) + : "r"(mbar_addr), "r"(dst_cta) + ); + asm volatile ("mbarrier.arrive.expect_tx.shared::cluster.b64 _, [%0], %1;\n" + :: "r"(neighbor_mbar_addr), "r"(bytes)); +} + + + + + + + + + + + + + + + +template +__device__ static inline void expect(semaphore& bar, const T& _1, const args&... _2) { + expect_bytes(bar, size_bytes); +} + + + + + + + + + + +__device__ static inline void arrive(semaphore& bar, int dst_cta, uint32_t count=1) { + uint32_t mbar_addr = static_cast(__cvta_generic_to_shared(&bar)); + uint32_t neighbor_mbar_addr; + asm volatile ( + "mapa.shared::cluster.u32 %0, %1, %2;\n" + : "=r"(neighbor_mbar_addr) + : "r"(mbar_addr), "r"(dst_cta) + ); + asm volatile ( + "mbarrier.arrive.shared::cluster.b64 _, [%0], %1;\n" + : + : "r"(neighbor_mbar_addr), "r" (count) + : "memory" + ); +} + + + +__device__ static inline void load_async(void *dst, void *src, uint32_t size_bytes, semaphore& bar, uint16_t cta_mask) { + asm volatile ( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster [%0], [%1], %2, [%3], %4;\n" + :: "r"(static_cast(__cvta_generic_to_shared(dst))), + "l"(src), "r"(size_bytes), + "r"(static_cast(__cvta_generic_to_shared(&bar))), + "h"(cta_mask) + : "memory" + ); +} +template +__device__ static inline void load_async(T &dst, T &src, uint32_t size_bytes, semaphore& bar, uint16_t cta_mask) { + load_async(reinterpret_cast(&dst), reinterpret_cast(&src), size_bytes, bar, cta_mask); +} + +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/util/util.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/util/util.cuh new file mode 100644 index 0000000..ad52da6 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/util/util.cuh @@ -0,0 +1,327 @@ + + + + + +#pragma once + +#include "sync.cuh" +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +#include "tma.cuh" +#endif + +namespace kittens { + + + +template struct move { + __device__ static inline void lds(T& dst, uint32_t src); + __device__ static inline void sts(uint32_t dst, const T& src); + __device__ static inline void ldg(T& dst, T* src); + __device__ static inline void stg(T* dst, const T& src); +}; + +template<> struct move { + __device__ static inline void lds(bf16& dst, uint32_t src) { + asm volatile("ld.shared.b16 %0, [%1];\n" : "=h"(*(uint16_t*)&dst) : "r"(src)); + } + __device__ static inline void sts(uint32_t dst, const bf16& src) { + asm volatile("st.shared.b16 [%1], %0;\n" : : "h"(*(uint16_t*)&src), "r"(dst)); + } + __device__ static inline void ldg(bf16& dst, bf16* src) { + asm volatile("ld.global.b16 %0, [%1];\n" : "=h"(*(uint16_t*)&dst) : "l"(src)); + } + __device__ static inline void stg(bf16* dst, const bf16& src) { + asm volatile("st.global.b16 [%1], %0;\n" : : "h"(*(uint16_t*)&src), "l"(dst)); + } +}; +template<> struct move { + __device__ static inline void lds(half& dst, uint32_t src) { + asm volatile("ld.shared.b16 %0, [%1];\n" : "=h"(*(uint16_t*)&dst) : "r"(src)); + } + __device__ static inline void sts(uint32_t dst, const half& src) { + asm volatile("st.shared.b16 [%1], %0;\n" : : "h"(*(uint16_t*)&src), "r"(dst)); + } + __device__ static inline void ldg(half& dst, half* src) { + asm volatile("ld.global.b16 %0, [%1];\n" : "=h"(*(uint16_t*)&dst) : "l"(src)); + } + __device__ static inline void stg(half* dst, const half& src) { + asm volatile("st.global.b16 [%1], %0;\n" : : "h"(*(uint16_t*)&src), "l"(dst)); + } +}; +template<> struct move { + __device__ static inline void lds(float& dst, uint32_t src) { + asm volatile("ld.shared.f32 %0, [%1];\n" : "=f"(dst) : "r"(src)); + } + __device__ static inline void sts(uint32_t dst, const float& src) { + asm volatile("st.shared.f32 [%1], %0;\n" : : "f"(src), "r"(dst)); + } + __device__ static inline void ldg(float& dst, float* src) { + asm volatile("ld.global.f32 %0, [%1];\n" : "=f"(dst) : "l"(src)); + } + __device__ static inline void stg(float* dst, const float& src) { + asm volatile("st.global.f32 [%1], %0;\n" : : "f"(src), "l"(dst)); + } +}; +template<> struct move { + __device__ static inline void lds(int& dst, uint32_t src) { + asm volatile("ld.shared.u32 %0, [%1];\n" : "=r"(dst) : "r"(src)); + } + __device__ static inline void sts(uint32_t dst, const int& src) { + asm volatile("st.shared.u32 [%1], %0;\n" : : "r"(src), "r"(dst)); + } + __device__ static inline void ldg(int& dst, int* src) { + asm volatile("ld.global.u32 %0, [%1];\n" : "=r"(dst) : "l"(src)); + } + __device__ static inline void stg(int* dst, const int& src) { + asm volatile("st.global.u32 [%1], %0;\n" : : "r"(src), "l"(dst)); + } +}; + +template<> struct move { + __device__ static inline void lds(bf16_2& dst, uint32_t src) { + asm volatile("ld.shared.b32 %0, [%1];\n" : "=r"(*(uint32_t*)&dst) : "r"(src)); + } + __device__ static inline void sts(uint32_t dst, const bf16_2& src) { + asm volatile("st.shared.b32 [%1], %0;\n" : : "r"(*(uint32_t*)&src), "r"(dst)); + } + __device__ static inline void ldg(bf16_2& dst, bf16_2* src) { + asm volatile("ld.global.b32 %0, [%1];\n" : "=r"(*(uint32_t*)&dst) : "l"(src)); + } + __device__ static inline void stg(bf16_2* dst, const bf16_2& src) { + asm volatile("st.global.b32 [%1], %0;\n" : : "r"(*(uint32_t*)&src), "l"(dst)); + } + __device__ static inline void ldsm4(bf16_2& dst1, bf16_2& dst2, bf16_2& dst3, bf16_2& dst4, uint32_t src) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" : + "=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src)); + } + __device__ static inline void ldsm4t(bf16_2& dst1, bf16_2& dst2, bf16_2& dst3, bf16_2& dst4, uint32_t src) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" : + "=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src)); + } + __device__ static inline void stsm4(uint32_t dst, bf16_2& src1, bf16_2& src2, bf16_2& src3, bf16_2& src4) { + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" :: + "r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst)); + } + __device__ static inline void stsm4t(uint32_t dst, bf16_2& src1, bf16_2& src2, bf16_2& src3, bf16_2& src4) { + asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" :: + "r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst)); + } +}; +template<> struct move { + __device__ static inline void lds(half_2& dst, uint32_t src) { + asm volatile("ld.shared.b32 %0, [%1];\n" : "=r"(*(uint32_t*)&dst) : "r"(src)); + } + __device__ static inline void sts(uint32_t dst, const half_2& src) { + asm volatile("st.shared.b32 [%1], %0;\n" : : "r"(*(uint32_t*)&src), "r"(dst)); + } + __device__ static inline void ldg(half_2& dst, half_2* src) { + asm volatile("ld.global.b32 %0, [%1];\n" : "=r"(*(uint32_t*)&dst) : "l"(src)); + } + __device__ static inline void stg(half_2* dst, const half_2& src) { + asm volatile("st.global.b32 [%1], %0;\n" : : "r"(*(uint32_t*)&src), "l"(dst)); + } + __device__ static inline void ldsm4(half_2& dst1, half_2& dst2, half_2& dst3, half_2& dst4, uint32_t src) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" : + "=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src)); + } + __device__ static inline void ldsm4t(half_2& dst1, half_2& dst2, half_2& dst3, half_2& dst4, uint32_t src) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" : + "=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src)); + } + __device__ static inline void stsm4(uint32_t dst, half_2& src1, half_2& src2, half_2& src3, half_2& src4) { + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" :: + "r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst)); + } + __device__ static inline void stsm4t(uint32_t dst, half_2& src1, half_2& src2, half_2& src3, half_2& src4) { + asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" :: + "r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst)); + } +}; +template<> struct move { + __device__ static inline void lds(float2& dst, uint32_t src) { + asm volatile("ld.shared.v2.f32 {%0, %1}, [%2];\n" : "=f"(dst.x), "=f"(dst.y) : "r"(src)); + } + __device__ static inline void sts(uint32_t dst, const float2& src) { + asm volatile("st.shared.v2.f32 [%2], {%0, %1};\n" : : "f"(src.x), "f"(src.y), "r"(dst)); + } + __device__ static inline void ldg(float2& dst, float2* src) { + asm volatile("ld.global.v2.f32 {%0, %1}, [%2];\n" : "=f"(dst.x), "=f"(dst.y) : "l"(src)); + } + __device__ static inline void stg(float2* dst, const float2& src) { + asm volatile("st.global.v2.f32 [%2], {%0, %1};\n" : : "f"(src.x), "f"(src.y), "l"(dst)); + } +}; +template<> struct move { + __device__ static inline void lds(float4& dst, uint32_t src) { + asm volatile("ld.shared.v4.f32 {%0, %1, %2, %3}, [%4];\n" : "=f"(dst.x), "=f"(dst.y), "=f"(dst.z), "=f"(dst.w) : "r"(src)); + } + __device__ static inline void sts(uint32_t dst, const float4& src) { + asm volatile("st.shared.v4.f32 [%4], {%0, %1, %2, %3};\n" : : "f"(src.x), "f"(src.y), "f"(src.z), "f"(src.w), "r"(dst)); + } + __device__ static inline void ldg(float4& dst, float4* src) { + asm volatile("ld.global.v4.f32 {%0, %1, %2, %3}, [%4];\n" : "=f"(dst.x), "=f"(dst.y), "=f"(dst.z), "=f"(dst.w) : "l"(src)); + } + __device__ static inline void stg(float4* dst, const float4& src) { + asm volatile("st.global.v4.f32 [%4], {%0, %1, %2, %3};\n" : : "f"(src.x), "f"(src.y), "f"(src.z), "f"(src.w), "l"(dst)); + } +}; +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template<> struct move { + __device__ static inline void ldsm4(fp8e4m3_4& dst1, fp8e4m3_4& dst2, fp8e4m3_4& dst3, fp8e4m3_4& dst4, uint32_t src) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" : + "=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src)); + } + __device__ static inline void stsm4(uint32_t dst, fp8e4m3_4& src1, fp8e4m3_4& src2, fp8e4m3_4& src3, fp8e4m3_4& src4) { + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" :: + "r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst)); + } + +}; +template<> struct move { + __device__ static inline void ldsm4(fp8e5m2_4& dst1, fp8e5m2_4& dst2, fp8e5m2_4& dst3, fp8e5m2_4& dst4, uint32_t src) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" : + "=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src)); + } + __device__ static inline void stsm4(uint32_t dst, fp8e5m2_4& src1, fp8e5m2_4& src2, fp8e5m2_4& src3, fp8e5m2_4& src4) { + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" :: + "r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst)); + } +}; +template<> struct move { + __device__ static inline void ldsm4(int8_4& dst1, int8_4& dst2, int8_4& dst3, int8_4& dst4, uint32_t src) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" : + "=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src)); + } + __device__ static inline void stsm4(uint32_t dst, int8_4& src1, int8_4& src2, int8_4& src3, int8_4& src4) { + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" :: + "r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst)); + } +}; +template<> struct move { + __device__ static inline void ldsm4(uint8_4& dst1, uint8_4& dst2, uint8_4& dst3, uint8_4& dst4, uint32_t src) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" : + "=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src)); + } + __device__ static inline void stsm4(uint32_t dst, uint8_4& src1, uint8_4& src2, uint8_4& src3, uint8_4& src4) { + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" :: + "r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst)); + } +}; +#endif + + + +enum cache_policy { + NORMAL = 0, + EVICT_FIRST = 1, + EVICT_LAST = 2 +}; +template __device__ inline uint64_t make_cache_policy() { + uint64_t cache_policy_val; + constexpr float fraction = 1.0f; + static_assert(policy == cache_policy::EVICT_FIRST || policy == cache_policy::EVICT_LAST, "Unexpected cache policy"); + if constexpr (policy == cache_policy::EVICT_FIRST) { + asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, %1;\n" : "=l"(cache_policy_val) : "f"(fraction)); + } + else { + asm volatile("createpolicy.fractional.L2::evict_last.b64 %0, %1;\n" : "=l"(cache_policy_val) : "f"(fraction)); + } + return cache_policy_val; +} + + + +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + +namespace clc { + +struct handle { + uint4 internal_value; +}; + +struct result { + uint32_t success; + uint32_t x; + uint32_t y; + uint32_t z; +}; + + + + + + + + +__device__ static inline void schedule(handle &h, semaphore &sem) { + asm volatile("{clusterlaunchcontrol.try_cancel.async.shared::cta.mbarrier::complete_tx::bytes.multicast::cluster::all.b128 [%0], [%1];}" + :: "r"(static_cast(__cvta_generic_to_shared(&h.internal_value))), "r"(static_cast(__cvta_generic_to_shared(&sem))) + : "memory" + ); +} + + + + + +__device__ static inline result query(handle &h) { + result r; + asm volatile( + "{\n" + ".reg .pred SUCCESS;\n" + ".reg .b128 CLC_HANDLE;\n" + "ld.shared.b128 CLC_HANDLE, [%4];\n" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 SUCCESS, CLC_HANDLE;\n" + "selp.u32 %0, 1, 0, SUCCESS;\n" + "@!SUCCESS bra.uni DONE;\n" + "clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%1, %2, %3, _}, CLC_HANDLE;\n" + "fence.proxy.async.shared::cta;\n" + "DONE:\n" + "}" + : "=r"(r.success), "=r"(r.x), "=r"(r.y), "=r"(r.z) + : "r"(static_cast(__cvta_generic_to_shared(&h.internal_value))) + : "memory" + ); + return r; +} + +} + +#endif + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + + +namespace pdl { + + + + + + + + +__device__ static inline void arrive() { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +} + + + + +__device__ static inline void wait() { + asm volatile("griddepcontrol.wait;" ::: "memory"); +} + +} + +#endif + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/broker.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/broker.cuh new file mode 100644 index 0000000..427debc --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/broker.cuh @@ -0,0 +1,549 @@ + + + + + + + + + + + + + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64) + #error "KittensBroker is not supported on Windows" +#endif + +namespace kittens { + +namespace detail { +namespace broker { + +static constexpr int MAX_LOCAL_WORLD_SIZE = 72; +static constexpr int VAULT_SIZE_PER_RANK = 64; + +struct KittensVault { + static constexpr int INIT_CODE = 0x43617473; + int init; + int barrier; + int sense; + uint8_t data[MAX_LOCAL_WORLD_SIZE * VAULT_SIZE_PER_RANK]; +}; + +static constexpr int SHM_SIZE = (sizeof(KittensVault) + 4095) / 4096 * 4096; + +__host__ inline static void init_sync( + int local_rank, + volatile KittensVault *vault +) { + if (local_rank == 0) { + + vault->barrier = 0; + vault->sense = 0; + __sync_synchronize(); + vault->init = KittensVault::INIT_CODE; + } else { + while (vault->init != KittensVault::INIT_CODE) usleep(1); + __sync_synchronize(); + } +} + +__host__ inline static void sync( + int local_world_size, + volatile KittensVault *vault +) { + if (vault->init != KittensVault::INIT_CODE) + throw std::runtime_error("KittensBroker: KittensVault not initialized"); + + + int arrived = __sync_add_and_fetch(&vault->barrier, 1); + if (arrived == local_world_size) vault->sense = 1; + while (!vault->sense) usleep(1); + + + __sync_synchronize(); + + + arrived = __sync_add_and_fetch(&vault->barrier, -1); + if (arrived == 0) vault->sense = 0; + while (vault->sense) usleep(1); +} + +__host__ inline void *create_shm(const char *key, size_t size) { + int shm_fd; + shm_fd = shm_open(key, O_RDWR | O_CREAT | O_EXCL | O_CLOEXEC, 0600); + + if (shm_fd < 0) { + if (errno == EEXIST) + throw std::runtime_error("KittensBroker: Named shared memory already exists"); + throw std::runtime_error("KittensBroker: Failed to create shared memory"); + } + + if (ftruncate(shm_fd, size) != 0) { + shm_unlink(key); + close(shm_fd); + throw std::runtime_error("KittensBroker: Failed to truncate shared memory"); + } + + void *addr = mmap(0, size, PROT_READ | PROT_WRITE, MAP_SHARED, shm_fd, 0); + close(shm_fd); + if (addr == MAP_FAILED) { + shm_unlink(key); + throw std::runtime_error("KittensBroker: Failed to map to shared memory"); + } + + return addr; +} + +__host__ inline void *open_shm(const char *key, size_t size) { + int shm_fd; + while (true) { + shm_fd = shm_open(key, O_RDWR | O_CLOEXEC, 0); + if (shm_fd >= 0) + break; + if (errno != ENOENT) + throw std::runtime_error("KittensBroker: Failed to open shared memory"); + usleep(1); + } + + struct stat shm_st; + do { + if (fstat(shm_fd, &shm_st) != 0) { + shm_unlink(key); + close(shm_fd); + throw std::runtime_error("KittensBroker: Failed to open shared memory stats"); + } + usleep(1); + } while ((size_t)shm_st.st_size < size); + + void *addr = mmap(0, size, PROT_READ | PROT_WRITE, MAP_SHARED, shm_fd, 0); + close(shm_fd); + if (addr == MAP_FAILED) { + shm_unlink(key); + throw std::runtime_error("KittensBroker: Failed to map to shared memory"); + } + + return addr; +} + +__host__ inline void unlink_shm(const char *key) { + shm_unlink(key); +} + +__host__ inline void unmap_shm(void *addr, size_t size) { + munmap(addr, size); +} + +__host__ inline int create_socket(const char *key, int local_rank) { + int sock_fd; + if ((sock_fd = socket(AF_UNIX, SOCK_DGRAM | SOCK_CLOEXEC, 0)) < 0) + throw std::runtime_error("KittensBroker: Socket creation error"); + + struct sockaddr_un addr; + memset(&addr, 0, sizeof(addr)); + addr.sun_family = AF_UNIX; + + char unique_key[64]; + int n = snprintf(unique_key, sizeof(unique_key), "%s%d", key, local_rank); + if (n < 0 || n >= (int)sizeof(unique_key)) { + close(sock_fd); + throw std::runtime_error("KittensBroker: Socket name too long"); + } + + size_t len = strnlen(unique_key, sizeof(addr.sun_path)); + if (len > (sizeof(addr.sun_path) - 1)) { + close(sock_fd); + throw std::runtime_error("KittensBroker: Socket name too long"); + } + strcpy(addr.sun_path, unique_key); + unlink(unique_key); + + if (bind(sock_fd, (struct sockaddr *)&addr, SUN_LEN(&addr)) < 0) { + close(sock_fd); + throw std::runtime_error("KittensBroker: Failed to bind socket"); + } + + return sock_fd; +} + +__host__ inline void send_fd( + int sock_fd, + int data_fd, + const char *dst_key, + int dst_local_rank, + int src_local_rank +) { + union { + struct cmsghdr cm; + char* control; + } control_un; + + size_t sizeof_control = CMSG_SPACE(sizeof(int)); + control_un.control = reinterpret_cast(malloc(sizeof_control)); + if (!control_un.control) { + close(sock_fd); + close(data_fd); + throw std::runtime_error("KittensBroker: Failed to allocate a control buffer"); + } + + struct msghdr msg {}; + msg.msg_control = control_un.control; + msg.msg_controllen = sizeof_control; + + struct cmsghdr *cmptr = CMSG_FIRSTHDR(&msg); + cmptr->cmsg_len = CMSG_LEN(sizeof(int)); + cmptr->cmsg_level = SOL_SOCKET; + cmptr->cmsg_type = SCM_RIGHTS; + memmove(CMSG_DATA(cmptr), &data_fd, sizeof(data_fd)); + + struct sockaddr_un addr {}; + addr.sun_family = AF_UNIX; + char dst_unique_key[64]; + int n = snprintf(dst_unique_key, sizeof(dst_unique_key), "%s%d", dst_key, dst_local_rank); + if (n < 0 || n >= (int)sizeof(dst_unique_key)) { + free(control_un.control); + close(sock_fd); + close(data_fd); + throw std::runtime_error("KittensBroker: dst path too long"); + } + strcpy(addr.sun_path, dst_unique_key); + msg.msg_name = (void *)&addr; + msg.msg_namelen = sizeof(struct sockaddr_un); + + int payload = src_local_rank; + struct iovec iov[1]; + iov[0].iov_base = &payload; + iov[0].iov_len = sizeof(payload); + msg.msg_iov = iov; + msg.msg_iovlen = 1; + + while (true) { + ssize_t sent = sendmsg(sock_fd, &msg, 0); + if (sent <= 0) { + if (errno == EINTR) continue; + close(sock_fd); + close(data_fd); + free(control_un.control); + throw std::runtime_error("KittensBroker: Failed to send FD over socket"); + } + break; + } + + free(control_un.control); +} + +__host__ inline void recv_fd(int sock_fd, int *data_fd, int *src_local_rank) { + union { + struct cmsghdr cm; + char* control; + } control_un; + + size_t sizeof_control = CMSG_SPACE(sizeof(int)); + control_un.control = reinterpret_cast(malloc(sizeof_control)); + if (!control_un.control) { + close(sock_fd); + throw std::runtime_error("KittensBroker: Failed to allocate a control buffer"); + } + + struct msghdr msg {}; + msg.msg_control = control_un.control; + msg.msg_controllen = sizeof_control; + + int payload = -1; + struct iovec iov[1]; + iov[0].iov_base = &payload; + iov[0].iov_len = sizeof(payload); + msg.msg_iov = iov; + msg.msg_iovlen = 1; + + while (true) { + ssize_t received = recvmsg(sock_fd, &msg, 0); + if (received < 0 && errno == EINTR) { + msg.msg_controllen = sizeof_control; + msg.msg_iovlen = 1; + continue; + } + if (received < static_cast(sizeof(*data_fd))) { + free(control_un.control); + close(sock_fd); + throw std::runtime_error("KittensBroker: Failed to receive data over socket"); + } + break; + } + + if (msg.msg_flags & MSG_CTRUNC) { + free(control_un.control); + close(sock_fd); + throw std::runtime_error("KittensBroker: Control data truncated"); + } + + struct cmsghdr *cmptr = CMSG_FIRSTHDR(&msg); + if (!cmptr || + cmptr->cmsg_len != CMSG_LEN(sizeof(int)) || + cmptr->cmsg_level != SOL_SOCKET || + cmptr->cmsg_type != SCM_RIGHTS) { + free(control_un.control); + close(sock_fd); + throw std::runtime_error("KittensBroker: Failed to receive data over socket"); + } + + memmove(data_fd, CMSG_DATA(cmptr), sizeof(*data_fd)); + free(control_un.control); + *src_local_rank = payload; +} + +__host__ inline void unlink_socket(const char *key, int local_rank) { + char unique_key[64]; + int n = snprintf(unique_key, sizeof(unique_key), "%s%d", key, local_rank); + if (n < 0 || n >= (int)sizeof(unique_key)) + throw std::runtime_error("KittensBroker: Socket name too long"); + unlink(unique_key); +} + +__host__ inline void close_socket(int sock_fd) { + close(sock_fd); +} + +} +} + + + + + + + + + + + + + + + +struct KittensBroker { + + static inline constexpr const char *SHM_KEY_ = "/kittens_broker_shm"; + static inline constexpr const char *SOCK_KEY_ = "/tmp/kittens_broker.sock"; + + int local_rank_; + int local_world_size_; + + void *shm_raw_; + volatile detail::broker::KittensVault *shm_; + int sock_; + + __host__ inline KittensBroker(int local_rank, int local_world_size) + : local_rank_(local_rank), + local_world_size_(local_world_size), + shm_raw_(nullptr), + shm_(nullptr), + sock_(-1) { + if (local_rank_ < 0) + throw std::runtime_error("KittensBroker: Local rank must be non-negative"); + if (local_rank_ >= local_world_size_) + throw std::runtime_error("KittensBroker: Local rank is greater than local world size"); + if (local_world_size_ > detail::broker::MAX_LOCAL_WORLD_SIZE) + throw std::runtime_error("KittensBroker: Local world size is greater than MAX_LOCAL_WORLD_SIZE"); + + if (local_rank_ == 0) { + shm_raw_ = detail::broker::create_shm(SHM_KEY_, sizeof(detail::broker::KittensVault)); + shm_ = reinterpret_cast(shm_raw_); + memset(shm_raw_, 0, sizeof(detail::broker::KittensVault)); + } else { + shm_raw_ = detail::broker::open_shm(SHM_KEY_, sizeof(detail::broker::KittensVault)); + shm_ = reinterpret_cast(shm_raw_); + } + detail::broker::init_sync(local_rank_, shm_); + detail::broker::sync(local_world_size_, shm_); + + if (local_rank_ ==0) + detail::broker::unlink_shm(SHM_KEY_); + detail::broker::sync(local_world_size_, shm_); + + sock_ = detail::broker::create_socket(SOCK_KEY_, local_rank_); + detail::broker::sync(local_world_size_, shm_); + } + + __host__ KittensBroker(const KittensBroker&) = delete; + __host__ KittensBroker& operator=(const KittensBroker&) = delete; + + __host__ inline KittensBroker(KittensBroker&& other) noexcept + : local_rank_(other.local_rank_), + local_world_size_(other.local_world_size_), + shm_raw_(other.shm_raw_), + shm_(other.shm_), + sock_(other.sock_) { + other.local_rank_ = -1; + other.local_world_size_ = -1; + other.shm_raw_ = nullptr; + other.shm_ = nullptr; + other.sock_ = -1; + } + + __host__ inline void destroy() { + if (shm_raw_) { + detail::broker::unmap_shm(shm_raw_, sizeof(detail::broker::KittensVault)); + shm_raw_ = nullptr; + shm_ = nullptr; + } + if (sock_ >= 0) { + detail::broker::unlink_socket(SOCK_KEY_, local_rank_); + detail::broker::close_socket(sock_); + sock_ = -1; + } + local_rank_ = -1; + local_world_size_ = -1; + } + + __host__ inline KittensBroker& operator=(KittensBroker&& other) noexcept { + if (this != &other) { + destroy(); + local_rank_ = other.local_rank_; + local_world_size_ = other.local_world_size_; + shm_raw_ = other.shm_raw_; + shm_ = other.shm_; + sock_ = other.sock_; + other.local_rank_ = -1; + other.local_world_size_ = -1; + other.shm_raw_ = nullptr; + other.shm_ = nullptr; + other.sock_ = -1; + } + return *this; + } + + __host__ inline ~KittensBroker() { + destroy(); + } + + __host__ inline void sync(int num_ranks = -1) { + if (num_ranks == -1) + num_ranks = local_world_size_; + else if (num_ranks < 0 || num_ranks > local_world_size_) + throw std::runtime_error("KittensBroker: Invalid number of ranks"); + + detail::broker::sync(num_ranks, shm_); + } + + __host__ inline void exchange_data(void *dst_, const void *src_, size_t size) { + if (size > detail::broker::VAULT_SIZE_PER_RANK) + throw std::runtime_error("KittensBroker: Size is greater than VAULT_SIZE_PER_RANK"); + + uint8_t *dst = reinterpret_cast(dst_); + const uint8_t *src = reinterpret_cast(src_); + + + sync(); + memcpy(const_cast(shm_->data) + local_rank_ * detail::broker::VAULT_SIZE_PER_RANK, src, size); + sync(); + + + for (int i = 0; i < local_world_size_; i++) + memcpy(dst + i * size, const_cast(shm_->data) + i * detail::broker::VAULT_SIZE_PER_RANK, size); + } + + __host__ inline void exchange_fds(int *dst, const int data_fd) { + if (dst == nullptr) + throw std::runtime_error("KittensBroker: dst is null"); + if (data_fd < 0) + throw std::runtime_error("KittensBroker: source fd is negative"); + + + for (int i = 0; i < local_world_size_; ++i) + dst[i] = -1; + + + sync(); + + if (local_rank_ == 0) { + + dst[0] = data_fd; + for (int i = 0; i < local_world_size_ - 1; i++) { + int received_fd; + int src_local_rank; + detail::broker::recv_fd(sock_, &received_fd, &src_local_rank); + if (received_fd < 0) + throw std::runtime_error("KittensBroker: Failed to receive FD over socket"); + if (src_local_rank == local_rank_) + throw std::runtime_error("KittensBroker: Invalid source rank"); + dst[src_local_rank] = received_fd; + } + for (int dst_local_rank = 1; dst_local_rank < local_world_size_; dst_local_rank++) { + for (int src_local_rank = 0; src_local_rank < local_world_size_; src_local_rank++) { + if (dst_local_rank == src_local_rank) + continue; + detail::broker::send_fd(sock_, dst[src_local_rank], SOCK_KEY_, dst_local_rank, src_local_rank); + } + } + close(dst[0]); + dst[0] = -1; + } else { + + detail::broker::send_fd(sock_, data_fd, SOCK_KEY_, 0, local_rank_); + close(data_fd); + for (int i = 0; i < local_world_size_ - 1; i++) { + int received_fd; + int src_local_rank; + detail::broker::recv_fd(sock_, &received_fd, &src_local_rank); + if (received_fd < 0) + throw std::runtime_error("KittensBroker: Failed to receive FD over socket"); + if (src_local_rank == local_rank_) + throw std::runtime_error("KittensBroker: Invalid source rank"); + dst[src_local_rank] = received_fd; + } + } + + + sync(); + } + + __host__ inline void broadcast_fd(int *dst, const int data_fd, const int src_local_rank) { + if (src_local_rank < 0 || src_local_rank >= local_world_size_) + throw std::runtime_error("KittensBroker: Invalid source rank"); + + + sync(); + + if (local_rank_ == src_local_rank) { + if (data_fd < 0) + throw std::runtime_error("KittensBroker: Source rank has invalid FD"); + for (int dst_local_rank = 0; dst_local_rank < local_world_size_; dst_local_rank++) { + if (dst_local_rank == src_local_rank) + continue; + detail::broker::send_fd(sock_, data_fd, SOCK_KEY_, dst_local_rank, src_local_rank); + } + close(data_fd); + } else { + if (!dst) + throw std::runtime_error("KittensBroker: Destination rank has invalid buffer"); + int _src_local_rank; + detail::broker::recv_fd(sock_, dst, &_src_local_rank); + if (*dst < 0) + throw std::runtime_error("KittensBroker: Failed to receive valid FD over socket"); + if (_src_local_rank != src_local_rank) + throw std::runtime_error("KittensBroker: Invalid source rank"); + } + + + sync(); + } +}; + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/club.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/club.cuh new file mode 100644 index 0000000..5493951 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/club.cuh @@ -0,0 +1,124 @@ +#pragma once + +#include +#include +#include +#include +#include + + + + + + + + + + + + + + + + + + + +class KittensClub { +public: + __host__ inline KittensClub(const int *device_ids, const int num_devices); + __host__ inline KittensClub(const int *device_ids, const cudaStream_t *streams, const int num_devices); + __host__ inline ~KittensClub(); + + + __host__ inline void execute(std::function task); + +private: + + bool stop; + std::vector task_available; + int n_task_done; + + + std::vector workers; + + + std::vector streams; + + + __host__ inline void worker(int worker_id, int device_id); + + + std::function current_task; + + + std::mutex mutex; + std::condition_variable cond_task_available; + std::condition_variable cond_task_done; +}; + +__host__ inline KittensClub::KittensClub(const int *device_ids, const int num_devices) : stop(false), n_task_done(0) { + for (size_t dev_idx = 0; dev_idx < num_devices; ++dev_idx) { + task_available.push_back(false); + streams.push_back(0); + workers.emplace_back([this, dev_idx, device_ids] { worker(dev_idx, device_ids[dev_idx]); }); + } +} + +__host__ inline KittensClub::KittensClub(const int *device_ids, const cudaStream_t *streams_in, const int num_devices) : stop(false), n_task_done(0) { + for (size_t dev_idx = 0; dev_idx < num_devices; ++dev_idx) { + task_available.push_back(false); + streams.push_back(streams_in[dev_idx]); + workers.emplace_back([this, dev_idx, device_ids] { worker(dev_idx, device_ids[dev_idx]); }); + } +} + +__host__ inline KittensClub::~KittensClub() { + { + std::lock_guard lock(mutex); + stop = true; + } + cond_task_available.notify_all(); + for (std::thread &worker : workers) { + worker.join(); + } +} + +__host__ inline void KittensClub::execute(std::function task) { + { + std::lock_guard lock(mutex); + current_task = task; + for (size_t i = 0; i < task_available.size(); ++i) + task_available[i] = true; + } + cond_task_available.notify_all(); + { + std::unique_lock lock(mutex); + cond_task_done.wait(lock, [this] { return n_task_done == workers.size(); }); + n_task_done = 0; + } +} + +__host__ inline void KittensClub::worker(int worker_id, int device_id) { + cudaSetDevice(device_id); + while (true) { + std::function task; + { + std::unique_lock lock(mutex); + cond_task_available.wait(lock, [this, worker_id] { return stop || task_available[worker_id]; }); + + if (stop) + return; + + task = current_task; + task_available[worker_id] = false; + } + task(worker_id, streams[worker_id]); + { + std::lock_guard lock(mutex); + ++n_task_done; + if (n_task_done == workers.size()) + cond_task_done.notify_one(); + } + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/parallel_tensor.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/parallel_tensor.cuh new file mode 100644 index 0000000..839a14c --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/parallel_tensor.cuh @@ -0,0 +1,336 @@ +#pragma once + +#include +#include +#include + +#include +#include +#include + +#include "../types/system/vmm.cuh" +#include "../types/system/ipc.cuh" +#include "broker.cuh" + +namespace kittens { +namespace py { + + + + + + +struct TKParallelTensor { + inline static std::map, KittensBroker> brokers_; + + at::Tensor data_; + std::vector shape_; + at::ScalarType dtype_; + + std::vector raw_ptrs_; + size_t allocated_size_; + + int local_rank_; + int local_world_size_; + + bool multicast_; + void *multicast_ptr_; + size_t multicast_allocated_size_; + + detail::ipc::flavor ipc_flavor_; + + __host__ inline TKParallelTensor( + const at::Tensor &tensor, + int local_rank, + int local_world_size, + bool multicast + ) : data_(tensor), + shape_(tensor.sizes().vec()), + dtype_(tensor.scalar_type()), + raw_ptrs_(local_world_size, nullptr), + allocated_size_(tensor.nbytes()), + local_rank_(local_rank), + local_world_size_(local_world_size), + multicast_(multicast), + multicast_ptr_(nullptr), + multicast_allocated_size_(0), + ipc_flavor_(detail::ipc::flavor::LEGACY) { + + TORCH_CHECK(tensor.is_cuda(), "Tensor must be on CUDA device"); + TORCH_CHECK(tensor.is_contiguous(), "Tensor must be contiguous"); + TORCH_CHECK(tensor.dim() <= 4, "Only tensors with dim <= 4 are supported for TKParallelTensor"); + TORCH_CHECK(tensor.device().index() == local_rank_, "Tensor device index must match local_rank"); + TORCH_CHECK(local_rank_ >= 0, "local_rank must be non-negative"); + TORCH_CHECK(local_rank_ < local_world_size_, "local_rank must be less than local_world_size"); + TORCH_CHECK(!multicast, "Multicast is not supported for pre-allocated tensors"); + + brokers_.try_emplace( + {local_rank_, local_world_size_}, + local_rank_, local_world_size_ + ); + + if (brokers_.size() > 1) + std::cerr << "WARNING: 2 KittensBroker instances created in the same process. This is not safe." << std::endl; + + c10::cuda::CUDAGuard device_guard(local_rank_); + exchange_ipc_handles(); + } + + __host__ inline TKParallelTensor( + const std::vector &shape, + const at::ScalarType dtype, + int local_rank, + int local_world_size, + bool multicast + ) : shape_(shape), + dtype_(dtype), + raw_ptrs_(local_world_size, nullptr), + allocated_size_(0), + local_rank_(local_rank), + local_world_size_(local_world_size), + multicast_(multicast), + multicast_ptr_(nullptr), + multicast_allocated_size_(0), + ipc_flavor_(detail::ipc::flavor::VMM) { + + TORCH_CHECK(local_rank_ >= 0, "local_rank must be non-negative"); + TORCH_CHECK(local_rank_ < local_world_size_, "local_rank must be less than local_world_size"); + + brokers_.try_emplace( + {local_rank_, local_world_size_}, + local_rank_, local_world_size_ + ); + + if (brokers_.size() > 1) + std::cerr << "WARNING: 2 KittensBroker instances created in the same process. This is not safe." << std::endl; + + c10::cuda::CUDAGuard device_guard(local_rank_); + create_shareable_cuda_tensor(); + exchange_ipc_handles(); + + if (multicast_) + initialize_multicast(); + } + + __host__ TKParallelTensor(const TKParallelTensor&) = delete; + __host__ TKParallelTensor& operator=(const TKParallelTensor&) = delete; + __host__ TKParallelTensor& operator=(TKParallelTensor&& other) = delete; + + __host__ inline TKParallelTensor(TKParallelTensor&& other) : + data_(std::move(other.data_)), + shape_(std::move(other.shape_)), + dtype_(std::move(other.dtype_)), + raw_ptrs_(std::move(other.raw_ptrs_)), + allocated_size_(other.allocated_size_), + local_rank_(other.local_rank_), + local_world_size_(other.local_world_size_), + multicast_(other.multicast_), + multicast_ptr_(other.multicast_ptr_), + multicast_allocated_size_(other.multicast_allocated_size_), + ipc_flavor_(other.ipc_flavor_) { + other.data_ = at::Tensor(); + other.shape_.clear(); + other.dtype_ = at::ScalarType::Undefined; + other.raw_ptrs_.clear(); + other.allocated_size_ = 0; + other.local_rank_ = -1; + other.local_world_size_ = -1; + other.multicast_ = false; + other.multicast_ptr_ = nullptr; + other.multicast_allocated_size_ = 0; + } + + __host__ inline ~TKParallelTensor() { + destroy(); + } + + __host__ inline at::Tensor data() const { + return data_; + } + + __host__ inline void create_shareable_cuda_tensor() { + c10::cuda::CUDAGuard device_guard(local_rank_); + + TORCH_CHECK(!shape_.empty(), "Shape must be non-empty"); + TORCH_CHECK(shape_.size() <= 4, "Shape must have at most 4 dimensions for TKParallelTensor"); + size_t size = c10::elementSize(dtype_); + for (auto dim : shape_) { + TORCH_CHECK(dim > 0, "Size dimensions must be positive"); + size *= static_cast(dim); + } + + void *raw_ptr; + detail::vmm::vm_alloc_map_set_access( + &raw_ptr, &allocated_size_, size, local_rank_, local_world_size_); + + + int local_rank = local_rank_; + size_t allocated_size = allocated_size_; + + auto deleter = [local_rank, raw_ptr, allocated_size](void* p) mutable { + if (!p) return; + c10::cuda::CUDAGuard device_guard(local_rank); + auto stream = c10::cuda::getCurrentCUDAStream().stream(); + CUDACHECK(cudaStreamSynchronize(stream)); + detail::vmm::vm_unmap(raw_ptr, allocated_size); + }; + + at::TensorOptions options = at::TensorOptions() + .dtype(dtype_) + .device(at::kCUDA, local_rank_); + + data_ = at::from_blob(raw_ptr, shape_, std::move(deleter), options); + } + + template + __host__ inline void exchange_ipc_handles() { + using handle_t = detail::ipc::handle; + + + detail::ipc::check_support(local_rank_); + void *raw_ptr = reinterpret_cast(data_.data_ptr()); + handle_t ipc_handle; + detail::ipc::export_handle(&ipc_handle, raw_ptr); + + + std::vector all_ipc_handles(local_world_size_); + if constexpr (IPC_FLAVOR == detail::ipc::flavor::LEGACY) { + brokers_.at({local_rank_, local_world_size_}).exchange_data( + reinterpret_cast(all_ipc_handles.data()), + reinterpret_cast(&ipc_handle), + sizeof(handle_t) + ); + } else if constexpr (IPC_FLAVOR == detail::ipc::flavor::VMM) { + brokers_.at({local_rank_, local_world_size_}).exchange_fds( + reinterpret_cast(all_ipc_handles.data()), + ipc_handle.handle_ + ); + } else { + throw std::runtime_error("Invalid IPC flavor"); + } + + + for (int i = 0; i < local_world_size_; i++) { + if (i == local_rank_) + raw_ptrs_[i] = raw_ptr; + else + detail::ipc::import_handle(&raw_ptrs_[i], all_ipc_handles[i], allocated_size_, local_world_size_); + } + } + + __host__ inline void initialize_multicast() { + using handle_t = detail::ipc::handle; + + detail::vmm::multicast_check(local_rank_); + detail::ipc::check_support(local_rank_); + detail::vmm::handle multicast_handle; + + if (local_rank_ == 0) { + + detail::vmm::multicast_create_handle( + &multicast_handle, + &multicast_allocated_size_, + allocated_size_, + local_world_size_ + ); + + + if (allocated_size_ != multicast_allocated_size_) + throw std::runtime_error("Multicast allocated size does not match memory allocated size"); + + + handle_t ipc_handle; + detail::ipc::export_handle(&ipc_handle, multicast_handle); + + + brokers_.at({local_rank_, local_world_size_}).broadcast_fd(nullptr, ipc_handle.handle_, 0); + } else { + + handle_t ipc_handle; + brokers_.at({local_rank_, local_world_size_}).broadcast_fd(&ipc_handle.handle_, -1, 0); + multicast_allocated_size_ = allocated_size_; + detail::ipc::import_handle(&multicast_handle, ipc_handle, multicast_allocated_size_, local_world_size_); + } + + + detail::vmm::multicast_bind_device(multicast_handle, local_rank_); + brokers_.at({local_rank_, local_world_size_}).sync(); + + + detail::vmm::handle memory_handle; + detail::vmm::vm_retrieve_handle(&memory_handle, raw_ptrs_[local_rank_]); + detail::vmm::multicast_bind_memory(multicast_handle, memory_handle, allocated_size_); + brokers_.at({local_rank_, local_world_size_}).sync(); + + + detail::vmm::vm_map(&multicast_ptr_, multicast_handle, multicast_allocated_size_); + detail::vmm::vm_set_access(multicast_ptr_, multicast_allocated_size_, local_world_size_); + + + detail::vmm::vm_free(multicast_handle); + detail::vmm::vm_free(memory_handle); + } + + __host__ inline void destroy() { + + if (multicast_ && multicast_ptr_) { + brokers_.at({local_rank_, local_world_size_}).sync(); + detail::vmm::handle multicast_handle; + detail::vmm::vm_retrieve_handle(&multicast_handle, multicast_ptr_); + detail::vmm::vm_unmap(multicast_ptr_, multicast_allocated_size_); + detail::vmm::multicast_unbind_device(multicast_handle, multicast_allocated_size_, local_rank_); + brokers_.at({local_rank_, local_world_size_}).sync(); + detail::vmm::vm_free(multicast_handle); + } + + + for (int i = 0; i < local_world_size_; i++) { + if (i != local_rank_ && i < raw_ptrs_.size()) { + if (ipc_flavor_ == detail::ipc::flavor::LEGACY) { + detail::ipc::free_handle(raw_ptrs_[i], allocated_size_); + } else if (ipc_flavor_ == detail::ipc::flavor::VMM) { + detail::ipc::free_handle(raw_ptrs_[i], allocated_size_); + } else { + throw std::runtime_error("Invalid IPC flavor"); + } + } + } + brokers_.at({local_rank_, local_world_size_}).sync(); + + + if (data_.defined()) + data_.reset(); + + + shape_.clear(); + dtype_ = at::ScalarType::Undefined; + raw_ptrs_.clear(); + allocated_size_ = 0; + local_rank_ = -1; + local_world_size_ = -1; + multicast_ = false; + multicast_ptr_ = nullptr; + multicast_allocated_size_ = 0; + } +}; + +} +} + +#define BIND_TK_PARALLEL_TENSOR(m) \ + pybind11::class_(m, "TKParallelTensor") \ + .def(pybind11::init(), \ + pybind11::arg("tensor"), \ + pybind11::arg("local_rank"), \ + pybind11::arg("local_world_size"), \ + pybind11::arg("multicast") = false) \ + .def(pybind11::init&, const at::ScalarType&, int, int, bool>(), \ + pybind11::arg("shape"), \ + pybind11::arg("dtype"), \ + pybind11::arg("local_rank"), \ + pybind11::arg("local_world_size"), \ + pybind11::arg("multicast") = false) \ + .def("data", &kittens::py::TKParallelTensor::data) \ + .def_readonly("data_", &kittens::py::TKParallelTensor::data_) \ + .def_readonly("local_rank_", &kittens::py::TKParallelTensor::local_rank_) \ + .def_readonly("local_world_size_", &kittens::py::TKParallelTensor::local_world_size_) diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/profiler.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/profiler.cuh new file mode 100644 index 0000000..6ecf712 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/profiler.cuh @@ -0,0 +1,148 @@ +#pragma once + +#include + +#ifdef KITTENS_SM10X + +#include +#include +#include + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +template +struct __align__(1024) TKProfiler { + static constexpr int TIMING_WIDTH = _TIMING_WIDTH; + int timings[TIMING_WIDTH]; + uint64_t start_time; + using timing_t = kittens::gl; + + + + + + __device__ inline void init() { + if (threadIdx.x < TIMING_WIDTH) + timings[threadIdx.x] = 0; + if (threadIdx.x == 0) { + uint64_t start_time_reg; + asm volatile("{mov.u64 %0, %%globaltimer;}" : "=l"(start_time_reg)); + uint32_t start_time_addr = static_cast(__cvta_generic_to_shared(&start_time)); + asm volatile("{st.shared.b64 [%0], %1;}" : : "r"(start_time_addr), "l"(start_time_reg)); + } + __syncthreads(); + } + + + + + + __device__ inline void record(int index) { + uint64_t current_time; + asm volatile("{mov.u64 %0, %%globaltimer;}" : "=l"(current_time)); + timings[index] = static_cast(current_time - start_time); + } + + + + + + + + __device__ inline void store_and_reset(const timing_t &timings_global, int i, int j) { + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(timings)); + uint64_t dst_ptr = reinterpret_cast(&timings_global[{i, j, 0}]); + asm volatile("{fence.proxy.async.shared::cta;}" ::: "memory"); + asm volatile("{cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;}" + :: "l"(dst_ptr), "r"(src_ptr), "n"(TIMING_WIDTH*sizeof(int)) : "memory"); + kittens::tma::store_commit_group(); + kittens::tma::store_async_read_wait(); + asm volatile("{st.bulk.weak [%0], %1, 0;}" :: "r"(src_ptr), "n"(TIMING_WIDTH*sizeof(int)) : "memory"); + } + + + + + + + + __host__ static inline timing_t allocate(int N, int M) { + const size_t timings_size = static_cast(N)*M*TIMING_WIDTH; + int *d_timings; + CUDACHECK(cudaMalloc(&d_timings, timings_size*sizeof(int))); + CUDACHECK(cudaMemset(d_timings, 0, timings_size*sizeof(int))); + return timing_t{d_timings, nullptr, static_cast(N), static_cast(M), nullptr}; + } + + + + + + __host__ static inline void deallocate(timing_t &timings_global) { + CUDACHECK(cudaFree(timings_global.raw_ptr)); + } + + + + + + + __host__ static inline void save(timing_t &timings_global, const std::string &filename) { + const int M = static_cast(timings_global.depth()); + const int N = static_cast(timings_global.rows()); + const size_t timings_size = static_cast(M)*N*TIMING_WIDTH; + + std::vector h_timings(timings_size); + CUDACHECK(cudaMemcpy(h_timings.data(), timings_global.raw_ptr, timings_size*sizeof(int), cudaMemcpyDeviceToHost)); + + std::ofstream outfile(filename); + if (outfile.is_open()) { + outfile << "# shape: (" << M << ", " << N << ", " << TIMING_WIDTH << ")\n"; + for (int i = 0; i < M; i++) { + for (int j = 0; j < N; j++) { + for (int k = 0; k < TIMING_WIDTH; k++) { + int idx = (i * N + j) * TIMING_WIDTH + k; + outfile << h_timings[idx]; + if (k < TIMING_WIDTH - 1) outfile << " "; + } + outfile << "\n"; + } + } + outfile.close(); + } + } +}; + +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/pyutils.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/pyutils.cuh new file mode 100644 index 0000000..4aad211 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/pyutils.cuh @@ -0,0 +1,94 @@ +#pragma once + +#include +#include + +namespace kittens { +namespace py { + +template struct from_object { + __host__ static T make(pybind11::object obj) { + return obj.cast(); + } +}; +template struct from_object { + __host__ static GL make(pybind11::object obj) { + + if (pybind11::hasattr(obj, "__class__") && + obj.attr("__class__").attr("__name__").cast() == "Tensor") { + + + if (!obj.attr("is_contiguous")().cast()) { + throw std::runtime_error("Tensor must be contiguous"); + } + if (obj.attr("device").attr("type").cast() == "cpu") { + throw std::runtime_error("Tensor must be on CUDA device"); + } + + + std::array shape = {1, 1, 1, 1}; + auto py_shape = obj.attr("shape").cast(); + size_t dims = py_shape.size(); + if (dims > 4) { + throw std::runtime_error("Expected Tensor.ndim <= 4"); + } + for (size_t i = 0; i < dims; ++i) { + shape[4 - dims + i] = pybind11::cast(py_shape[i]); + } + + + uint64_t data_ptr = obj.attr("data_ptr")().cast(); + + + return make_gl(data_ptr, shape[0], shape[1], shape[2], shape[3]); + } + throw std::runtime_error("Expected a torch.Tensor"); + } +}; + + +template struct first_arg; +template struct first_arg { using type = T; }; +template using first_arg_t = typename first_arg::type; + + +template +__host__ constexpr MT TGlobal::* member_cast(MT TBase::* p) { + return static_cast(p); +} + +template struct trait; +template using object = pybind11::object; +template struct trait { using member_type = MT; using type = T; }; +template concept has_dynamic_shared_memory = requires(T t) { { t.dynamic_shared_memory() } -> std::convertible_to; }; +template __host__ static void bind_kernel(auto m, auto name, MemberPtrs... member_ptrs) { + using TGlobal = first_arg_t; + ((void)member_cast(member_ptrs), ...); + m.def(name, [=](object... args, pybind11::kwargs kwargs) { + TGlobal __g__ {from_object::member_type>::make(args)...}; + cudaStream_t raw_stream = nullptr; + if (kwargs.contains("stream")) { + + uintptr_t stream_ptr = kwargs["stream"].attr("cuda_stream").cast(); + raw_stream = reinterpret_cast(stream_ptr); + } + if constexpr (has_dynamic_shared_memory) { + int __dynamic_shared_memory__ = (int)__g__.dynamic_shared_memory(); + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, __dynamic_shared_memory__); + kernel<<<__g__.grid(), __g__.block(), __dynamic_shared_memory__, raw_stream>>>(__g__); + } else { + kernel<<<__g__.grid(), __g__.block(), 0, raw_stream>>>(__g__); + } + }); +} +template __host__ static void bind_function(auto m, auto name, MemberPtrs... member_ptrs) { + m.def(name, [=](object... args) { + using TGlobal = first_arg_t; + ((void)member_cast(member_ptrs), ...); + TGlobal __g__ {from_object::member_type>::make(args)...}; + function(__g__); + }); +} + +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/torchutils.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/torchutils.cuh new file mode 100644 index 0000000..59763a5 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/torchutils.cuh @@ -0,0 +1,227 @@ +#pragma once + +#include +#include + +#include +#include "parallel_tensor.cuh" + +#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor") +#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") +#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) + +namespace kittens { +namespace py { + +template +concept has_min_blocks_per_sm = requires { std::integral_constant{}; }; + +template +__host__ consteval int min_blocks_per_sm() { + if constexpr(has_min_blocks_per_sm) + return Config::MIN_BLOCKS_PER_SM; + else + return 1; +} + +template +__global__ +__launch_bounds__(Config::NUM_THREADS, min_blocks_per_sm()) +void global_kernel(const __grid_constant__ Globals G) { + Kernel(G); +} + +template +__host__ static inline void tensor_check(const at::Tensor &t) { + TORCH_CHECK(t.is_cuda() || t.is_pinned(), "Tensor must be on CUDA device or pinned memory") + TORCH_CHECK(t.is_contiguous(), "Tensor must be contiguous") + TORCH_CHECK(t.dim() <= 4, "Expected Tensor.dim() <= 4"); + + if constexpr (!TypeCheck) { + return; + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Bool, "Tensor has invalid dtype (expected bool)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Char, "Tensor has invalid dtype (expected int8)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Byte, "Tensor has invalid dtype (expected uint8)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Short, "Tensor has invalid dtype (expected int16)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::UInt16, "Tensor has invalid dtype (expected uint16)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Int, "Tensor has invalid dtype (expected int32)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::UInt32, "Tensor has invalid dtype (expected uint32)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Long, "Tensor has invalid dtype (expected int64)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::UInt64, "Tensor has invalid dtype (expected uint64)"); +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Float8_e4m3fn, "Tensor has invalid dtype (expected fp8e4m3)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Float8_e5m2, "Tensor has invalid dtype (expected fp8e5m2)"); +#endif +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Float8_e8m0fnu || t.dtype() == at::ScalarType::Byte, "Tensor has invalid dtype (expected fp8e8m0)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Float4_e2m1fn_x2, "Tensor has invalid dtype (expected fp4_2)"); +#endif + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::BFloat16, "Tensor has invalid dtype (expected bfloat16)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Half, "Tensor has invalid dtype (expected float16)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Float, "Tensor has invalid dtype (expected float32)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Double, "Tensor has invalid dtype (expected float64)"); + } else { + TORCH_CHECK(false, "Unsupported dtype"); + } +} + +template +__host__ static inline void parallel_tensor_check(const TKParallelTensor& t) { + tensor_check(t.data_); + TORCH_CHECK(t.data_.sizes().vec() == t.shape_, "Shape mismatch between TKParallelTensor and the underlying tensor"); + TORCH_CHECK(t.data_.dtype() == t.dtype_, "Dtype mismatch between TKParallelTensor and the underlying tensor"); + TORCH_CHECK(t.raw_ptrs_.size() == PGL::num_devices, "Number of devices mismatch between PGL and TKParallelTensor"); + TORCH_CHECK(t.local_rank_ == t.data_.device().index(), "Current tensor device index mismatch within TKParallelTensor"); + TORCH_CHECK(t.local_world_size_ == PGL::num_devices, "Number of devices mismatch between PGL and TKParallelTensor"); + TORCH_CHECK(t.multicast_ == PGL::multicast, "Multicast mismatch between PGL and TKParallelTensor"); + TORCH_CHECK(t.raw_ptrs_[t.local_rank_] == reinterpret_cast(t.data_.data_ptr()), "Current tensor data pointer not found in TKParallelTensor's raw_ptrs_"); +} + +template +__host__ static inline GL tensor_to_gl(const at::Tensor &t) { + tensor_check(t); + + std::array shape = {1, 1, 1, 1}; + for (int i = 0; i < static_cast(t.dim()); ++i) + shape[4 - t.dim() + i] = static_cast(t.size(i)); + + uint64_t data_ptr = reinterpret_cast(t.data_ptr()); + + return ::kittens::make_gl(data_ptr, shape[0], shape[1], shape[2], shape[3]); +} + +template +__host__ static inline GL tensor_to_gl(const at::Tensor &t, int B, int D, int R, int C) { + tensor_check(t); + + return ::kittens::make_gl(reinterpret_cast(t.data_ptr()), B, D, R, C); +} + +template +__host__ static inline PGL parallel_tensor_to_pgl(TKParallelTensor &t) { + parallel_tensor_check(t); + + std::array shape = {1, 1, 1, 1}; + for (int i = 0; i < static_cast(t.data_.dim()); ++i) { + shape[4 - t.data_.dim() + i] = static_cast(t.data_.size(i)); + } + + if constexpr (PGL::multicast) + return ::kittens::make_pgl( + reinterpret_cast(t.multicast_ptr_), reinterpret_cast(t.raw_ptrs_.data()), shape[0], shape[1], shape[2], shape[3]); + else + return ::kittens::make_pgl( + reinterpret_cast(t.raw_ptrs_.data()), shape[0], shape[1], shape[2], shape[3]); +} + +template +__host__ static inline PGL parallel_tensor_to_pgl(TKParallelTensor &t, int B, int D, int R, int C) { + parallel_tensor_check(t); + + if constexpr (PGL::multicast) + return ::kittens::make_pgl( + reinterpret_cast(t.multicast_ptr_), reinterpret_cast(t.raw_ptrs_.data()), B, D, R, C); + else + return ::kittens::make_pgl( + reinterpret_cast(t.raw_ptrs_.data()), B, D, R, C); +} + +template +__host__ static inline GL make_fake_gl(const int batch, const int depth, const int rows, const int cols) { + return ::kittens::make_gl(reinterpret_cast(nullptr), batch, depth, rows, cols); +} + +__host__ static inline void _device_check(const at::Tensor& first, const at::Tensor& second) { + TORCH_CHECK(first.device() == second.device(), "All tensors must be on the same device"); +} + +template +__host__ static inline void device_check(const T1& first, const Ts&... rest) { + (_device_check(first, rest), ...); +} + +__host__ static inline void _parallel_tensor_check(const TKParallelTensor& first, const TKParallelTensor& second) { + TORCH_CHECK(first.local_rank_ == second.local_rank_, "All parallel tensors must have the same local_rank"); + TORCH_CHECK(first.local_world_size_ == second.local_world_size_, "All parallel tensors must have the same local_world_size"); +} + +template +__host__ static inline void parallel_tensor_check(const T1& first, const Ts&... rest) { + (_parallel_tensor_check(first, rest), ...); +} + +template +concept static_grid = requires { Config::NUM_BLOCKS; }; + +template +concept static_block = requires { Config::NUM_THREADS; }; + +template +concept static_dynamic_shared_memory = requires { Config::DYNAMIC_SHARED_MEMORY; }; + +template +concept has_pdl_config = requires { { Config::USE_PDL } -> std::convertible_to; }; +template +inline constexpr bool use_pdl = false; +template requires has_pdl_config +inline constexpr bool use_pdl = Config::USE_PDL; + +template +__host__ static inline void launch_kernel(const Globals &G) { + dim3 grid; + if constexpr (static_grid) + grid = dim3{Config::NUM_BLOCKS, 1, 1}; + else + grid = G.grid(); + + dim3 block; + if constexpr (static_block) + block = dim3{Config::NUM_THREADS, 1, 1}; + else + block = G.block(); + + int dynamic_shared_memory; + if constexpr (static_dynamic_shared_memory) + dynamic_shared_memory = static_cast(Config::DYNAMIC_SHARED_MEMORY); + else + dynamic_shared_memory = G.dynamic_shared_memory(); + + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + +#if defined(KITTENS_SM90) + static_assert(Config::CLUSTER_SIZE <= 8, "Cluster size must be less than or equal to 8 for Hopper"); +#elif defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert(Config::CLUSTER_SIZE <= 16, "Cluster size must be less than or equal to 16 for Blackwell"); + if constexpr (Config::CLUSTER_SIZE > 8) + CUDACHECK(cudaFuncSetAttribute(global_kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1)); +#endif + CUDACHECK(cudaFuncSetAttribute(global_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dynamic_shared_memory)); + + if constexpr (Config::CLUSTER_SIZE <= 1) { + LaunchConfig> launch_config(grid, block, dynamic_shared_memory, stream); + CUDACHECK(cudaLaunchKernelEx(launch_config, global_kernel, G)); + } else { + LaunchConfig> launch_config(grid, block, dynamic_shared_memory, stream, Config::CLUSTER_SIZE); + CUDACHECK(cudaLaunchKernelEx(launch_config, global_kernel, G)); + } +} + +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/global/cgl.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/global/cgl.cuh new file mode 100644 index 0000000..c285239 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/global/cgl.cuh @@ -0,0 +1,50 @@ + + + + + +#pragma once + +#include "gl.cuh" + +namespace kittens { + + + +namespace ducks { +namespace cgl { +struct identifier {}; +} +} + + + + + + +template +struct cgl { + using identifier = ducks::cgl::identifier; + using component = _GL; + using T = component::T; + using T2 = component::T2; + using dtype = component::dtype; + component real, imag; +}; + +namespace ducks { +namespace cgl { + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v; +} +} + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/global/gl.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/global/gl.cuh new file mode 100644 index 0000000..16ed80b --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/global/gl.cuh @@ -0,0 +1,250 @@ + + + + + +#pragma once + +#include "../../common/common.cuh" +#include "../shared/shared.cuh" +#include "util.cuh" +#if (defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120)) && !defined(KITTENS_NO_HOST) +#include "tma.cuh" +#endif + +namespace kittens { + + + +struct dim { + static constexpr int BATCH = 0; + static constexpr int DEPTH = 1; + static constexpr int ROW = 2; + static constexpr int COL = 3; +}; + + + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +namespace ducks { +namespace tma { +namespace descriptor { +struct identifier {}; +template concept all = requires { + typename T::identifier; +} && std::is_same_v; +} +} +} +namespace detail { +namespace tma { +template struct descriptor_copy_helper {}; +template struct descriptor_copy_helper<_T> { static constexpr int value = _T::axis; using T = _T::T; }; +template struct descriptor_copy_helper<_T> { static constexpr int value = 2; using T = _T; }; +template struct descriptor_copy_helper<_T> { static constexpr int value = -1; using T = _T; }; +template using descriptor_copy_helper_t = descriptor_copy_helper::T; +template static constexpr int descriptor_copy_helper_v = descriptor_copy_helper::value; +} +} +namespace tma { +template struct descriptor { + using identifier = ducks::tma::descriptor::identifier; + using T = detail::tma::descriptor_copy_helper_t<_T>; + static_assert(ducks::st::all || ducks::sv::all || ducks::tma::descriptor::all, "Must be a shared TK type to generate a TMA descriptor."); + static constexpr int axis = ( + ducks::tma::descriptor::all<_T> ? detail::tma::descriptor_copy_helper_v<_T> : + (_axis != -9999) ? _axis : detail::tma::descriptor_copy_helper_v<_T>); + static_assert((kittens::ducks::st::all && axis >= 0 && axis <= 2) || (kittens::ducks::sv::all && axis == -1), "Internal template error detected."); +}; +} +#endif + +namespace detail { +template +struct descriptor_dict { +#ifndef KITTENS_NO_HOST + __host__ descriptor_dict() {} + template __host__ descriptor_dict(T _, int b, int d, int r, int c) {} +#endif + __host__ __device__ descriptor_dict(const descriptor_dict &other) {} +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + template __device__ const CUtensorMap* get() const { + static_assert( + std::is_same_v && std::is_same_v, + "SKILL ISSUE: Requested a TMA descriptor for a type not initialized in the global layout." + ); + return nullptr; + } +#endif +}; + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template +struct descriptor_dict<_T, Args...> { + static_assert(ducks::sv::all<_T> || ducks::st::all<_T> || ducks::tma::descriptor::all<_T>, "Must be a shared TK type to generate a TMA descriptor."); + using DESC = kittens::tma::descriptor<_T>; + CUtensorMap tma_desc; + descriptor_dict other_descs; +#ifndef KITTENS_NO_HOST + __host__ descriptor_dict() {} + __host__ descriptor_dict(typename DESC::T::dtype *data, int b, int d, int r, int c): other_descs(data, b, d, r, c) { + kittens::detail::tma::create_tensor_map(&tma_desc, data, b, d, r, c); + } +#endif + __host__ __device__ inline descriptor_dict(const descriptor_dict &other) : + tma_desc(other.tma_desc), other_descs(other.other_descs) {} + template __device__ inline const CUtensorMap* get() const { + if constexpr (std::is_same_v && DESC::axis == axis) { return &tma_desc; } + else { return other_descs.template get(); } + } +}; +#endif +} + + + +namespace ducks { +namespace gl { +struct identifier {}; +} +} + +template +struct gl { + static_assert(!std::is_same_v<_T, char>, "Bare 'char' is not permitted because its signedness is ISA-dependent."); +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert(!std::is_same_v<_T, fp4e2m1>, "For FP4 types, you must use a packed type (i.e., fp4e2m1_2 or fp4e2m1_4)."); +#endif + using identifier = ducks::gl::identifier; + + using T = base_types::packing<_T>::unpacked_type; + using T2 = base_types::packing<_T>::packed_type; + using dtype = T; + + T* raw_ptr; + + static constexpr int __b__ = b, __d__ = d, __r__ = r, __c__ = c; + + ducks::gl::make_dim_t batch_internal; + ducks::gl::make_dim_t depth_internal; + ducks::gl::make_dim_t rows_internal; + ducks::gl::make_dim_t cols_internal; + +#ifdef KITTENS_NO_HOST + __device__ constexpr int batch() const { if constexpr (b > 0) return b; else return batch_internal; } + __device__ constexpr int depth() const { if constexpr (d > 0) return d; else return depth_internal; } + __device__ constexpr int rows() const { if constexpr (r > 0) return r; else return rows_internal; } + __device__ constexpr int cols() const { if constexpr (c > 0) return c; else return cols_internal; } +#else + template __device__ __host__ static constexpr std::enable_if_t<(B > 0), int> batch() { return B; } + template __device__ __host__ std::enable_if_t<(B == -1), int> batch() const { return batch_internal; } + template __device__ __host__ static constexpr std::enable_if_t<(D > 0), int> depth() { return D; } + template __device__ __host__ std::enable_if_t<(D == -1), int> depth() const { return depth_internal; } + template __device__ __host__ static constexpr std::enable_if_t<(R > 0), int> rows() { return R; } + template __device__ __host__ std::enable_if_t<(R == -1), int> rows() const { return rows_internal; } + template __device__ __host__ static constexpr std::enable_if_t<(C > 0), int> cols() { return C; } + template __device__ __host__ std::enable_if_t<(C == -1), int> cols() const { return cols_internal; } +#endif + __device__ __host__ inline size_t numel() const { return static_cast(batch()) * depth() * rows() * cols(); } + + detail::descriptor_dict tma_descs; + + __host__ __device__ gl() = delete; +#ifndef KITTENS_NO_HOST + __host__ inline gl(T *_data, + ducks::gl::make_arg_t _batch, + ducks::gl::make_arg_t _depth, + ducks::gl::make_arg_t _rows, + ducks::gl::make_arg_t _cols) : + raw_ptr(_data), batch_internal(_batch), depth_internal(_depth), rows_internal(_rows), cols_internal(_cols) { + tma_descs = detail::descriptor_dict(raw_ptr, batch_internal, depth_internal, rows_internal, cols_internal); + } +#endif + __host__ __device__ inline gl(const gl &other) : + raw_ptr(other.raw_ptr), batch_internal(other.batch_internal), depth_internal(other.depth_internal), rows_internal(other.rows_internal), cols_internal(other.cols_internal), tma_descs(other.tma_descs) {} +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + template __device__ inline const CUtensorMap* get_tma() const { + return tma_descs.template get(); + } + template __device__ inline void prefetch_tma() const { + const CUtensorMap *tma_desc = tma_descs.template get(); + asm volatile ("{prefetch.tensormap [%0];}" :: "l"(reinterpret_cast(tma_desc)) : "memory"); + } +#endif + __device__ inline T& operator[](const coord &idx) const { + return raw_ptr[(((size_t)idx.b*depth() + idx.d)*rows() + idx.r)*cols() + idx.c]; + } + template __device__ inline size_t shape() const { + static_assert(axis==0 || axis==1 || axis==2 || axis==3, "Axis must be 0, 1, 2, or 3."); + if constexpr (axis==0) { return size_t(batch()); } + else if constexpr (axis==1) { return size_t(depth()); } + else if constexpr (axis==2) { return size_t(rows()); } + else if constexpr (axis==3) { return size_t(cols()); } + } + template __device__ inline size_t stride() const { + static_assert(axis==0 || axis==1 || axis==2 || axis==3, "Axis must be 0, 1, 2, or 3."); + if constexpr (axis==0) { return (size_t)depth()*rows()*cols(); } + else if constexpr (axis==1) { return (size_t)rows()*cols(); } + else if constexpr (axis==2) { return (size_t)cols(); } + else if constexpr (axis==3) { return 1; } + } +}; + +template using gl3 = gl<_T, 1, d, r, c, TMA_Types...>; +template using gl2 = gl<_T, 1, 1, r, c, TMA_Types...>; +template using gl1 = gl<_T, 1, 1, 1, c, TMA_Types...>; + +namespace ducks { +namespace gl { + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v; +} +} + +#ifndef KITTENS_NO_HOST + + + + + + + +template __host__ auto make_unsafe_gl_arg(int param) { + if constexpr (N > 0) { return nullptr; } + else { return param; } +} +template __host__ inline GL make_gl(uint64_t data, int b, int d, int r, int c) { + if constexpr (safe) { + if(GL::__b__ > 0 && b != GL::__b__) { + throw std::runtime_error("Batch dimension mismatch. Expected: " + std::to_string(GL::__b__) + ", Got: " + std::to_string(b)); + } + if(GL::__d__ > 0 && d != GL::__d__) { + throw std::runtime_error("Depth dimension mismatch. Expected: " + std::to_string(GL::__d__) + ", Got: " + std::to_string(d)); + } + if(GL::__r__ > 0 && r != GL::__r__) { + throw std::runtime_error("Row dimension mismatch. Expected: " + std::to_string(GL::__r__) + ", Got: " + std::to_string(r)); + } + if(GL::__c__ > 0 && c != GL::__c__) { + throw std::runtime_error("Column dimension mismatch. Expected: " + std::to_string(GL::__c__) + ", Got: " + std::to_string(c)); + } + } + return GL( + reinterpret_cast(data), + make_unsafe_gl_arg(b), + make_unsafe_gl_arg(d), + make_unsafe_gl_arg(r), + make_unsafe_gl_arg(c) + ); +} + +#endif + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/global/global.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/global/global.cuh new file mode 100644 index 0000000..c7380ce --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/global/global.cuh @@ -0,0 +1,10 @@ + + + + + +#pragma once + +#include "util.cuh" +#include "gl.cuh" +#include "cgl.cuh" diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/global/tma.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/global/tma.cuh new file mode 100644 index 0000000..b58e428 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/global/tma.cuh @@ -0,0 +1,337 @@ +#pragma once + +#include "../../common/common.cuh" +#include "../shared/shared.cuh" +#include "util.cuh" + +namespace kittens { +namespace detail { +namespace tma { + + + + + + + + + + + + + + + +template +__host__ static inline void create_tensor_map( + CUtensorMap *tma_map, const typename ST::dtype *src, int batch, int depth, int rows, int cols +) { + using dtype = typename ST::dtype; + static_assert(axis==0 || axis==1 || axis==2, "axis must be 0, 1, or 2"); +#ifdef KITTENS_SM10X + static_assert(!(std::is_same_v && axis != 2), "Axes 0 and 1 are not yet supported for FP4 type"); +#endif + + constexpr uint32_t tma_dim = ST::swizzle ? 5 : 4; + void *global_addr = (void*)(src); + + constexpr CUtensorMapDataType tma_format = ( + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_BFLOAT16 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_FLOAT16 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_FLOAT32 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_INT32 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : +#ifdef KITTENS_SM10X + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : +#endif + CUtensorMapDataType(-1) + ); + constexpr CUtensorMapInterleave tma_interleave = CU_TENSOR_MAP_INTERLEAVE_NONE; + constexpr CUtensorMapL2promotion tma_l2Promotion = CU_TENSOR_MAP_L2_PROMOTION_NONE; + constexpr CUtensorMapFloatOOBfill tma_oobFill = CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE; + constexpr CUtensorMapSwizzle tma_swizzle = ST::swizzle ? ( + ST::swizzle_bytes == 32 ? CU_TENSOR_MAP_SWIZZLE_32B : + ST::swizzle_bytes == 64 ? CU_TENSOR_MAP_SWIZZLE_64B : + ST::swizzle_bytes == 128 ? CU_TENSOR_MAP_SWIZZLE_128B : + CU_TENSOR_MAP_SWIZZLE_NONE + ) : CU_TENSOR_MAP_SWIZZLE_NONE; + + uint64_t gmem_shape [5] = {0, 0, 0, 0, 0}; + uint64_t gmem_stride[4] = {0, 0, 0, 0}; + uint32_t smem_shape [5] = {0, 0, 0, 0, 0}; + uint32_t smem_stride[5] = {1, 1, 1, 1, 1}; + + constexpr uint64_t shared_tile_height = ST::rows; + constexpr uint64_t shared_tile_width = ST::cols; + + + constexpr int swizzle_elements = ST::swizzle_bytes / sizeof(dtype); + + if constexpr (ST::swizzle) { + if constexpr (axis == 2) { + gmem_shape[0] = swizzle_elements; + gmem_shape[1] = (uint64_t)rows; + gmem_shape[2] = (uint64_t)(cols+swizzle_elements-1) / swizzle_elements; + gmem_shape[3] = (uint64_t)depth; + gmem_shape[4] = (uint64_t)batch; + + gmem_stride[0] = (uint64_t)cols * sizeof(dtype); + gmem_stride[1] = ST::swizzle_bytes; + gmem_stride[2] = (uint64_t)rows * cols * sizeof(dtype); + gmem_stride[3] = (uint64_t)depth * rows * cols * sizeof(dtype); + } + else if constexpr (axis == 1) { + gmem_shape[0] = swizzle_elements; + gmem_shape[1] = (uint64_t)depth; + gmem_shape[2] = (uint64_t)(cols+swizzle_elements-1) / swizzle_elements; + gmem_shape[3] = (uint64_t)rows; + gmem_shape[4] = (uint64_t)batch; + + gmem_stride[0] = (uint64_t)rows * cols * sizeof(dtype); + gmem_stride[1] = ST::swizzle_bytes; + gmem_stride[2] = (uint64_t)cols * sizeof(dtype); + gmem_stride[3] = (uint64_t)depth * rows * cols * sizeof(dtype); + + } + else { + gmem_shape[0] = swizzle_elements; + gmem_shape[1] = (uint64_t)batch; + gmem_shape[2] = (uint64_t)(cols+swizzle_elements-1) / swizzle_elements; + gmem_shape[3] = (uint64_t)rows; + gmem_shape[4] = (uint64_t)depth; + + gmem_stride[0] = (uint64_t)depth * rows * cols * sizeof(dtype); + gmem_stride[1] = ST::swizzle_bytes; + gmem_stride[2] = (uint64_t)cols * sizeof(dtype); + gmem_stride[3] = (uint64_t)rows * cols * sizeof(dtype); + } + smem_shape[0] = swizzle_elements; + smem_shape[1] = shared_tile_height; + smem_shape[2] = shared_tile_width / swizzle_elements; + smem_shape[3] = 1; + smem_shape[4] = 1; + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + + gmem_shape[0] = (uint64_t)cols; + gmem_shape[1] = (uint64_t)rows; + gmem_shape[2] = (uint64_t)depth; + gmem_shape[3] = (uint64_t)batch; + + gmem_stride[0] = (uint64_t)cols * sizeof(dtype); + gmem_stride[1] = (uint64_t)rows * cols * sizeof(dtype); + gmem_stride[2] = (uint64_t)depth * rows * cols * sizeof(dtype); + + smem_shape[0] = shared_tile_width; + smem_shape[1] = shared_tile_height; + smem_shape[2] = 1; + smem_shape[3] = 1; + } + + + assert((reinterpret_cast(global_addr) & 0b1111) == 0); + + assert(gmem_stride[0] % 16 == 0); + assert(gmem_stride[1] % 16 == 0); + assert(gmem_stride[2] % 16 == 0); + assert(gmem_stride[3] % 16 == 0); + + assert(smem_shape[0] <= 256); + assert(smem_shape[1] <= 256); + assert(smem_shape[2] <= 256); + + assert((smem_shape[0]*sizeof(dtype)) % 16 == 0); + + assert(smem_stride[0] <= 8); + assert(smem_stride[1] <= 8); + assert(smem_stride[2] <= 8); + assert(smem_stride[3] <= 8); + assert(smem_stride[4] <= 8); + + assert(smem_stride[0] == 1); + + if constexpr (tma_interleave == CU_TENSOR_MAP_INTERLEAVE_NONE && tma_swizzle != CU_TENSOR_MAP_SWIZZLE_NONE) { + assert(smem_shape[0] * sizeof(dtype) <= ST::swizzle_bytes); + } + + const uint64_t *gmem_shape_ptr = &gmem_shape[0]; + const uint64_t *gmem_stride_ptr = &gmem_stride[0]; + const uint32_t *smem_shape_ptr = &smem_shape[0]; + const uint32_t *smem_stride_ptr = &smem_stride[0]; + + CUresult result = cuTensorMapEncodeTiled( + tma_map, + tma_format, + tma_dim, + global_addr, + gmem_shape_ptr, + gmem_stride_ptr, + smem_shape_ptr, + smem_stride_ptr, + tma_interleave, + tma_swizzle, + tma_l2Promotion, + tma_oobFill); + + const char *error_string; + CUresult res = cuGetErrorString(result, &error_string); + if (result != CUDA_SUCCESS) { + std::string error_msg = format_tma_error( + "tile", error_string, + batch, depth, rows, cols, + tma_map, tma_format, tma_dim, global_addr, + gmem_shape_ptr, gmem_stride_ptr, + smem_shape_ptr, smem_stride_ptr, + 5, 4, 5, 5, + tma_interleave, tma_swizzle, tma_l2Promotion, tma_oobFill, + "ST::rows: " + std::to_string(ST::rows) + "\n ST::cols: " + std::to_string(ST::cols) + ); + throw std::runtime_error(error_msg); + } +} + + + + + + + + + + + + +template +__host__ static inline CUtensorMap* allocate_and_create_tensor_map(const typename ST::dtype *src, int batch, int depth, int rows, int cols) { + CUtensorMap *tma_map_d; + cudaMalloc(&tma_map_d, sizeof(CUtensorMap)); + CUtensorMap tma_map_host; + create_tensor_map(&tma_map_host, src, batch, depth, rows, cols); + cudaMemcpy(tma_map_d, &tma_map_host, sizeof(CUtensorMap), cudaMemcpyHostToDevice); + return tma_map_d; +} + + + + + + + + + + + + + + + +template +__host__ static inline void create_tensor_map(CUtensorMap *tma_map, const typename SV::dtype *src, int batch, int depth, int rows, int cols) { + using dtype = typename SV::dtype; + static_assert(axis == -1, "for vector TMA, row axis must be -1 as it's unused"); + static_assert(SV::length <= 256 || (SV::length*sizeof(dtype)) % 128 == 0); + + + + constexpr uint32_t tma_dim = 4; + void *global_addr = (void*)(src); + + constexpr CUtensorMapDataType tma_format = ( + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_BFLOAT16 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_FLOAT16 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_FLOAT32 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_INT32 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : +#ifdef KITTENS_SM10X + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : +#endif + CUtensorMapDataType(-1) + ); + constexpr CUtensorMapInterleave tma_interleave = CU_TENSOR_MAP_INTERLEAVE_NONE; + constexpr CUtensorMapL2promotion tma_l2Promotion = CU_TENSOR_MAP_L2_PROMOTION_NONE; + constexpr CUtensorMapFloatOOBfill tma_oobFill = CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE; + constexpr CUtensorMapSwizzle swizzle = CU_TENSOR_MAP_SWIZZLE_NONE; + + constexpr uint64_t dim1 = sv_tma_dim1; + + + uint64_t gmem_shape [4] = {(uint64_t)cols, (uint64_t)rows, (uint64_t)depth, (uint64_t)batch}; + uint64_t gmem_stride[3] = {(uint64_t)cols*sizeof(dtype), (uint64_t)cols*rows*sizeof(dtype), (uint64_t)cols*rows*depth*sizeof(dtype)}; + uint32_t smem_shape [4] = {(uint32_t)dim1, 1, 1, 1}; + uint32_t smem_stride[4] = {1, 1, 1, 1}; + + + assert((reinterpret_cast(global_addr) & 0b1111) == 0); + + assert(smem_shape[0] <= 256); + + const uint64_t *gmem_shape_ptr = &gmem_shape[0]; + const uint64_t *gmem_stride_ptr = &gmem_stride[0]; + const uint32_t *smem_shape_ptr = &smem_shape[0]; + const uint32_t *smem_stride_ptr = &smem_stride[0]; + + CUresult result = cuTensorMapEncodeTiled( + tma_map, + tma_format, + tma_dim, + global_addr, + gmem_shape_ptr, + gmem_stride_ptr, + smem_shape_ptr, + smem_stride_ptr, + tma_interleave, + swizzle, + tma_l2Promotion, + tma_oobFill + ); + + const char *error_string; + CUresult res = cuGetErrorString(result, &error_string); + if (result != CUDA_SUCCESS) { + std::string error_msg = format_tma_error( + "vector", error_string, + batch, depth, rows, cols, + tma_map, tma_format, tma_dim, global_addr, + gmem_shape_ptr, gmem_stride_ptr, + smem_shape_ptr, smem_stride_ptr, + 4, 3, 4, 4, + tma_interleave, swizzle, tma_l2Promotion, tma_oobFill, + "SV::length: " + std::to_string(SV::length) + ); + throw std::runtime_error(error_msg); + } +}; + + + + + + + + + + + + + +template +__host__ static inline CUtensorMap* allocate_and_create_tensor_map(const typename SV::dtype *src, int batch, int depth, int rows, int cols) { + CUtensorMap *tma_map_d; + cudaMalloc(&tma_map_d, sizeof(CUtensorMap)); + CUtensorMap tma_map_host; + create_tensor_map(&tma_map_host, src, batch, depth, rows, cols); + cudaMemcpy(tma_map_d, &tma_map_host, sizeof(CUtensorMap), cudaMemcpyHostToDevice); + return tma_map_d; +} + +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/global/util.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/global/util.cuh new file mode 100644 index 0000000..7dd5236 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/global/util.cuh @@ -0,0 +1,180 @@ +#pragma once + +namespace kittens { +namespace ducks { +namespace gl { + +template concept cdim = (d > 0); +template concept rdim = (d == -1); +template struct compiled_dim { + static_assert(cdim<_v>, "Invalid compile-time dimension value"); + static constexpr size_t v = _v; + __host__ __device__ inline compiled_dim(const std::nullptr_t &_) {} + __host__ __device__ inline constexpr operator size_t() const { return v; } +}; +struct runtime_dim { + size_t v; + __host__ __device__ inline runtime_dim(const size_t &_v) : v(_v) {} + __host__ __device__ inline operator size_t() const { return v; } +}; +template using make_dim_t = std::conditional_t, runtime_dim, compiled_dim>; +template using make_arg_t = std::conditional_t, size_t, std::nullptr_t>; +} +} + +namespace detail { +template concept tile = ducks::st::all || ducks::rt::all || ducks::cst::all || ducks::crt::all; +template concept vec = ducks::sv::all || ducks::rv::all || ducks::csv::all || ducks::crv::all; + +#if (defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120)) && !defined(KITTENS_NO_HOST) +namespace tma { + +__host__ static inline std::string format_tma_error( + const char* error_type, + const char* error_string, + int batch, int depth, int rows, int cols, + CUtensorMap* tma_map, + CUtensorMapDataType tma_format, + uint32_t tma_dim, + void* global_addr, + const uint64_t* gmem_shape, + const uint64_t* gmem_stride, + const uint32_t* smem_shape, + const uint32_t* smem_stride, + size_t gmem_shape_size, + size_t gmem_stride_size, + size_t smem_shape_size, + size_t smem_stride_size, + CUtensorMapInterleave tma_interleave, + CUtensorMapSwizzle tma_swizzle, + CUtensorMapL2promotion tma_l2Promotion, + CUtensorMapFloatOOBfill tma_oobFill, + const std::string& extra_info = "" +) { + std::string msg; + msg += std::string("Error in TMA descriptor creation (") + + (error_type ? error_type : "Unknown") + "): " + + (error_string ? error_string : "Unknown CUDA error") + "\n"; + msg += "Parameters:\n"; + msg += " batch: " + std::to_string(batch) + "\n"; + msg += " depth: " + std::to_string(depth) + "\n"; + msg += " rows: " + std::to_string(rows) + "\n"; + msg += " cols: " + std::to_string(cols) + "\n"; + if (!extra_info.empty()) msg += " " + extra_info + "\n"; + msg += "cuTensorMapEncodeTiled arguments:\n"; + msg += " tma_map: " + std::to_string(reinterpret_cast(tma_map)) + "\n"; + msg += " tma_format: " + std::to_string(tma_format) + "\n"; + msg += " tma_dim: " + std::to_string(tma_dim) + "\n"; + msg += " global_addr: " + std::to_string(reinterpret_cast(global_addr)) + "\n"; + cudaPointerAttributes attributes; + cudaError_t err = cudaPointerGetAttributes(&attributes, global_addr); + msg += " global_addr memory type: "; + if (err == cudaSuccess) { + if (attributes.type == cudaMemoryTypeDevice) msg += "valid device memory\n"; + else if (attributes.type == cudaMemoryTypeHost) msg += "host memory (invalid for TMA)\n"; + else if (attributes.type == cudaMemoryTypeManaged) msg += "managed memory\n"; + else msg += "unknown memory type\n"; + } else { + msg += "unable to determine (error: " + std::string(cudaGetErrorString(err)) + ")\n"; + } + msg += " gmem_shape: " + std::to_string(reinterpret_cast(gmem_shape)) + " ["; + for (size_t i = 0; i < gmem_shape_size; ++i) { + msg += std::to_string(gmem_shape[i]); + if (i < gmem_shape_size - 1) msg += ", "; + } + msg += "]\n"; + msg += " gmem_stride: " + std::to_string(reinterpret_cast(gmem_stride)) + " ["; + for (size_t i = 0; i < gmem_stride_size; ++i) { + msg += std::to_string(gmem_stride[i]); + if (i < gmem_stride_size - 1) msg += ", "; + } + msg += "]\n"; + msg += " smem_shape: " + std::to_string(reinterpret_cast(smem_shape)) + " ["; + for (size_t i = 0; i < smem_shape_size; ++i) { + msg += std::to_string(smem_shape[i]); + if (i < smem_shape_size - 1) msg += ", "; + } + msg += "]\n"; + msg += " smem_stride: " + std::to_string(reinterpret_cast(smem_stride)) + " ["; + for (size_t i = 0; i < smem_stride_size; ++i) { + msg += std::to_string(smem_stride[i]); + if (i < smem_stride_size - 1) msg += ", "; + } + msg += "]\n"; + msg += " tma_interleave: " + std::to_string(tma_interleave) + "\n"; + msg += " tma_swizzle: " + std::to_string(tma_swizzle) + "\n"; + msg += " tma_l2Promotion: " + std::to_string(tma_l2Promotion) + "\n"; + msg += " tma_oobFill: " + std::to_string(tma_oobFill) + "\n"; + return msg; +} + +} +#endif +} + +namespace ducks { +namespace coord { +struct identifier {}; +} +} +template struct coord { + using identifier = ducks::coord::identifier; + using BASE = _T; + static_assert(std::is_same_v || detail::tile || detail::vec); + int b, d, r, c; + __device__ inline coord(int _b, int _d, int _r, int _c) : b(_b), d(_d), r(_r), c(_c) {} + __device__ inline coord( int _d, int _r, int _c) : b( 0), d(_d), r(_r), c(_c) {} + __device__ inline coord( int _r, int _c) : b( 0), d( 0), r(_r), c(_c) {} + __device__ inline coord( int _c) : b( 0), d( 0), r( 0), c(_c) {} + __device__ inline coord( ) : b( 0), d( 0), r( 0), c( 0) {} + template __device__ inline coord(const coord &other) : b(other.b), d(other.d), r(other.r), c(other.c) {} + __device__ inline coord(const int4 &other) : b(other.x), d(other.y), r(other.z), c(other.w) {} + __device__ inline operator int4() const { return int4(b, d, r, c); } + template __device__ inline coord unit_coord() const { + if constexpr (detail::tile) { + static_assert(row_axis != col_axis, "row and column axes must be different"); + static_assert(row_axis >= 0 && row_axis <= 3, "row axis must be between 0 and 3"); + static_assert(col_axis >= 0 && col_axis <= 3, "column axis must be between 0 and 3"); + static_assert(col_axis == 3, "for now, column axis must be 3"); + return coord( + row_axis == 0 ? b*BASE::rows : b, + row_axis == 1 ? d*BASE::rows : d, + row_axis == 2 ? r*BASE::rows : r, + c*BASE::cols + ); + } + else if constexpr (detail::vec) { + static_assert(row_axis == -1, "row axis must be be -1 for a vector coordinate to be converted to a unit coordinate"); + static_assert(col_axis >= 0 && col_axis <= 3, "column axis must be between 0 and 3"); + static_assert(col_axis == 3, "for now, column axis must be 3"); + return coord(b, d, r, c*BASE::length); + } + else { + return coord(*this); + } + } + template __device__ inline int dim() const { + static_assert(axis >= 0 && axis <= 3, "axis must be between 0 and 3"); + if constexpr (axis == 0) { return b; } + else if constexpr (axis == 1) { return d; } + else if constexpr (axis == 2) { return r; } + else { return c; } + } +}; +namespace ducks { +namespace coord { + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v; +template concept tile = all && (std::is_same_v || detail::tile); +template concept vec = all && (std::is_same_v || detail::vec); +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/crt.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/crt.cuh new file mode 100644 index 0000000..fdfa55f --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/crt.cuh @@ -0,0 +1,95 @@ + + + + + +#pragma once + +#include "rt.cuh" +#include "crv.cuh" + +namespace kittens { + +namespace ducks { +namespace crt { + + + + + + +struct identifier {}; + + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v && ducks::rt::all; + + + + + + +template +concept row_layout = all && std::is_same_v; + + + + + + + + +template +concept col_layout = all && std::is_same_v; +} +} + + + + + + + + + + + + + + +template +struct crt { + using identifier = ducks::crt::identifier; + using component = rt<_T, _rows, _cols, _layout>; + using layout = component::layout; + using T = component::T; + using T2 = component::T2; + using dtype = component::dtype; + + static constexpr int rows = component::rows; + static constexpr int cols = component::cols; + static constexpr int height = component::height; + static constexpr int width = component::width; + + + component real; + component imag; + + using row_vec = crv::row_vec_layout>; + using col_vec = crv::col_vec_layout>; +}; + +template using crt_fl = crt; +template using crt_bf = crt; +template using crt_hf = crt; + + + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/crv.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/crv.cuh new file mode 100644 index 0000000..e3432b0 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/crv.cuh @@ -0,0 +1,85 @@ + + + + + +#pragma once + +#include "../../common/common.cuh" +#include "rv.cuh" + +namespace kittens { + + + + +namespace ducks { + + + + + +namespace crv { + + + + + + +struct identifier {}; + + + + + + + +template +concept all = requires { + typename T::identifier; +} && std::is_same_v; + +template concept naive_layout = all && std::is_same_v; +template concept align_layout = all && std::is_same_v; +template concept ortho_layout = all && std::is_same_v; +template concept tile_layout = align_layout || ortho_layout; +} +} + + + + + + + + + + + + + + +template +struct crv { + using identifier = ducks::crv::identifier; + using component = rv<_T, _length, _layout>; + using layout = component::layout; + + using T = component::T; + using T2 = component::T2; + using dtype = component::dtype; + + static constexpr int length = component::length; + static constexpr int tiles = component::tiles; + + + component real; + component imag; +}; + + +template using crv_fl = crv; +template using crv_bf = crv; +template using crv_hf = crv; + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/register.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/register.cuh new file mode 100644 index 0000000..35e7b5f --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/register.cuh @@ -0,0 +1,14 @@ + + + + + +#pragma once + +#include "rv_layout.cuh" +#include "rt_base.cuh" +#include "rv.cuh" +#include "rt.cuh" + +#include "crv.cuh" +#include "crt.cuh" diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/rt.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/rt.cuh new file mode 100644 index 0000000..7ba92fe --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/rt.cuh @@ -0,0 +1,274 @@ + + + + + +#pragma once + +#include "../../common/common.cuh" + +#include "rt_layout.cuh" +#include "rt_base.cuh" +#include "rv.cuh" + +namespace kittens { + + + + +namespace ducks { + + + + + +namespace rt { + + + + + + +struct identifier {}; + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v; + + + + + + + + +template +concept row_layout = all && std::is_same_v; + + + + + + + + +template +concept col_layout = all && std::is_same_v; +} +} + + + + + + + + + + + + + + + + +template +struct rt { + static_assert(!std::is_same_v<_T, char>, "Bare 'char' is not permitted because its signedness is ISA-dependent."); + using identifier = ducks::rt::identifier; + using layout = _layout; + static_assert(kittens::ducks::base_types::T1<_T>); + using T = kittens::base_types::packing<_T>::unpacked_type; + using T2 = kittens::base_types::packing<_T>::packed_type; + using dtype = T2; + + static constexpr int rows = _rows; + static_assert(rows % rt_base::tile_size_row == 0, "Rows must be divisible by the tile size"); + static constexpr int cols = _cols; + static_assert(cols % rt_base::tile_size_col == 0, "Columns must be divisible by the tile size"); + static constexpr int height = rows / rt_base::tile_size_row; + static constexpr int width = cols / rt_base::tile_size_col; + static constexpr int tile_size_row = rt_base::tile_size_row; + static constexpr int tile_size_col = rt_base::tile_size_col; + static constexpr int num_elements = rt_base::num_elements * width * height; + static constexpr int elements_per_thread = rt_base::elements_per_thread * width * height; + static constexpr int packed_per_thread = rt_base::packed_per_thread * width * height; + static constexpr int packed_per_tile = rt_base::packed_per_thread; + + rt_base tiles[height][width]; + + using row_vec = rv::row_vec_layout>; + using col_vec = rv::col_vec_layout>; + + __device__ inline void operator=(const T &value) { + T2 value2 = base_types::packing::pack(value); + #pragma unroll + for(int i = 0; i < height; i++) { + #pragma unroll + for(int j = 0; j < width; j++) { + #pragma unroll + for(int k = 0; k < packed_per_tile; k++) { + tiles[i][j].data[k] = value2; + } + } + } + } + template + __device__ inline void operator=(const rt &other) { + using U2 = base_types::packing::packed_type; + #pragma unroll + for(int i = 0; i < height; i++) { + #pragma unroll + for(int j = 0; j < width; j++) { + #pragma unroll + for(int k = 0; k < packed_per_tile; k++) { + tiles[i][j].data[k] = base_types::convertor::convert(other.tiles[i][j].data[k]); + } + } + } + } +}; + + + + + + +template using rt_fl = rt; +template using rt_bf = rt; +template using rt_hf = rt; +template using rt_int8 = rt; +template using rt_uint8 = rt; +template using rt_int = rt; +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template using rt_fp8e4m3 = rt; +template using rt_fp8e5m2 = rt; +#endif +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template using rt_fp8e8m0 = rt; +template using rt_fp4e2m1_2 = rt; +#endif + + + + + + +template +__device__ constexpr const char* get_rt_type_name() { + if constexpr (std::is_same_v) { + return "rt_fl"; + } else if constexpr (std::is_same_v) { + return "rt_hf"; + } else if constexpr (std::is_same_v) { + return "rt_bf"; +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + return "rt_fp4_e2m1_2"; + } else if constexpr (std::is_same_v) { + return "rt_fp8_e8m0"; +#endif +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + return "rt_fp8_e4m3"; + } else if constexpr (std::is_same_v) { + return "rt_fp8_e5m2"; +#endif + } else { + return "rt_unknown"; + } +} + + + + + + + + + + +template +__device__ inline void print(const RT& tile) { + if (laneid() == 0) { + printf("Block %d, Warp %d: Register Tile %dx%d (Type: %s<%d,%d>) - Distributed View:\n", + blockIdx.x, threadIdx.x / WARP_THREADS, RT::rows, RT::cols, + get_rt_type_name(), RT::rows, RT::cols); + printf("Each thread holds %d elements (%d packed)\n", + RT::elements_per_thread, RT::packed_per_thread); + printf("\n"); + } + __syncwarp(); + + + for (int tid = 0; tid < WARP_THREADS; tid++) { + if (laneid() == tid) { + printf("Thread %2d: ", tid); + + + for (int i = 0; i < RT::height; i++) { + for (int j = 0; j < RT::width; j++) { + printf("Subtile[%d][%d]: ", i, j); + for (int k = 0; k < RT::packed_per_tile && k < 4; k++) { + auto packed_val = tile.tiles[i][j].data[k]; + + if constexpr (std::is_same_v) { + + if constexpr (std::is_same_v) { + printf("%.3f ", packed_val); + } else if constexpr (std::is_same_v) { + printf("%.3f ", __half2float(packed_val)); + } else if constexpr (std::is_same_v) { + printf("%.3f ", __bfloat162float(packed_val)); +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + printf("%.3f ", (float)packed_val); +#endif + } else { + printf("%.3f ", (float)packed_val); + } + } else { + + if constexpr (std::is_same_v) { + printf("[%.3f, %.3f] ", packed_val.x, packed_val.y); + } else if constexpr (std::is_same_v) { + + printf("[%.3f, %.3f] ", __bfloat162float(packed_val.x), __bfloat162float(packed_val.y)); +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + + __nv_fp8_e8m0 *vals = reinterpret_cast<__nv_fp8_e8m0*>(const_cast(&packed_val)); + printf("[%.3f,%.3f,%.3f,%.3f] ", + (float)vals[0], (float)vals[1], (float)vals[2], (float)vals[3]); + } else if constexpr (std::is_same_v) { + + uint8_t *vals = reinterpret_cast(const_cast(&packed_val)); + printf("[%.3f,%.3f,%.3f,%.3f] ", (float)fp4e2m1(vals[0] & 0xF), (float)fp4e2m1((vals[0] >> 4) & 0xF), (float)fp4e2m1(vals[1] & 0xF), (float)fp4e2m1((vals[1] >> 4) & 0xF)); +#endif + } else { + + printf("0x%x ", *(uint32_t*)&packed_val); + } + } + } + if (RT::packed_per_tile > 4) printf("... "); + } + } + printf("\n"); + } + __syncwarp(); + } + + if (laneid() == 0) { + printf("\n"); + } + __syncwarp(); +} + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/rt_base.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/rt_base.cuh new file mode 100644 index 0000000..4a36591 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/rt_base.cuh @@ -0,0 +1,124 @@ + + + + + +#pragma once + +#include "../../common/common.cuh" +#include "rt_layout.cuh" +#include "rv_layout.cuh" + +namespace kittens { + + + +namespace ducks { + + + + + +namespace rt_base { + + + + + + +struct identifier {}; +} +} + + + + + + + + + + + + +template struct rt_base { + using identifier = ducks::rt_base::identifier; + using layout = _layout; + static_assert(kittens::ducks::base_types::T1<_T>); + using T = kittens::base_types::packing<_T>::unpacked_type; + using T2 = kittens::base_types::packing<_T>::packed_type; + using dtype = T2; + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert( + std::is_same_v || std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + || std::is_same_v || std::is_same_v +#endif + , + "rt_base was provided an unsupported type." + ); +#else + static_assert( + std::is_same_v || std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v || std::is_same_v, + "rt_base was provided an unsupported type." + ); +#endif + + static constexpr int tile_size_row = kittens::TILE_ROW_DIM; + static constexpr int tile_size_col = kittens::TILE_COL_DIM; + static constexpr int rows = tile_size_row; + static constexpr int cols = tile_size_col; + static constexpr int num_elements = rows*cols; + static constexpr int elements_per_thread = num_elements / 32; + + static constexpr int packed_per_thread = (elements_per_thread / base_types::packing::num()) ; + static constexpr int registers_per_thread = packed_per_thread * sizeof(dtype) / 4; + + using row_vec_layout = std::conditional_t, ducks::rv_layout::align, ducks::rv_layout::ortho>; + using col_vec_layout = std::conditional_t, ducks::rv_layout::ortho, ducks::rv_layout::align>; + + dtype data[packed_per_thread]; +}; + + + + + + + +namespace ducks { +namespace rt_base { + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v; +} +} + + + +template using rt_base_fl = rt_base; +template using rt_base_bf = rt_base; +template using rt_base_hf = rt_base; +template using rt_base_int8 = rt_base; +template using rt_base_uint8 = rt_base; +template using rt_base_int = rt_base; +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template using rt_base_fp8e4m3 = rt_base; +template using rt_base_fp8e5m2 = rt_base; +#endif +#ifdef KITTENS_SM10X +template using rt_base_fp8e8m0 = rt_base; +template using rt_base_fp4e2m1_2 = rt_base; +#endif +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/rt_layout.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/rt_layout.cuh new file mode 100644 index 0000000..251b09a --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/rt_layout.cuh @@ -0,0 +1,40 @@ + + + + + +#pragma once + +namespace kittens { +namespace ducks { + + + + + +namespace rt_layout { + + + + +struct row {}; + + + +struct col {}; + + + + +template +concept all = std::is_same_v || std::is_same_v; + + + + +template struct transpose { using type = col; }; +template<> struct transpose { using type = row; }; + +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/rv.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/rv.cuh new file mode 100644 index 0000000..c384d7b --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/rv.cuh @@ -0,0 +1,230 @@ + + + + + +#pragma once + +#include "../../common/common.cuh" +#include "rv_layout.cuh" + +namespace kittens { + + + + +namespace ducks { + + + + + +namespace rv { + + + + + + +struct identifier {}; + + + + + + + +template +concept all = requires { + typename T::identifier; +} && std::is_same_v; + +template concept naive_layout = all && std::is_same_v; +template concept align_layout = all && std::is_same_v; +template concept ortho_layout = all && std::is_same_v; +template concept tile_layout = align_layout || ortho_layout; +} +} + + + + + + + + + + + + + +template +struct rv { + static_assert(!std::is_same_v<_T, char>, "Bare 'char' is not permitted because its signedness is ISA-dependent."); + using identifier = ducks::rv::identifier; + static_assert(kittens::ducks::base_types::T1<_T>); + using layout = _layout; + static constexpr bool is_naive = std::is_same_v; + using T = kittens::base_types::packing<_T>::unpacked_type; + using T2 = kittens::base_types::packing<_T>::packed_type; + using dtype = std::conditional_t; + + static constexpr int length = _length; + static_assert(length % kittens::TILE_ROW_DIM == 0, "Length must be divisible by the tile dimension"); + static constexpr int tiles = _length / kittens::TILE_ROW_DIM; + static constexpr int inner_dim = layout::inner_dim; + static constexpr int outer_dim = is_naive ? (tiles+1)/2 : tiles; + #if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert(!std::is_same_v && !std::is_same_v, "Unsupported type for fp8"); + #endif + + dtype data[outer_dim][inner_dim]; + + __device__ inline dtype* operator[](size_t idx) { return &data[idx][0]; } + __device__ inline const dtype* operator[](size_t idx) const { return &data[idx][0]; } + __device__ inline dtype& operator[](int2 outin) { return data[outin.x][outin.y]; } + __device__ inline const dtype& operator[](int2 outin) const { return data[outin.x][outin.y]; } + + __device__ inline void operator=(const T &value) { + dtype value2; + if constexpr(is_naive) { + value2 = value; + } else { + value2 = base_types::packing::pack(value); + } + #pragma unroll + for(int i = 0; i < outer_dim; i++) { + #pragma unroll + for(int j = 0; j < inner_dim; j++) { + data[i][j] = value2; + } + } + } + template + __device__ inline void operator=(const rv &other) { + using U2 = base_types::packing::packed_type; + #pragma unroll + for(int i = 0; i < outer_dim; i++) { + #pragma unroll + for(int j = 0; j < inner_dim; j++) { + data[i][j] = base_types::convertor::convert(other.data[i][j]); + } + } + } +}; + +template using rv_fl = rv; +template using rv_bf = rv; +template using rv_hf = rv; + + + + + + + + + + + +template +__device__ void print(const RV &vec) { + if (laneid() == 0) { + printf("Block %d, Warp %d: Register Vector %d (Type: %s, Layout: %s) - Distributed View:\n", + blockIdx.x, threadIdx.x / WARP_THREADS, RV::length, + std::is_same_v ? "float" : + std::is_same_v ? "bf16" : + std::is_same_v ? "half" : +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + std::is_same_v ? "fp8e8m0" : +#endif +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + std::is_same_v ? "fp8e4m3" : + std::is_same_v ? "fp8e5m2" : "unknown", +#endif + RV::is_naive ? "naive" : "tile"); + printf("Each thread holds %dx%d elements\n", RV::outer_dim, RV::inner_dim); + printf("\n"); + } + __syncwarp(); + + + for (int tid = 0; tid < WARP_THREADS; tid++) { + if (laneid() == tid) { + printf("Thread %2d: ", tid); + + + for (int i = 0; i < RV::outer_dim; i++) { + printf("Outer[%d]: ", i); + for (int j = 0; j < RV::inner_dim; j++) { + auto value = vec.data[i][j]; + + if constexpr (std::is_same_v) { + + if constexpr (std::is_same_v) { + printf("%.3f ", value); + } else if constexpr (std::is_same_v) { + printf("%.3f ", __half2float(value)); + } else if constexpr (std::is_same_v) { + printf("%.3f ", __bfloat162float(value)); +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + printf("%.3f ", (float)value); +#endif +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + printf("%.3f ", (float)value); + } else if constexpr (std::is_same_v) { + printf("%.3f ", (float)value); +#endif + } else { + printf("%.3f ", (float)value); + } + } else { + + if constexpr (std::is_same_v) { + printf("[%.3f, %.3f] ", value.x, value.y); + } else if constexpr (std::is_same_v) { + + printf("[%.3f, %.3f] ", __bfloat162float(value.x), __bfloat162float(value.y)); + } else if constexpr (std::is_same_v) { + + printf("[%.3f, %.3f] ", __half2float(value.x), __half2float(value.y)); +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + + __nv_fp8_e8m0 *vals = reinterpret_cast<__nv_fp8_e8m0*>(const_cast(&value)); + printf("[%.3f,%.3f,%.3f,%.3f] ", + (float)vals[0], (float)vals[1], (float)vals[2], (float)vals[3]); +#endif +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + + __nv_fp8_e4m3 *vals = reinterpret_cast<__nv_fp8_e4m3*>(const_cast(&value)); + printf("[%.3f,%.3f,%.3f,%.3f] ", + (float)vals[0], (float)vals[1], (float)vals[2], (float)vals[3]); + } else if constexpr (std::is_same_v) { + + __nv_fp8_e5m2 *vals = reinterpret_cast<__nv_fp8_e5m2*>(const_cast(&value)); + printf("[%.3f,%.3f,%.3f,%.3f] ", + (float)vals[0], (float)vals[1], (float)vals[2], (float)vals[3]); +#endif + } else { + + printf("0x%x ", *(uint32_t*)&value); + } + } + } + printf(" "); + } + printf("\n"); + } + __syncwarp(); + } + + if (laneid() == 0) { + printf("\n"); + } +} + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/rv_layout.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/rv_layout.cuh new file mode 100644 index 0000000..75ac192 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/rv_layout.cuh @@ -0,0 +1,38 @@ + + + + + +#pragma once + +namespace kittens { +namespace ducks { + + + + + +namespace rv_layout { + + + + +struct align { constexpr static int inner_dim = 2; }; + + + +struct ortho { constexpr static int inner_dim = 1; }; + + + +struct naive { constexpr static int inner_dim = 1; }; + + + + +template +concept all = std::is_same_v || std::is_same_v || std::is_same_v; + +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/cst.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/cst.cuh new file mode 100644 index 0000000..b4f867f --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/cst.cuh @@ -0,0 +1,78 @@ + + + + + +#pragma once + +#include "st.cuh" +#include "csv.cuh" + +namespace kittens { + +namespace ducks { +namespace cst { + + + + + + +struct identifier {}; + + + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v && ducks::st::all; + +} +} + + + + + + + + + + + + + +template +struct cst { + using identifier = ducks::cst::identifier; + using component = st<_T, _rows, _cols>; + using T = component::T; + using T2 = component::T2; + using dtype = component::dtype; + + static constexpr int rows = component::rows; + static constexpr int cols = component::cols; + + + + + component real; + component imag; + + + using col_vec = csv; + using row_vec = csv; +}; + + + +template using cst_bf = cst; +template using cst_hf = cst; +template using cst_fl = cst; + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/csv.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/csv.cuh new file mode 100644 index 0000000..a961774 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/csv.cuh @@ -0,0 +1,72 @@ + + + + + +#pragma once + +#include "sv.cuh" + +namespace kittens { + +namespace ducks { +namespace csv { + + + + + + +struct identifier {}; + + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v && ducks::sv::all; + +} +} + + + + + + + + + + + + +template +struct csv { + using identifier = ducks::csv::identifier; + using component = sv<_T, _length>; + using T = component::T; + using T2 = component::T2; + using dtype = component::dtype; + + static constexpr int length = component::length; + static constexpr int tiles = component::tiles; + + + + + component real; + component imag; +}; + + + + +template using csv_bf = csv; +template using csv_hf = csv; +template using csv_fl = csv; + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/descriptor.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/descriptor.cuh new file mode 100644 index 0000000..d5b955f --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/descriptor.cuh @@ -0,0 +1,142 @@ + + + + + +#pragma once + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + +#include "../../common/common.cuh" +#include "st.cuh" +#include "cst.cuh" + +namespace kittens { +namespace ducks { +namespace st_descriptor { +struct identifier {}; +} +} + +namespace detail { + + +__device__ static inline uint64_t matrix_descriptor_encode(uint64_t x) { return (((x) & 0x3FFFF) >> 0x4); } + +template +__device__ static inline uint64_t matrix_descriptor_raw( + T *addr, + uint32_t leading_dim_offset, + uint32_t stride_dim_offset, + uint32_t swizzle_mode +) { +#ifdef KITTENS_SM10X + + return matrix_descriptor_encode(reinterpret_cast(addr)) | + (1llu << 46) | +#else + return matrix_descriptor_encode(reinterpret_cast(addr)) | +#endif + matrix_descriptor_encode((uint64_t)leading_dim_offset) << 16 | + matrix_descriptor_encode((uint64_t)stride_dim_offset) << 32 | + (uint64_t)swizzle_mode << 62; +} + +} + +template +struct st_descriptor { + using identifier = ducks::st_descriptor::identifier; + using ST = _ST; + using T = typename ST::T; + static constexpr int rows = ST::rows; + static constexpr int cols = ST::cols; + static constexpr bool swizzle = ST::swizzle; + static_assert(swizzle, "Non-swizzled descriptor is not supported yet."); + uint64_t base_desc; + __device__ inline st_descriptor(const ST &tile) { + + if constexpr (MN_major) { + if constexpr (ST::swizzle_bytes == 128) + base_desc = detail::matrix_descriptor_raw(&tile.data[0], 2048*ST::rows/TILE_ROW_DIM, 1024, 1); + else if constexpr (ST::swizzle_bytes == 64) + base_desc = detail::matrix_descriptor_raw(&tile.data[0], 1024*ST::rows/TILE_ROW_DIM, 512, 2); + else + base_desc = detail::matrix_descriptor_raw(&tile.data[0], 512*ST::rows/TILE_ROW_DIM, 256, 3); + } + else { + if constexpr (ST::swizzle_bytes == 128) + base_desc = detail::matrix_descriptor_raw(&tile.data[0], 16 , 1024, 1); + else if constexpr (ST::swizzle_bytes == 64) + base_desc = detail::matrix_descriptor_raw(&tile.data[0], 16 , 512, 2); + else + base_desc = detail::matrix_descriptor_raw(&tile.data[0], 16 , 256, 3); + } + } + __device__ inline st_descriptor(const st_descriptor &other) : base_desc(other.base_desc) {} + __device__ inline uint64_t chunk_descriptor(int chunk_idx) { + + + + + + + if constexpr (MN_major) { + if constexpr (ST::swizzle_bytes == 128) { + return base_desc + detail::matrix_descriptor_encode(chunk_idx*2048); + } + else if constexpr (ST::swizzle_bytes == 64) { + return base_desc + detail::matrix_descriptor_encode(chunk_idx*1024); + } + else { + return base_desc + detail::matrix_descriptor_encode(chunk_idx*512); + } + } + else { + if constexpr (ST::swizzle_bytes == 128) { + + return base_desc + detail::matrix_descriptor_encode((chunk_idx%4)*32 + (chunk_idx/4)*(ST::rows/TILE_ROW_DIM)*2048); + } + else if constexpr (ST::swizzle_bytes == 64) { + + return base_desc + detail::matrix_descriptor_encode((chunk_idx%2)*32 + (chunk_idx/2)*(ST::rows/TILE_ROW_DIM)*1024); + } + else { + + return base_desc + detail::matrix_descriptor_encode(chunk_idx*(ST::rows/TILE_ROW_DIM)*512); + } + } + } +#ifdef KITTENS_SM103 + __device__ inline uint64_t chunk_descriptor_k96(int chunk_idx) { + static_assert(!MN_major, "K96 chunk descriptors are only supported for K-major tcgen05 operands."); + static_assert(std::is_same_v, "K96 chunk descriptors are only supported for packed FP4."); + static_assert(ST::swizzle_bytes == 128, "K96 chunk descriptors require 128B swizzle mode."); + const int start_byte = chunk_idx*48; + const int edge_gap = 128 - start_byte%128; + const int cont_byte = start_byte + (edge_gap < 48 ? edge_gap : 32); + + const uint64_t start_desc = base_desc + detail::matrix_descriptor_encode(start_byte%128 + (start_byte/128)*(ST::rows/TILE_ROW_DIM)*2048); + const uint64_t cont_desc = base_desc + detail::matrix_descriptor_encode(cont_byte%128 + (cont_byte/128)*(ST::rows/TILE_ROW_DIM)*2048); + return (start_desc & ~(0x3FFFull << 16)) | ((cont_desc & 0x3FFFull) << 16) | (1ull << 52); + } +#endif +}; + +namespace ducks { +namespace st_descriptor { + +template concept input = ducks::st::all || (requires {typename T::identifier;} && std::is_same_v); +template concept complex_input = ducks::cst::all; +namespace detail { +template struct st_getter { using type = typename T::ST; }; +template struct st_getter { using type = T; }; +template struct st_getter { using type = T::component; }; +template using get_st = typename st_getter::type; +} +} +} + +} + +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/shared.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/shared.cuh new file mode 100644 index 0000000..3755973 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/shared.cuh @@ -0,0 +1,14 @@ + + + + + +#pragma once + +#include "sv.cuh" +#include "st.cuh" + +#include "csv.cuh" +#include "cst.cuh" + +#include "descriptor.cuh" diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/st.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/st.cuh new file mode 100644 index 0000000..438299e --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/st.cuh @@ -0,0 +1,658 @@ + + + + + +#pragma once + +#include "../../common/common.cuh" +#include "sv.cuh" + + + + +namespace kittens { +namespace ducks { + + + + + +namespace st { + + + + + + + +struct identifier {}; + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v; +} +} + + +template< + typename ST, + int _subtile_height, + int _subtile_width +> +struct st_subtile; + + + + + + + + +template +struct KITTENS_DEFAULT_ALIGN st { + static_assert(!std::is_same_v<_T, char>, "Bare 'char' is not permitted because its signedness is ISA-dependent."); +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert(!std::is_same_v<_T, fp4e2m1>, "For FP4 types, you must use a packed type (i.e., fp4e2m1_2 or fp4e2m1_4)."); +#endif + using identifier = ducks::st::identifier; + using T = base_types::packing<_T>::unpacked_type; + using T2 = base_types::packing<_T>::packed_type; + using dtype = T; + + static constexpr bool swizzle = _swizzle; + + + static constexpr int underlying_rows = _rows; + static constexpr int underlying_cols = _cols; + static constexpr int underlying_num_elements = underlying_rows * underlying_cols; + + static constexpr int rows = _rows; + static constexpr int cols = _cols; + static constexpr int num_elements = rows * cols; + + static_assert(!swizzle || (rows % kittens::TILE_ROW_DIM == 0), "Rows must be divisible by the tile dimension"); + static_assert((swizzle && (cols % kittens::TILE_COL_DIM == 0)) || (!swizzle && (cols % kittens::BASE_TILE_DIM == 0)), "Cols must be divisible by the tile dimension"); + +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + static_assert(base_types::packing::num() == 1 || std::is_same_v); +#else + static_assert(base_types::packing::num() == 1); +#endif + + + static_assert(_swizzle_bytes == 0 || _swizzle_bytes == 32 || _swizzle_bytes == 64 || _swizzle_bytes == 128); + static constexpr int swizzle_bytes = _swizzle_bytes > 0 ? _swizzle_bytes : ( + sizeof(dtype) == 1 ? ( + (cols/kittens::TILE_COL_DIM)%4 == 0 ? 128 : + (cols/kittens::TILE_COL_DIM)%2 == 0 ? 64 : 32 + ) : + sizeof(dtype) == 2 ? ( + (cols/kittens::TILE_COL_DIM)%4 == 0 ? 128 : + (cols/kittens::TILE_COL_DIM)%2 == 0 ? 64 : 32 + ) : + sizeof(dtype) == 4 ? ( + (cols/kittens::TILE_COL_DIM)%2 == 0 ? 128 : 64 + ) : -1 + ); + + dtype data[rows*cols]; + + __device__ static inline T* idx(T *ptr, int2 coord) { + int r = coord.x, c = coord.y; + if constexpr (swizzle) { + static constexpr int swizzle_repeat = swizzle_bytes * 8; + static constexpr int subtile_cols = swizzle_bytes / sizeof(T); + const int outer_idx = c/subtile_cols; + const uint64_t addr = (uint64_t)(&ptr[outer_idx*rows*subtile_cols + r*subtile_cols + c%subtile_cols]); + const int swizzle = ((addr % swizzle_repeat) >> 7) << 4; + return (T*)(addr ^ swizzle); + } else { + return &ptr[r*cols + c]; + } + } + __device__ static inline uint32_t idx(uint32_t ptr, int2 coord) { + int r = coord.x, c = coord.y; + if constexpr (swizzle) { + static constexpr int swizzle_repeat = swizzle_bytes * 8; + static constexpr int subtile_cols = swizzle_bytes / sizeof(T); + const int outer_idx = c/subtile_cols; + const uint32_t addr = ptr + sizeof(T)*(outer_idx*rows*subtile_cols + r*subtile_cols + c%subtile_cols); + const int swizzle = ((addr % swizzle_repeat) >> 7) << 4; + return (addr ^ swizzle); + } else { + return ptr + sizeof(T)*(r*cols + c); + } + } + + + + + + + __device__ inline dtype& operator[](const int2 &rowcol) { + return *idx(data, rowcol); + } + __device__ inline const dtype& operator[](const int2 &rowcol) const { + return *(const dtype*)idx((dtype*)data, rowcol); + } + __device__ inline dtype& operator[](int idx) { + return data[idx]; + } + __device__ inline const dtype& operator[](int idx) const { + return data[idx]; + } + + template + __device__ inline st_subtile, subtile_rows, subtile_cols> subtile(int2 rowcol); + + + + + + template + __device__ inline st<_T, _rows, subtile_cols, _swizzle, swizzle_bytes > &subtile(int idx) { + static_assert(swizzle_bytes > 0, "Parent shared tile must have an explicit swizzle_bytes."); + constexpr int swizzle_elements = swizzle_bytes / sizeof(T); + static_assert(subtile_cols >= 0 && subtile_cols % swizzle_elements == 0); + return *reinterpret_cast *>( + &data[rows*swizzle_elements*(subtile_cols/swizzle_elements)*idx] + ); + } + + + using col_vec = sv; + using row_vec = sv; +}; + + + + + + + + + + + + + +template< + typename _ST, + int _subtile_rows, + int _subtile_cols +> +struct st_subtile { + using identifier = ducks::st::identifier; + using ST = _ST; + using T = ST::T; + using T2 = ST::T2; + using dtype = T; + + static constexpr bool swizzle = ST::swizzle; + + static constexpr int underlying_rows = ST::underlying_rows; + static_assert(underlying_rows % kittens::TILE_ROW_DIM == 0, "Underlying rows must be divisible by the tile dimension"); + static constexpr int underlying_cols = ST::underlying_cols; + static_assert(underlying_cols % kittens::TILE_COL_DIM == 0, "Underlying cols must be divisible by the tile dimension"); + static constexpr int underlying_num_elements = ST::underlying_num_elements; + + static constexpr int rows = _subtile_rows; + static_assert(rows % kittens::TILE_ROW_DIM == 0, "Rows must be divisible by the tile dimension"); + static constexpr int cols = _subtile_cols; + static_assert(cols % kittens::TILE_COL_DIM == 0, "Cols must be divisible by the tile dimension"); + static constexpr int num_elements = rows * cols; + + static constexpr int swizzle_bytes = ST::swizzle_bytes; + + dtype *data; + int row_offset, col_offset; + + __device__ st_subtile(ST &src, int2 rowcol) { + data = &src.data[0]; + row_offset = rowcol.x * rows; + col_offset = rowcol.y * cols; + } + + __device__ inline T* idx(T *ptr, const int2 coord) const { + int r = coord.x+row_offset, c = coord.y+col_offset; + if constexpr (swizzle) { + static constexpr int swizzle_repeat = swizzle_bytes * 8; + static constexpr int subtile_cols = swizzle_bytes / sizeof(T); + const int outer_idx = c/subtile_cols; + const uint64_t addr = (uint64_t)(&ptr[outer_idx*underlying_rows*subtile_cols + r*subtile_cols + c%subtile_cols]); + const int swizzle = ((addr % swizzle_repeat) >> 7) << 4; + return (T*)(addr ^ swizzle); + } else { + return &ptr[r*cols + c]; + } + } + __device__ inline uint32_t idx(uint32_t ptr, const int2 coord) const { + int r = coord.x+row_offset, c = coord.y+col_offset; + if constexpr(swizzle) { + static constexpr int swizzle_repeat = swizzle_bytes * 8; + static constexpr int subtile_cols = swizzle_bytes / sizeof(T); + const int outer_idx = c/subtile_cols; + const uint32_t addr = ptr + sizeof(T)*(outer_idx*underlying_rows*subtile_cols + r*subtile_cols + c%subtile_cols); + const int swizzle = ((addr % swizzle_repeat) >> 7) << 4; + return (addr ^ swizzle); + } else { + return ptr + sizeof(T)*(r*cols + c); + } + } + + + + + + + __device__ inline dtype& operator[](const int2 &rowcol) { + return *idx(data, rowcol); + } + __device__ inline const dtype& operator[](const int2 &rowcol) const { + return *(const dtype*)idx((dtype*)data, rowcol); + } + + + + + + using col_vec = sv; + using row_vec = sv; + + __device__ inline void operator=(const dtype &value) { + #pragma unroll + for(int i = kittens::laneid(); i < num_elements; i += WARP_THREADS) { + data[i] = value; + } + } +}; + +template +template +__device__ inline st_subtile, subtile_rows, subtile_cols> +st<_T, _rows, _cols, _swizzle, _swizzle_bytes>::subtile(int2 rowcol) +{ + + using ST_t = st<_T, _rows, _cols>; + using dtype = typename ST_t::dtype; + + + static_assert(subtile_rows > 0 && subtile_cols > 0, "Subtile dimensions must be positive."); + static_assert(subtile_rows % kittens::TILE_ROW_DIM == 0, + "Subtile rows must be divisible by the base tile row dimension."); + static_assert(subtile_cols % kittens::TILE_COL_DIM == 0, + "Subtile cols must be divisible by the base tile col dimension."); + + + static_assert(ST_t::rows % subtile_rows == 0, + "Parent tile rows must be divisible by subtile rows."); + static_assert(ST_t::cols % subtile_cols == 0, + "Parent tile cols must be divisible by subtile cols."); + + + + static_assert(ST_t::rows == ST_t::underlying_rows && ST_t::cols == ST_t::underlying_cols, + "Cannot create a subtile from an object that appears to be a subtile view (rows/cols mismatch underlying)."); + + + + + return st_subtile(*this, rowcol); +} + + + +template +using st_bf = st; +template +using st_hf = st; +template +using st_fl = st; +template +using st_int8 = st; +template +using st_uint8 = st; +template +using st_int = st; +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template +using st_fp8e4m3 = st; +template +using st_fp8e5m2 = st; +#endif +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template +using st_fp8e8m0 = st; +template +using st_fp4e2m1_2 = st; +#endif + + + + + + +template +__device__ constexpr const char* get_tile_type_name() { + if constexpr (std::is_same_v) { + return "st_fl"; + } else if constexpr (std::is_same_v) { + return "st_hf"; + } else if constexpr (std::is_same_v) { + return "st_bf"; +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + return "st_fp4_e2m1_2"; + } else if constexpr (std::is_same_v) { + return "st_fp8_e8m0"; +#endif +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + return "st_fp8_e4m3"; + } else if constexpr (std::is_same_v) { + return "st_fp8_e5m2"; +#endif + } else { + return "st_unknown"; + } +} + +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + +template +__device__ inline void print_fp4(const ST& tile) { + if (std::is_same_v) { + + constexpr int cols = ST::cols * 2; + printf("Block %d: Shared Tile %dx%d (Type: %s<%d,%d>):\n", blockIdx.x, ST::rows, cols, get_tile_type_name(), ST::rows, cols); + + + printf(" "); + for (int c = 0; c < cols; c++) { + printf("%8d ", c); + } + printf("\n"); + + + printf(" "); + for (int c = 0; c < cols; c++) { + printf("--------+"); + } + printf("\n"); + + + for (int r = 0; r < ST::rows; r++) { + printf("%3d |", r); + for (int c = 0; c < cols; c += 2) { + uint8_t *vals = reinterpret_cast(const_cast(&tile[{r,c/2}])); + + + float f1 = static_cast(fp4e2m1(vals[0] & 0xF)); + float f2 = static_cast(fp4e2m1((vals[0] >> 4) & 0xF)); + + printf("%8.3f %8.3f ", f1, f2); + + } + printf("\n"); + } + printf("\n"); + } else { + printf("Type must be FP4 in this function\n"); + } +} +#endif + + + + + + + + + +template +__device__ inline void print(const ST& tile) { + printf("Block %d: Shared Tile %dx%d (Type: %s<%d,%d>):\n", blockIdx.x, ST::rows, ST::cols, get_tile_type_name(), ST::rows, ST::cols); + + + printf(" "); + for (int c = 0; c < ST::cols; c++) { + printf("%8d ", c); + } + printf("\n"); + + + printf(" "); + for (int c = 0; c < ST::cols; c++) { + printf("--------+"); + } + printf("\n"); + + + for (int r = 0; r < ST::rows; r++) { + printf("%3d |", r); + for (int c = 0; c < ST::cols; c++) { + if constexpr (std::is_same_v) { + printf("%8.3f ", tile[{r,c}]); + } else if constexpr (std::is_same_v) { + printf("%8.3f ", __bfloat162float(tile[{r,c}])); + } else if constexpr (std::is_integral_v) { + printf("%8d ", (int)tile[{r,c}]); +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + printf("%8.3f ", static_cast(tile[{r,c}])); + } else if constexpr (std::is_same_v) { + printf("%8.3f ", static_cast(tile[{r,c}])); +#endif +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + printf("%8.3f ", static_cast(tile[{r,c}])); +#endif + } else { + printf("%8.3f ", (float)tile[{r,c}]); + } + } + printf("\n"); + } + printf("\n"); +} + + + + + + + + + +template +__device__ inline void print(const st_subtile& subtile) { + printf("Block %d: Shared Subtile %dx%d (offset: [%d,%d], Type: %s<%d,%d> from %s<%d,%d>):\n", + blockIdx.x, subtile.rows, subtile.cols, + subtile.row_offset, subtile.col_offset, + get_tile_type_name(), subtile.rows, subtile.cols, + get_tile_type_name(), ST::rows, ST::cols); + + + printf(" "); + for (int c = 0; c < subtile.cols; c++) { + printf("%8d ", c); + } + printf("\n"); + + + printf(" "); + for (int c = 0; c < subtile.cols; c++) { + printf("--------+"); + } + printf("\n"); + + + for (int r = 0; r < subtile.rows; r++) { + printf("%3d |", r); + for (int c = 0; c < subtile.cols; c++) { + if constexpr (std::is_same_v) { + printf("%8.3f ", subtile[{r,c}]); + } else if constexpr (std::is_same_v) { + printf("%8.3f ", __bfloat162float(subtile[{r,c}])); + } else if constexpr (std::is_integral_v) { + printf("%8d ", (int)subtile[{r,c}]); + } else { + printf("%8.3f ", (float)subtile[{r,c}]); + } + } + printf("\n"); + } + printf("\n"); +} + + + + + + + + + +template +__device__ inline void fill_value(ST& tile, float value) { + printf("Filling Tile %dx%d with %f:\n", ST::rows, ST::cols, value); + + + for (int r = 0; r < ST::rows; r++) { + for (int c = 0; c < ST::cols; c++) { + if constexpr (std::is_same_v) { + tile[{r,c}] = value; + } else if constexpr (std::is_same_v) { + tile[{r,c}] = __float2bfloat16(value); +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + tile.data[r*ST::cols + c] = fp4e2m1(value); + } else if constexpr (std::is_same_v) { + tile.data[r*ST::cols + c] = fp8e8m0(value); +#endif + } else { + tile[{r,c}] = value; + } + } + } +} + + + + + + + + + +template +__device__ inline void fill_identity(ST& tile) { + printf("Filling Tile %dx%d with identity:\n", ST::rows, ST::cols); + + + for (int r = 0; r < ST::rows; r++) { + for (int c = 0; c < ST::cols; c++) { + if constexpr (std::is_same_v) { + + tile[{r,c}] = 1.0f; + } else if constexpr (std::is_same_v) { + + tile[{r,c}] = __float2bfloat16(1.0f); +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + if(r == c){ + + tile.data[r*ST::cols + c] = std::bit_cast(uint8_t(0xFF)); + } else { + tile.data[r*ST::cols + c] = std::bit_cast(uint8_t(0x00)); + } +#endif + } else { + + } + } + } +} + + + + + + + + + + +template +__device__ inline void print_bits(const ST& tile, bool unswizzle = false) { + printf("Block %d: Shared Tile %dx%d (Type: %s<%d,%d>):\n", blockIdx.x, ST::rows, ST::cols, get_tile_type_name(), ST::rows, ST::cols); + + + printf(" "); + for (int c = 0; c < ST::cols; c++) { + printf("%11d ", c); + } + printf("\n"); + + + printf(" "); + for (int c = 0; c < ST::cols; c++) { + printf("-----------+"); + } + printf("\n"); + + + for (int r = 0; r < ST::rows; r++) { + printf("%3d |", r); + for (int c = 0; c < ST::cols; c++) { + if constexpr (std::is_same_v) { + printf("%8.3f ", tile[{r,c}]); + } else if constexpr (std::is_same_v) { + printf("%8.3f ", __bfloat162float(tile[{r,c}])); + + +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v || std::is_same_v) { + + + uint8_t bits; + if(unswizzle){ + bits = *reinterpret_cast(&tile.data[r*ST::cols + c]); + } else { + bits = *reinterpret_cast(&tile[{r,c}]); + } + + printf("0b"); + + for (int bit = 7; bit >= 0; bit--) { + printf("%d", (bits >> bit) & 0x1); + if (bit % 4 == 0 && bit > 0) printf("_"); + } + printf(" "); +#endif + } else { + printf("%8.3f ", (float)tile[{r,c}]); + } + } + printf("\n"); + } + printf("\n"); +} + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/sv.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/sv.cuh new file mode 100644 index 0000000..1bf6120 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/sv.cuh @@ -0,0 +1,155 @@ + + + + + +#pragma once + +#include "../../common/common.cuh" + +namespace kittens { + + + +namespace ducks { + + + + + +namespace sv { + + + + + + +struct identifier {}; + + + + + + + +template +concept all = requires { + typename T::identifier; +} && std::is_same_v; +} +} + + + + + + + + + + + +template +struct KITTENS_DEFAULT_ALIGN sv { + static_assert(!std::is_same_v<_T, char>, "Bare 'char' is not permitted because its signedness is ISA-dependent."); + using identifier = ducks::sv::identifier; + using T = base_types::packing<_T>::unpacked_type; + using T2 = base_types::packing<_T>::packed_type; + using dtype = T; + + static constexpr int length = _length; + static_assert(length % TILE_ROW_DIM == 0, "Length must be divisible by the tile dimension"); + static constexpr int tiles = length / TILE_ROW_DIM; + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static constexpr int num_alloc_elements = ((length * sizeof(dtype) + 127) / 128) * (128 / sizeof(dtype)); +#else + static constexpr int num_alloc_elements = length; +#endif + dtype data[num_alloc_elements]; + + __device__ static inline T* idx(T *ptr, int idx) { + return ptr[idx]; + } + + __device__ inline dtype& operator[](size_t idx) { return data[idx]; } + __device__ inline const dtype& operator[](size_t idx) const { return data[idx]; } + + template __device__ inline sv<_T, sub_length> &subvec(int idx) { + return *(sv*)&data[idx * sub_length]; + } + template __device__ inline const sv<_T, sub_length> &subvec(int idx) const { + return *(sv*)&data[idx * sub_length]; + } + + __device__ inline void operator=(const dtype &value) { + #pragma unroll + for(int i = kittens::laneid(); i < length; i += WARP_THREADS) { + data[i] = value; + } + } +}; + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +namespace detail { +namespace tma { + + + +template struct find_vector_divider { + static constexpr int value = (SV::length % (16*D) == 0 && (SV::length < 256 || ((16*D)*sizeof(typename SV::dtype)) % 128 == 0)) ? + 16*D : find_vector_divider::value; +}; +template struct find_vector_divider { static constexpr int value = 16; }; +template constexpr int sv_tma_dim1 = find_vector_divider::value; +template constexpr int sv_tma_dim2 = (SV::length / sv_tma_dim1); +} +} +#endif + + + + +template using sv_bf = sv; +template using sv_hf = sv; +template using sv_fl = sv; +template using sv_int8 = sv; +template using sv_uint8 = sv; +template using sv_int = sv; +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template using sv_fp8e4m3 = sv; +template using sv_fp8e5m2 = sv; +#endif +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template using sv_fp8e8m0 = sv; +template using sv_fp4e2m1_2 = sv; +#endif + + + +template +__device__ inline void print(const SV& sv) { + printf("Shared Vector %d:\n", SV::length); + for(int i = 0; i < SV::length; i++) { + if constexpr (std::is_same_v) { + printf("%f ", __bfloat162float(sv[i])); + } else if constexpr (std::is_same_v) { + printf("%f ", __half2float(sv[i])); + } else if constexpr (std::is_same_v) { + printf("%f ", sv[i]); +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + printf("%f ", static_cast(sv[i])); +#endif +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + printf("%f ", static_cast(sv[i])); +#endif + } else { + printf("%d ", (int)(sv[i])); + } + } + printf("\n"); +} + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/system/ipc.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/system/ipc.cuh new file mode 100644 index 0000000..91463fc --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/system/ipc.cuh @@ -0,0 +1,143 @@ +#pragma once + +#include "../../common/common.cuh" +#include "vmm.cuh" + +namespace kittens { +namespace ducks { +namespace ipc { +namespace handle { + +struct identifier {}; + +template concept all = requires { + typename T::identifier; +} && std::is_same_v; + +} +} +} + +namespace detail { +namespace ipc { + +enum flavor { + LEGACY = 0, + VMM = 1 +}; + +template +struct handle; + +template<> +struct handle { + using identifier = ducks::ipc::handle::identifier; + static constexpr flavor flavor_ = flavor::LEGACY; + cudaIpcMemHandle_t handle_ {}; +}; + +template<> +struct handle { + using identifier = ducks::ipc::handle::identifier; + static constexpr flavor flavor_ = flavor::VMM; + int handle_; +}; + +__host__ inline static void check_support(const int device_id) { + CUdevice device; + CUCHECK(cuDeviceGet(&device, device_id)); + + int ipc_supported = 0; + CUDACHECK(cudaDeviceGetAttribute(&ipc_supported, cudaDevAttrIpcEventSupport, device_id)); + int ipc_handle_supported = 0; + CUCHECK(cuDeviceGetAttribute(&ipc_handle_supported, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR_SUPPORTED, device)); + + if (!ipc_supported || !ipc_handle_supported) + throw std::runtime_error("CUDA IPC is not supported on this device"); +} + +template +__host__ inline static void export_handle( + IPC_HANDLE *ipc_handle, + void *ptr +) { + if constexpr (IPC_HANDLE::flavor_ == flavor::LEGACY) { + CUDACHECK(cudaIpcGetMemHandle(&ipc_handle->handle_, ptr)); + } else if constexpr (IPC_HANDLE::flavor_ == flavor::VMM) { + CUmemGenericAllocationHandle memory_handle; + detail::vmm::vm_retrieve_handle(&memory_handle, ptr); + + CUCHECK(cuMemExportToShareableHandle(&ipc_handle->handle_, memory_handle, detail::vmm::HANDLE_TYPE, 0)); + detail::vmm::vm_free(memory_handle); + } else { + throw std::runtime_error("Invalid IPC handle type"); + } +} + +template +__host__ inline static void export_handle( + IPC_HANDLE *ipc_handle, + CUmemGenericAllocationHandle &memory_handle +) { + if constexpr (IPC_HANDLE::flavor_ == flavor::VMM) { + CUCHECK(cuMemExportToShareableHandle(&ipc_handle->handle_, memory_handle, detail::vmm::HANDLE_TYPE, 0)); + } else { + throw std::runtime_error("Invalid IPC handle type"); + } +} + +template +__host__ inline static void import_handle ( + void **ptr, + IPC_HANDLE &ipc_handle, + const size_t size, + int local_world_size +) { + if constexpr (IPC_HANDLE::flavor_ == flavor::LEGACY) { + CUDACHECK(cudaIpcOpenMemHandle(ptr, ipc_handle.handle_, cudaIpcMemLazyEnablePeerAccess)); + } else if constexpr (IPC_HANDLE::flavor_ == flavor::VMM) { + CUmemGenericAllocationHandle memory_handle; + CUCHECK(cuMemImportFromShareableHandle(&memory_handle, reinterpret_cast(static_cast(ipc_handle.handle_)), detail::vmm::HANDLE_TYPE)); + detail::vmm::vm_map(ptr, memory_handle, size); + detail::vmm::vm_set_access(*ptr, size, local_world_size); + detail::vmm::vm_free(memory_handle); + close(ipc_handle.handle_); + ipc_handle.handle_ = -1; + } else { + throw std::runtime_error("Invalid IPC handle type"); + } +} + +template +__host__ inline static void import_handle ( + CUmemGenericAllocationHandle *memory_handle, + IPC_HANDLE &ipc_handle, + const size_t size, + int local_world_size +) { + if constexpr (IPC_HANDLE::flavor_ == flavor::VMM) { + CUCHECK(cuMemImportFromShareableHandle(memory_handle, reinterpret_cast(static_cast(ipc_handle.handle_)), detail::vmm::HANDLE_TYPE)); + close(ipc_handle.handle_); + ipc_handle.handle_ = -1; + } else { + throw std::runtime_error("Invalid IPC handle type"); + } +} + +template +__host__ inline static void free_handle( + void *ptr, + const size_t size +) { + if constexpr (_flavor == flavor::LEGACY) { + CUDACHECK(cudaIpcCloseMemHandle(ptr)); + } else if constexpr (_flavor == flavor::VMM) { + detail::vmm::vm_unmap(ptr, size); + } else { + throw std::runtime_error("Invalid IPC handle type"); + } +} + +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/system/pgl.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/system/pgl.cuh new file mode 100644 index 0000000..d3a0492 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/system/pgl.cuh @@ -0,0 +1,182 @@ + + + + + +#pragma once + +#include "../../common/common.cuh" +#include "../shared/shared.cuh" +#include "../global/global.cuh" + +namespace kittens { + + + +namespace ducks { +namespace pgl { + +struct identifier {}; + + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v; + +} +} + + + + + + + + + +template +struct pgl { + using identifier = ducks::pgl::identifier; + using GL = _GL; + using T = GL::dtype; + static_assert(!std::is_same_v, "Bare 'char' is not permitted because its signedness is ISA-dependent."); + using dtype = T; + + static constexpr int num_devices = NUM_DEVICES; + static constexpr bool multicast = MULTICAST; + + T *mc_ptr; + GL gls[NUM_DEVICES]; + + detail::descriptor_dict tma_descs; + + __host__ __device__ const GL &operator[](int idx) const { return gls[idx]; } + __device__ inline T* mc_ptr_at(const coord &idx) const { + static_assert(MULTICAST, "Multicast is not enabled for this PGL."); + const GL &gl = gls[0]; + return &mc_ptr[((idx.b * static_cast(gl.depth()) + idx.d) * gl.rows() + idx.r) * gl.cols() + idx.c]; + } + + __host__ __device__ pgl() = delete; +#ifndef KITTENS_NO_HOST + __host__ inline pgl(T **_data, + ducks::gl::make_arg_t _batch, + ducks::gl::make_arg_t _depth, + ducks::gl::make_arg_t _rows, + ducks::gl::make_arg_t _cols) : + pgl(std::make_index_sequence{}, _data, _batch, _depth, _rows, _cols) { } + __host__ inline pgl(T *_mc_ptr, + T **_data, + ducks::gl::make_arg_t _batch, + ducks::gl::make_arg_t _depth, + ducks::gl::make_arg_t _rows, + ducks::gl::make_arg_t _cols) : + pgl(std::make_index_sequence{}, _mc_ptr, _data, _batch, _depth, _rows, _cols) { } + template + __host__ inline pgl(std::index_sequence, + T **_data, + ducks::gl::make_arg_t _batch, + ducks::gl::make_arg_t _depth, + ducks::gl::make_arg_t _rows, + ducks::gl::make_arg_t _cols) : + mc_ptr(nullptr), gls{GL(_data[I], _batch, _depth, _rows, _cols)...} { + static_assert(!MULTICAST, "Multicast pointer not passed to multicast-enabled PGL."); + } + template + __host__ inline pgl(std::index_sequence, + T *_mc_ptr, + T **_data, + ducks::gl::make_arg_t _batch, + ducks::gl::make_arg_t _depth, + ducks::gl::make_arg_t _rows, + ducks::gl::make_arg_t _cols) : + mc_ptr(_mc_ptr), gls{GL(_data[I], _batch, _depth, _rows, _cols)...} { + static_assert(MULTICAST, "Multicast pointer passed to multicast-disabled PGL."); + tma_descs = detail::descriptor_dict( + mc_ptr, gls[0].batch_internal, gls[0].depth_internal, gls[0].rows_internal, gls[0].cols_internal); + } +#endif + + template + __device__ inline const CUtensorMap* get_tma() const { + return tma_descs.template get(); + } + + __host__ __device__ inline auto batch() const { return gls[0].batch(); } + __host__ __device__ inline auto depth() const { return gls[0].depth(); } + __host__ __device__ inline auto rows() const { return gls[0].rows(); } + __host__ __device__ inline auto cols() const { return gls[0].cols(); } + __host__ __device__ inline auto numel() const { return gls[0].numel(); } + + template __device__ inline size_t shape() const { return gls[0].template shape(); } + template __device__ inline size_t stride() const { return gls[0].template stride(); } +}; + +#ifndef KITTENS_NO_HOST + +template __host__ inline PGL make_pgl( + uint64_t *data, int b, int d, int r, int c +) { + if constexpr (safe) { + if (PGL::GL::__b__ > 0 && b != PGL::GL::__b__) { + throw std::runtime_error("Batch dimension mismatch. Expected: " + std::to_string(PGL::GL::__b__) + ", Got: " + std::to_string(b)); + } + if (PGL::GL::__d__ > 0 && d != PGL::GL::__d__) { + throw std::runtime_error("Depth dimension mismatch. Expected: " + std::to_string(PGL::GL::__d__) + ", Got: " + std::to_string(d)); + } + if (PGL::GL::__r__ > 0 && r != PGL::GL::__r__) { + throw std::runtime_error("Row dimension mismatch. Expected: " + std::to_string(PGL::GL::__r__) + ", Got: " + std::to_string(r)); + } + if (PGL::GL::__c__ > 0 && c != PGL::GL::__c__) { + throw std::runtime_error("Column dimension mismatch. Expected: " + std::to_string(PGL::GL::__c__) + ", Got: " + std::to_string(c)); + } + } + return PGL( + reinterpret_cast(data), + make_unsafe_gl_arg(b), + make_unsafe_gl_arg(d), + make_unsafe_gl_arg(r), + make_unsafe_gl_arg(c) + ); +} + +template __host__ inline PGL make_pgl( + uint64_t mc_ptr, uint64_t *data, int b, int d, int r, int c +) { + if constexpr (safe) { + if (PGL::GL::__b__ > 0 && b != PGL::GL::__b__) { + throw std::runtime_error("Batch dimension mismatch. Expected: " + std::to_string(PGL::GL::__b__) + ", Got: " + std::to_string(b)); + } + if (PGL::GL::__d__ > 0 && d != PGL::GL::__d__) { + throw std::runtime_error("Depth dimension mismatch. Expected: " + std::to_string(PGL::GL::__d__) + ", Got: " + std::to_string(d)); + } + if (PGL::GL::__r__ > 0 && r != PGL::GL::__r__) { + throw std::runtime_error("Row dimension mismatch. Expected: " + std::to_string(PGL::GL::__r__) + ", Got: " + std::to_string(r)); + } + if (PGL::GL::__c__ > 0 && c != PGL::GL::__c__) { + throw std::runtime_error("Column dimension mismatch. Expected: " + std::to_string(PGL::GL::__c__) + ", Got: " + std::to_string(c)); + } + } + return PGL( + reinterpret_cast(mc_ptr), + reinterpret_cast(data), + make_unsafe_gl_arg(b), + make_unsafe_gl_arg(d), + make_unsafe_gl_arg(r), + make_unsafe_gl_arg(c) + ); +} + +#endif + + +template +using barrier_t = pgl, NUM_DEVICES, true>; + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/system/system.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/system/system.cuh new file mode 100644 index 0000000..067ca81 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/system/system.cuh @@ -0,0 +1,14 @@ + + + + + +#pragma once + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +#ifndef KITTENS_NO_HOST +#include "ipc.cuh" +#include "vmm.cuh" +#endif +#include "pgl.cuh" +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/system/vmm.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/system/vmm.cuh new file mode 100644 index 0000000..b11cd85 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/system/vmm.cuh @@ -0,0 +1,176 @@ +#pragma once + +#include "../../common/common.cuh" + +namespace kittens { +namespace detail { +namespace vmm { + + + +static constexpr CUmemAllocationHandleType HANDLE_TYPE = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR; + +typedef CUmemGenericAllocationHandle handle; + +__host__ inline static void vm_alloc( + CUmemGenericAllocationHandle *handle, + size_t *allocated_size, + const size_t size, + const int device_id +) { + CUmemAllocationProp prop = {}; + prop.location.id = device_id; + prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + prop.requestedHandleTypes = HANDLE_TYPE; + prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; + + size_t granularity; + CUCHECK(cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED)); + *allocated_size = (size + granularity - 1) / granularity * granularity; + + CUCHECK(cuMemCreate(handle, *allocated_size, &prop, 0)); +} + +__host__ inline static void vm_map( + void **ptr, + const CUmemGenericAllocationHandle &handle, + const size_t size +) { + CUdeviceptr device_ptr; + CUCHECK(cuMemAddressReserve(&device_ptr, size, 0, 0, 0)); + CUCHECK(cuMemMap(device_ptr, size, 0, handle, 0)); + *ptr = (void *)device_ptr; +} + +__host__ inline static void vm_set_access( + void *ptr, + const size_t size, + const int num_devices +) { + std::vector descs(num_devices); + for (int i = 0; i < num_devices; i++) { + descs[i].location.id = i; + descs[i].location.type = CU_MEM_LOCATION_TYPE_DEVICE; + descs[i].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; + } + CUCHECK(cuMemSetAccess(reinterpret_cast(ptr), size, descs.data(), num_devices)); +} + +__host__ inline static void vm_retrieve_handle( + CUmemGenericAllocationHandle *handle, + void *ptr +) { + + CUCHECK(cuMemRetainAllocationHandle(handle, ptr)); +} + +__host__ inline static void vm_unmap( + void *ptr, + const size_t size +) { + CUCHECK(cuMemUnmap(reinterpret_cast(ptr), size)); + CUCHECK(cuMemAddressFree(reinterpret_cast(ptr), size)); +} + +__host__ inline static void vm_free(CUmemGenericAllocationHandle &handle) { + + + CUCHECK(cuMemRelease(handle)); +} + +__host__ inline static void vm_alloc_map_set_access( + void **ptr, + size_t *allocated_size, + const size_t size, + const int device_id, + const int num_devices +) { + CUmemGenericAllocationHandle handle; + vm_alloc(&handle, allocated_size, size, device_id); + vm_map(ptr, handle, *allocated_size); + vm_set_access(*ptr, *allocated_size, num_devices); + vm_free(handle); +} + +__host__ inline static void multicast_check(const int device_id) { + CUdevice device; + CUCHECK(cuDeviceGet(&device, device_id)); + + int multicast_supported; + CUresult result = cuDeviceGetAttribute( + &multicast_supported, + CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, + device + ); + + if (!multicast_supported) + throw std::runtime_error("Device does not support multicast"); +} + +__host__ inline static void multicast_create_handle( + CUmemGenericAllocationHandle *handle, + size_t *allocated_size, + const size_t size, + const int num_devices +) { + if (num_devices <= 1) + throw std::runtime_error("Multicast requires at least 2 devices"); + + CUmulticastObjectProp prop = {}; + prop.numDevices = num_devices; + prop.handleTypes = HANDLE_TYPE; + + size_t granularity; + CUCHECK(cuMulticastGetGranularity(&granularity, &prop, CU_MULTICAST_GRANULARITY_RECOMMENDED)); + *allocated_size = (size + granularity - 1) / granularity * granularity; + prop.size = *allocated_size; + + + cuMulticastCreate(handle, &prop); +} + +__host__ inline static void multicast_bind_device( + const CUmemGenericAllocationHandle &handle, + const int device_id +) { + + CUdevice device; + CUCHECK(cuDeviceGet(&device, device_id)); + CUCHECK(cuMulticastAddDevice(handle, device)); +} + +__host__ inline static void multicast_bind_memory( + const CUmemGenericAllocationHandle &multicast_handle, + const CUmemGenericAllocationHandle &memory_handle, + const size_t size +) { + + CUCHECK(cuMulticastBindMem(multicast_handle, 0, memory_handle, 0, size, 0)); +} + +__host__ inline static void multicast_bind_address( + const CUmemGenericAllocationHandle &multicast_handle, + void *ptr, + const size_t size +) { + + CUmemGenericAllocationHandle memory_handle; + vm_retrieve_handle(&memory_handle, ptr); + multicast_bind_memory(multicast_handle, memory_handle, size); + vm_free(memory_handle); +} + +__host__ inline static void multicast_unbind_device( + const CUmemGenericAllocationHandle &handle, + const size_t size, + const int device_id +) { + + CUdevice device; + CUCHECK(cuDeviceGet(&device, device_id)); + CUCHECK(cuMulticastUnbind(handle, device, 0, size)); +} + +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/tensor/tensor.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/tensor/tensor.cuh new file mode 100644 index 0000000..052967c --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/tensor/tensor.cuh @@ -0,0 +1,152 @@ + + + + + +#pragma once + +#include "tt.cuh" + + +namespace kittens { +namespace ducks { + + + + + +namespace tensor_allocator { + + + +struct identifier {}; + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v; +} +} + +template struct tensor_allocator { + static_assert(_nblocks_per_sm == 1 || _nblocks_per_sm == 2, "nblocks_per_sm must be 1 or 2"); + static_assert(_ncta == 1 || _ncta == 2, "ncta must be 1 or 2"); + + using identifier = ducks::tensor_allocator::identifier; + + static constexpr int nblocks_per_sm = _nblocks_per_sm; + static constexpr int cols =((MAX_TENSOR_COLS/nblocks_per_sm) / 32) * 32; + static_assert(cols>0 && cols%32==0, "cols must be a multiple of 32"); + static constexpr int ncta = _ncta; + static constexpr bool managed = _managed; + + uint32_t addr; + + template __device__ inline void check_bounds() { + static_assert(col_offset >= 0 && col_offset + TT::cols <= cols, "Tile allocation extends out of bounds of the tensor allocator!"); + } + + __device__ inline void set_addr(uint32_t _addr) { + addr = _addr; + } + + __device__ inline uint32_t get_addr(int superlane, int col_offset) const { + return addr + ((superlane*16) << 16) + col_offset; + } + + __device__ inline uint32_t get_addr(int col_offset) const { + return addr + col_offset; + } + + __device__ inline void provision(uint32_t &shared_addr) { + + + + + if constexpr (ncta == 1) { + asm volatile( + "tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;\n" + :: "l"(reinterpret_cast(&shared_addr)), "n"(cols) + ); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;\n"); + } + else { + asm volatile( + "tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;\n" + :: "l"(reinterpret_cast(&shared_addr)), "n"(cols) + ); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;\n"); + } + } + + __device__ inline tensor_allocator() { + if constexpr (managed) { + __shared__ uint32_t shared_addr; + if (warpid() == 0) provision(shared_addr); + asm volatile("tcgen05.fence::before_thread_sync;\n"); + asm volatile("bar.sync 0;\n"); + asm volatile("tcgen05.fence::after_thread_sync;\n"); + set_addr(shared_addr); + } + } + + template __device__ inline auto allocate(int superlane, int col_offset) { +#ifndef NDEBUG + int allocate_cols = std::is_same_v ? TT::cols/4 : TT::cols; + if(col_offset + allocate_cols > cols) { + if(laneid() == 0) printf("Tile allocation extends out of bounds of the tensor allocator! col_offset: %d, TT::cols: %d, allocator cols: %d\n", col_offset, TT::cols, cols); + asm volatile("trap;"); + } + if(superlane < 0 || superlane > 1) { + printf("Superlane must be 0 or 1! superlane: %d\n", superlane); + asm volatile("trap;"); + } +#endif + return TT(get_addr(superlane, col_offset)); + } + + template __device__ inline auto allocate(int col_offset) { +#ifndef NDEBUG + int allocate_cols = std::is_same_v ? TT::cols/4 : TT::cols; + if(col_offset + allocate_cols > cols) { + if(laneid() == 0) printf("Tile allocation extends out of bounds of the tensor allocator! col_offset: %d, TT::cols: %d, allocator cols: %d\n", col_offset, TT::cols, cols); + asm volatile("trap;"); + } +#endif + return TT(get_addr(0, col_offset)); + } + + __device__ inline void deprovision() { + + + + if constexpr (ncta == 1) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;\n" + :: "r"(addr), "n"(cols) + ); + } else { + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;\n" + :: "r"(addr), "n"(cols) + ); + } + } + + __device__ inline ~tensor_allocator() { + if constexpr (managed) { + if (warpid() == 0) { + if (ncta == 2) { + asm volatile ("barrier.cluster.arrive.release.aligned;\n"); + asm volatile ("barrier.cluster.wait.acquire.aligned;\n"); + } + deprovision(); + } + } + } +}; + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/tensor/tt.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/tensor/tt.cuh new file mode 100644 index 0000000..7d37e42 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/tensor/tt.cuh @@ -0,0 +1,135 @@ + + + + + +#pragma once + +#include "../../common/common.cuh" + + + + +namespace kittens { + +constexpr int MAX_TENSOR_ROWS = 128; +constexpr int MAX_TENSOR_COLS = 512; + +namespace ducks { + + + + + +namespace tt { + + + +struct identifier {}; + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v; +template concept half = all && T::rows == MAX_TENSOR_ROWS / 2; +template concept full = all && T::rows == MAX_TENSOR_ROWS; +} +} + + + + + + + + +template +struct tt { + static_assert(!std::is_same_v<_T, char>, "Bare 'char' is not permitted because its signedness is ISA-dependent."); + using identifier = ducks::tt::identifier; + using T = base_types::packing<_T>::unpacked_type; + using T2 = base_types::packing<_T>::packed_type; + using dtype = T; + + static constexpr int rows = _rows; + static constexpr int cols = _cols; + + static_assert(rows / (4 / sizeof(T)) <= MAX_TENSOR_ROWS, "Row dimension must be less than or equal to MAX_TENSOR_ROWS"); + static_assert(cols / (4 / sizeof(T)) <= MAX_TENSOR_COLS, "Column dimension must be less than or equal to MAX_TENSOR_COLS"); + static_assert(rows % kittens::BASE_TILE_DIM == 0, "Row dimension must be divisible by the 16"); + static_assert(cols % kittens::BASE_TILE_DIM == 0, "Column dimension must be divisible by the 16"); + + uint32_t addr; + + __device__ inline tt() : addr(0) {} + __device__ inline tt(uint32_t addr) : addr(addr) {} + + template __device__ inline TT subtile(int row_offset, int col_offset) const { +#ifndef NDEBUG + if(row_offset < 0 || row_offset+TT::rows > rows || col_offset < 0 || col_offset+TT::cols > cols) { + printf("Subtile out of bounds! full tile rows: %d, full tile cols: %d, subtile rows: %d, subtile cols: %d, row_offset: %d, col_offset: %d\n", rows, cols, TT::rows, TT::cols, row_offset, col_offset); + asm volatile("trap;"); + } +#endif + return TT(addr + (row_offset<<16) + col_offset/(4/(uint32_t)sizeof(T))); + } + template __device__ inline TT subtile(int col_offset) const { + return TT(addr + col_offset/(4/(uint32_t)sizeof(T))); + } + template __device__ inline uint32_t chunk_addr(int chunk) const { + if constexpr (transpose) { + static_assert(!transpose, "Transposed TMEM chunk addressing is not supported due to Blackwell tensor-core limitations."); + } + else { + if constexpr (std::is_same_v || std::is_same_v) { + return addr + (16 * chunk / (4/(uint32_t)sizeof(T))); + } + else if constexpr (std::is_same_v || std::is_same_v || std::is_same_v || std::is_same_v) { + return addr + (32 * chunk / (4/(uint32_t)sizeof(T))); + } + else { + static_assert(sizeof(T) == 999, "Currently unsupported type for input to an mma."); + } + } + } +}; + + + +template using tt_bf = tt; +template using tt_hf = tt; +template using tt_fl = tt; +template using tt_int8 = tt; +template using tt_uint8 = tt; +template using tt_int = tt; +template using tt_fp8e4m3 = tt; +template using tt_fp8e5m2 = tt; +template using tt_fp8e8m0 = tt; +template using tt_fp4e2m1_2 = tt; +template using half_tt_bf = tt; +template using half_tt_hf = tt; +template using half_tt_fl = tt; +template using half_tt_int8 = tt; +template using half_tt_uint8 = tt; +template using half_tt_int = tt; +template using half_tt_fp8e4m3 = tt; +template using half_tt_fp8e5m2 = tt; +template using half_tt_fp8e8m0 = tt; +template using half_tt_fp4e2m1_2 = tt; +template using full_tt_bf = tt; +template using full_tt_hf = tt; +template using full_tt_fl = tt; +template using full_tt_int8 = tt; +template using full_tt_uint8 = tt; +template using full_tt_int = tt; +template using full_tt_fp8e4m3 = tt; +template using full_tt_fp8e5m2 = tt; +template using full_tt_fp8e8m0 = tt; +template using full_tt_fp4e2m1_2 = tt; + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/types.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/types.cuh new file mode 100644 index 0000000..d665854 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/types.cuh @@ -0,0 +1,67 @@ + + + + + +#pragma once + +#include "register/register.cuh" +#include "shared/shared.cuh" +#include "global/global.cuh" +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +#include "system/system.cuh" +#endif +#ifdef KITTENS_SM10X +#include "tensor/tensor.cuh" +#endif + + + +namespace kittens { + + + + + + + + + + + + + + + +template +using row_vec = T::row_vec; + + + + + + + + + + + + + + + +template +using col_vec = T::col_vec; + + + + +using row_l = ducks::rt_layout::row; +using col_l = ducks::rt_layout::col; + + +using align_l = ducks::rv_layout::align; +using ortho_l = ducks::rv_layout::ortho; +using naive_l = ducks::rv_layout::naive; + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/LICENSE b/deploy/nemotron_native/kernels/third_party/thunderkittens/LICENSE new file mode 100644 index 0000000..79e0fdb --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024-2026 HazyResearch + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/common/base_ops.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/common/base_ops.cuh new file mode 100644 index 0000000..a5a36da --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/common/base_ops.cuh @@ -0,0 +1,446 @@ + + + + + +#pragma once + +#include "base_types.cuh" + +namespace kittens { + + + + + + +namespace base_ops { + + + + + + + + + + + +struct zero { + template __device__ static inline constexpr T op(args... _) { return base_types::constants::zero(); } +}; + + + + + + + + +struct one { + template __device__ static inline constexpr T op(args... _) { return base_types::constants::one(); } +}; + + + + + + + + +struct pos_infty { + template __device__ static inline constexpr T op(args... _) { return base_types::constants::pos_infty(); } +}; + + + + + + + + +struct neg_infty { + template __device__ static inline constexpr T op(args... _) { return base_types::constants::neg_infty(); } +}; + + + + + + + + + + + + + +struct exp { + template static __device__ inline T op(const T &x) { return exp(x); } +}; +template<> __device__ inline float exp::op (const float &x ) { return __expf(x); } +template<> __device__ inline float2 exp::op(const float2 &x) { return float2{__expf(x.x), __expf(x.y)}; } +template<> __device__ inline bf16 exp::op (const bf16 &x ) { return hexp(x); } +template<> __device__ inline bf16_2 exp::op(const bf16_2 &x) { return h2exp(x); } +template<> __device__ inline half exp::op (const half &x ) { return hexp(x); } +template<> __device__ inline half_2 exp::op(const half_2 &x) { return h2exp(x); } + + + + + + + + + + +struct exp2 { + template static __device__ inline T op(const T &x) { return exp2f(x); } +}; +template<> __device__ inline float exp2::op (const float &x ) { return exp2f(x); } +template<> __device__ inline float2 exp2::op(const float2 &x) { return float2{exp2f(x.x), exp2f(x.y)}; } +template<> __device__ inline bf16 exp2::op (const bf16 &x ) { return hexp2(x); } +template<> __device__ inline bf16_2 exp2::op(const bf16_2 &x) { return h2exp2(x); } +template<> __device__ inline half exp2::op (const half &x ) { return hexp2(x); } +template<> __device__ inline half_2 exp2::op(const half_2 &x) { return h2exp2(x); } + + + + + + + + + +struct log { + template static __device__ inline T op(const T &x) { return log(x); } +}; +template<> __device__ inline float log::op (const float &x ) { return __logf(x); } +template<> __device__ inline float2 log::op(const float2 &x) { return float2{__logf(x.x), __logf(x.y)}; } +template<> __device__ inline bf16 log::op (const bf16 &x ) { return hlog(x); } +template<> __device__ inline bf16_2 log::op(const bf16_2 &x) { return h2log(x); } +template<> __device__ inline half log::op (const half &x ) { return hlog(x); } +template<> __device__ inline half_2 log::op(const half_2 &x) { return h2log(x); } + + + + + + + + + +struct log2 { + template static __device__ inline T op(const T &x) { return log2(x); } +}; +template<> __device__ inline float log2::op (const float &x ) { return __log2f(x); } +template<> __device__ inline float2 log2::op(const float2 &x) { return float2{__log2f(x.x), __log2f(x.y)}; } +template<> __device__ inline bf16 log2::op (const bf16 &x ) { return hlog2(x); } +template<> __device__ inline bf16_2 log2::op(const bf16_2 &x) { return h2log2(x); } +template<> __device__ inline half log2::op (const half &x ) { return hlog2(x); } +template<> __device__ inline half_2 log2::op(const half_2 &x) { return h2log2(x); } + + + + + + + + + +struct abs { + template static __device__ inline T op(const T &x) { return abs(x); } +}; +template<> __device__ inline float abs::op (const float &x ) { return fabsf(x); } +template<> __device__ inline float2 abs::op(const float2 &x) { return float2{fabsf(x.x), fabsf(x.y)}; } +template<> __device__ inline bf16 abs::op (const bf16 &x ) { return __habs(x); } +template<> __device__ inline bf16_2 abs::op(const bf16_2 &x) { return __habs2(x); } +template<> __device__ inline half abs::op (const half &x ) { return __habs(x); } +template<> __device__ inline half_2 abs::op(const half_2 &x) { return __habs2(x); } + + + + + + + + + + +struct relu { + template static __device__ inline T op(const T &x) { return max(x, base_types::constants::zero()); } +}; +template<> __device__ inline float relu::op (const float &x ) { return max(x, 0.f); } +template<> __device__ inline float2 relu::op(const float2 &x) { return float2{max(x.x, 0.f), max(x.y, 0.f)}; } +template<> __device__ inline bf16 relu::op (const bf16 &x ) { return __hmax(x, base_types::constants::zero()); } +template<> __device__ inline bf16_2 relu::op(const bf16_2 &x) { return __hmax2(x, base_types::constants::zero()); } +template<> __device__ inline half relu::op (const half &x ) { return __hmax(x, base_types::constants::zero()); } +template<> __device__ inline half_2 relu::op(const half_2 &x) { return __hmax2(x, base_types::constants::zero()); } + + + + + + + +struct neg { + template static __device__ inline T op(const T &x) { return -x; } +}; +template<> __device__ inline float neg::op (const float &x ) { return -x; } +template<> __device__ inline float2 neg::op (const float2 &x ) { return float2{-x.x, -x.y}; } +template<> __device__ inline bf16 neg::op (const bf16 &x ) { return __hneg(x); } +template<> __device__ inline bf16_2 neg::op (const bf16_2 &x ) { return __hneg2(x); } +template<> __device__ inline half neg::op (const half &x ) { return __hneg(x); } +template<> __device__ inline half_2 neg::op (const half_2 &x ) { return __hneg2(x); } + + + + + + + +struct sqrt { + template static __device__ inline T op(const T &x) { return sqrt(x); } +}; +template<> __device__ inline float sqrt::op (const float &x ) { return sqrtf(x); } +template<> __device__ inline float2 sqrt::op (const float2 &x ) { return float2{sqrtf(x.x), sqrtf(x.y)}; } +template<> __device__ inline bf16 sqrt::op (const bf16 &x ) { return hsqrt(x); } +template<> __device__ inline bf16_2 sqrt::op (const bf16_2 &x ) { return h2sqrt(x); } +template<> __device__ inline half sqrt::op (const half &x ) { return hsqrt(x); } +template<> __device__ inline half_2 sqrt::op (const half_2 &x ) { return h2sqrt(x); } + + + + + + + +struct rsqrt { + template static __device__ inline T op(const T &x) { return rsqrt(x); } +}; +template<> __device__ inline float rsqrt::op (const float &x ) { return rsqrtf(x); } +template<> __device__ inline float2 rsqrt::op (const float2 &x ) { return float2{rsqrtf(x.x), rsqrtf(x.y)}; } +template<> __device__ inline bf16 rsqrt::op (const bf16 &x ) { return hrsqrt(x); } +template<> __device__ inline bf16_2 rsqrt::op (const bf16_2 &x ) { return h2rsqrt(x); } +template<> __device__ inline half rsqrt::op (const half &x ) { return hrsqrt(x); } +template<> __device__ inline half_2 rsqrt::op (const half_2 &x ) { return h2rsqrt(x); } + + + + + + + + + +struct copy { + template static __device__ inline T op(const T &a) { return a; } +}; + + + + + + + + + + + + + + +struct copy2 { + template static __device__ inline T op(const T &a, const T &b) { return b; } +}; + + + + + + + + + + +struct sum { + template static __device__ inline T op(const T &a, const T &b) { return a+b; } +}; +template<> __device__ inline float2 sum::op(const float2 &a, const float2 &b) { +#ifdef KITTENS_SM10X + float2 c; + asm volatile("add.f32x2 %0, %1, %2;" : "=l"(*(uint64_t*)&c) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b)); + return c; +#else + return float2{a.x+b.x, a.y+b.y}; +#endif +} +template<> __device__ inline bf16 sum::op (const bf16 &a, const bf16 &b) { return __hadd(a, b); } +template<> __device__ inline bf16_2 sum::op(const bf16_2 &a, const bf16_2 &b) { return __hadd2(a, b); } +template<> __device__ inline half sum::op (const half &a, const half &b) { return __hadd(a, b); } +template<> __device__ inline half_2 sum::op(const half_2 &a, const half_2 &b) { return __hadd2(a, b); } + + + + + + + + + + +struct sub { + template static __device__ inline T op(const T &a, const T &b) { return a-b; } +}; +template<> __device__ inline float2 sub::op(const float2 &a, const float2 &b) { +#ifdef KITTENS_SM10X + float2 c; + asm volatile("sub.f32x2 %0, %1, %2;" : "=l"(*(uint64_t*)&c) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b)); + return c; +#else + return float2{a.x-b.x, a.y-b.y}; +#endif +} +template<> __device__ inline bf16 sub::op (const bf16 &a, const bf16 &b) { return __hsub(a, b); } +template<> __device__ inline bf16_2 sub::op(const bf16_2 &a, const bf16_2 &b) { return __hsub2(a, b); } +template<> __device__ inline half sub::op (const half &a, const half &b) { return __hsub(a, b); } +template<> __device__ inline half_2 sub::op(const half_2 &a, const half_2 &b) { return __hsub2(a, b); } + + + + + + + + + + +struct mul { + template static __device__ inline T op(const T &a, const T &b) { return a*b; } +}; +template<> __device__ inline float2 mul::op(const float2 &a, const float2 &b) { +#ifdef KITTENS_SM10X + float2 c; + asm volatile("mul.f32x2 %0, %1, %2;" : "=l"(*(uint64_t*)&c) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b)); + return c; +#else + return float2{a.x*b.x, a.y*b.y}; +#endif +} +template<> __device__ inline bf16 mul::op (const bf16 &a, const bf16 &b) { return __hmul(a, b); } +template<> __device__ inline bf16_2 mul::op(const bf16_2 &a, const bf16_2 &b) { return __hmul2(a, b); } +template<> __device__ inline half mul::op (const half &a, const half &b) { return __hmul(a, b); } +template<> __device__ inline half_2 mul::op(const half_2 &a, const half_2 &b) { return __hmul2(a, b); } + + + + + + + + + + +struct div { + template static __device__ inline T op(const T &a, const T &b) { return a/b; } +}; +template<> __device__ inline float2 div::op(const float2 &a, const float2 &b) { return float2{a.x/b.x, a.y/b.y}; } +template<> __device__ inline bf16 div::op (const bf16 &a, const bf16 &b) { return __hdiv(a, b); } +template<> __device__ inline bf16_2 div::op(const bf16_2 &a, const bf16_2 &b) { return __h2div(a, b); } +template<> __device__ inline half div::op (const half &a, const half &b) { return __hdiv(a, b); } +template<> __device__ inline half_2 div::op(const half_2 &a, const half_2 &b) { return __h2div(a, b); } + + + + + + + + + + + struct max { + template static __device__ inline T op(const T &a, const T &b) { return ::max(a, b); } +}; +template<> __device__ inline float2 max::op(const float2 &a, const float2 &b) { return float2{::max(a.x, b.x), ::max(a.y, b.y)}; } +template<> __device__ inline bf16 max::op (const bf16 &a, const bf16 &b) { return __hmax(a, b); } +template<> __device__ inline bf16_2 max::op(const bf16_2 &a, const bf16_2 &b) { return __hmax2(a, b); } +template<> __device__ inline half max::op (const half &a, const half &b) { return __hmax(a, b); } +template<> __device__ inline half_2 max::op(const half_2 &a, const half_2 &b) { return __hmax2(a, b); } + + + + + + + + + + +struct min { + template static __device__ inline T op(const T &a, const T &b) { return ::min(a, b); } +}; +template<> __device__ inline float2 min::op(const float2 &a, const float2 &b) { return float2{::min(a.x, b.x), ::min(a.y, b.y)}; } +template<> __device__ inline bf16 min::op (const bf16 &a, const bf16 &b) { return __hmin(a, b); } +template<> __device__ inline bf16_2 min::op(const bf16_2 &a, const bf16_2 &b) { return __hmin2(a, b); } +template<> __device__ inline half min::op (const half &a, const half &b) { return __hmin(a, b); } +template<> __device__ inline half_2 min::op(const half_2 &a, const half_2 &b) { return __hmin2(a, b); } + + + + + + + + + + + + + + + +struct fma_AxBtC { + template static __device__ inline T op(const T &a, const T &b, const T &c) { + return sum::op(mul::op(a, b), c); + } +}; +template<> __device__ inline float2 fma_AxBtC::op(const float2 &a, const float2 &b, const float2 &c) { +#ifdef KITTENS_SM10X + float2 d; + asm volatile("fma.rn.f32x2 %0, %1, %2, %3;" : "=l"(*(uint64_t*)&d) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b), "l"(*(uint64_t*)&c)); + return d; +#else + return float2{a.x*b.x+c.x, a.y*b.y+c.y}; +#endif +} + + + + + + + + + + + + +struct fma_AxCtB { + template static __device__ inline T op(const T &a, const T &b, const T &c) { + return sum::op(mul::op(a, c), b); + } +}; +template<> __device__ inline float2 fma_AxCtB::op(const float2 &a, const float2 &b, const float2 &c) { +#ifdef KITTENS_SM10X + float2 d; + asm volatile("fma.rn.f32x2 %0, %1, %2, %3;" : "=l"(*(uint64_t*)&d) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&c), "l"(*(uint64_t*)&b)); + return d; +#else + return float2{a.x*c.x+b.x, a.y*c.y+b.y}; +#endif +} + +} + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/common/base_types.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/common/base_types.cuh new file mode 100644 index 0000000..3a8efbb --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/common/base_types.cuh @@ -0,0 +1,777 @@ + + + + + + + + + + +#pragma once + +namespace kittens { + + + + +using bf16 = __nv_bfloat16; + + + +using half = __half; + + + +using int8 = signed char; + + + +using uint8 = unsigned char; + + + + +using bf16_2 = __nv_bfloat162; + + + +using half_2 = __half2; + + + + +using int8_2 = char2; +using int8_4 = char4; + + + +using uint8_2 = uchar2; +using uint8_4 = uchar4; + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + +using fp8e4m3 = __nv_fp8_e4m3; +using fp8e5m2 = __nv_fp8_e5m2; +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) +using fp8e8m0 = __nv_fp8_e8m0; +#endif + + + +using fp8e4m3_2 = __nv_fp8x2_e4m3; +using fp8e5m2_2 = __nv_fp8x2_e5m2; +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) +using fp8e8m0_2 = __nv_fp8x2_e8m0; +#endif + + + +using fp8e4m3_4 = __nv_fp8x4_e4m3; +using fp8e5m2_4 = __nv_fp8x4_e5m2; +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) +using fp8e8m0_4 = __nv_fp8x4_e8m0; +#endif +#endif + +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + +using fp4e2m1 = __nv_fp4_e2m1; + + + +using fp4e2m1_2 = __nv_fp4x2_e2m1; + + + +using fp4e2m1_4 = __nv_fp4x4_e2m1; +#endif + +namespace ducks { + + + + + +namespace base_types { + +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template +concept T2 = std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v; +template +concept T1 = std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v; +#elif defined(KITTENS_SM90) +template +concept T2 = std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v; +template +concept T1 = std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v; +#else +template +concept T2 = std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v; +template +concept T1 = std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v; +#endif + +} +} + + + + + + +namespace base_types { + + + + + + +template struct constants { + + + + + static __device__ inline constexpr T zero() { return T{0}; } + + + + + static __device__ inline constexpr T one() { return T{1}; } + + + + + static __device__ inline constexpr T pos_infty() { static_assert(sizeof(T) == 9999, "pos_infty not defined for this type"); return T{}; } + + + + + static __device__ inline constexpr T neg_infty() { static_assert(sizeof(T) == 9999, "neg_infty not defined for this type"); return T{}; } +}; +template<> struct constants { + static __device__ inline constexpr float zero() { return 0.f; } + static __device__ inline constexpr float one() { return 1.f; } + static __device__ inline constexpr float pos_infty() { return std::bit_cast(uint32_t(0x7f800000)); } + static __device__ inline constexpr float neg_infty() { return std::bit_cast(uint32_t(0xff800000)); } +}; +template<> struct constants { + static __device__ inline constexpr float2 zero() { return float2{0.f, 0.f}; } + static __device__ inline constexpr float2 one() { return float2{1.f, 1.f}; } + static __device__ inline constexpr float2 pos_infty() { return float2{constants::pos_infty(), constants::pos_infty()}; } + static __device__ inline constexpr float2 neg_infty() { return float2{constants::neg_infty(), constants::neg_infty()}; } +}; +template<> struct constants { + static __device__ inline constexpr bf16 zero() { return std::bit_cast<__nv_bfloat16>(uint16_t(0x0000)); } + static __device__ inline constexpr bf16 one() { return std::bit_cast<__nv_bfloat16>(uint16_t(0x3F80)); } + static __device__ inline constexpr bf16 pos_infty() { return std::bit_cast<__nv_bfloat16>(uint16_t(0x7F80)); } + static __device__ inline constexpr bf16 neg_infty() { return std::bit_cast<__nv_bfloat16>(uint16_t(0xFF80)); } +}; +template<> struct constants { + static __device__ inline constexpr bf16_2 zero() { return bf16_2{constants::zero(), constants::zero()}; } + static __device__ inline constexpr bf16_2 one() { return bf16_2{constants::one(), constants::one()}; } + static __device__ inline constexpr bf16_2 pos_infty() { return bf16_2{constants::pos_infty(), constants::pos_infty()}; } + static __device__ inline constexpr bf16_2 neg_infty() { return bf16_2{constants::neg_infty(), constants::neg_infty()}; } +}; +template<> struct constants { + static __device__ inline constexpr half zero() { return std::bit_cast<__half>(uint16_t(0x0000)); } + static __device__ inline constexpr half one() { return std::bit_cast<__half>(uint16_t(0x3C00)); } + static __device__ inline constexpr half pos_infty() { return std::bit_cast<__half>(uint16_t(0x7C00)); } + static __device__ inline constexpr half neg_infty() { return std::bit_cast<__half>(uint16_t(0xFC00)); } +}; +template<> struct constants { + static __device__ inline constexpr half_2 zero() { return half_2{constants::zero(), constants::zero()}; } + static __device__ inline constexpr half_2 one() { return half_2{constants::one(), constants::one()}; } + static __device__ inline constexpr half_2 pos_infty() { return half_2{constants::pos_infty(), constants::pos_infty()}; } + static __device__ inline constexpr half_2 neg_infty() { return half_2{constants::neg_infty(), constants::neg_infty()}; } +}; +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template<> struct constants { + static __device__ inline constexpr fp8e4m3 zero() { return std::bit_cast<__nv_fp8_e4m3>(uint8_t(0x00)); } + static __device__ inline constexpr fp8e4m3 one() { return std::bit_cast<__nv_fp8_e4m3>(uint8_t(0x38)); } +}; +template<> struct constants { + static __device__ inline constexpr fp8e4m3_2 zero() { return std::bit_cast(uint16_t(0x0000)); } + static __device__ inline constexpr fp8e4m3_2 one() { return std::bit_cast(uint16_t(0x3838)); } +}; +template<> struct constants { + static __device__ inline constexpr fp8e4m3_4 zero() { return std::bit_cast(uint32_t(0x00000000)); } + static __device__ inline constexpr fp8e4m3_4 one() { return std::bit_cast(uint32_t(0x38383838)); } +}; +template<> struct constants { + static __device__ inline constexpr fp8e5m2 zero() { return std::bit_cast<__nv_fp8_e5m2>(uint8_t(0x00)); } + static __device__ inline constexpr fp8e5m2 one() { return std::bit_cast<__nv_fp8_e5m2>(uint8_t(0x3C)); } +}; +template<> struct constants { + static __device__ inline constexpr fp8e5m2_2 zero() { return std::bit_cast(uint16_t(0x0000)); } + static __device__ inline constexpr fp8e5m2_2 one() { return std::bit_cast(uint16_t(0x3C3C)); } +}; +template<> struct constants { + static __device__ inline constexpr fp8e5m2_4 zero() { return std::bit_cast(uint32_t(0x00000000)); } + static __device__ inline constexpr fp8e5m2_4 one() { return std::bit_cast(uint32_t(0x3C3C3C3C)); } +}; +#endif +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template<> struct constants { + static __device__ inline constexpr fp4e2m1 zero() { return std::bit_cast(uint8_t(0x00)); } + static __device__ inline constexpr fp4e2m1 one() { return std::bit_cast(uint8_t(0x02)); } +}; +template<> struct constants { + static __device__ inline constexpr fp4e2m1_2 zero() { return std::bit_cast(uint8_t(0x00)); } + static __device__ inline constexpr fp4e2m1_2 one() { return std::bit_cast(uint8_t(0x22)); } +}; +template<> struct constants { + static __device__ inline constexpr fp4e2m1_4 zero() { return std::bit_cast(uint16_t(0x0000)); } + static __device__ inline constexpr fp4e2m1_4 one() { return std::bit_cast(uint16_t(0x2222)); } +}; +template<> struct constants { + static __device__ inline constexpr fp8e8m0 zero() { return std::bit_cast(uint8_t(0x00)); } + static __device__ inline constexpr fp8e8m0 one() { return std::bit_cast(uint8_t(0x7F)); } +}; +template<> struct constants { + static __device__ inline constexpr fp8e8m0_2 zero() { return std::bit_cast(uint16_t(0x0000)); } + static __device__ inline constexpr fp8e8m0_2 one() { return std::bit_cast(uint16_t(0x7F7F)); } +}; +template<> struct constants { + static __device__ inline constexpr fp8e8m0_4 zero() { return std::bit_cast(uint32_t(0x00000000)); } + static __device__ inline constexpr fp8e8m0_4 one() { return std::bit_cast(uint32_t(0x7F7F7F7F)); } +}; +#endif + +template<> struct constants { + static __device__ inline constexpr int zero() { return 0; } + static __device__ inline constexpr int one() { return 1; } +}; +template<> struct constants { + static __device__ inline constexpr int2 zero() { return int2{0, 0}; } + static __device__ inline constexpr int2 one() { return int2{1, 1}; } +}; + + + + + + +template struct packing { + + + + + + static __device__ inline constexpr int num() { return 1; } + + + + + + + static __device__ inline constexpr T pack(const bf16 &i); +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = bf16; + using packed_type = bf16_2; + static __device__ inline constexpr bf16_2 pack(const bf16 &i) { return bf16_2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = bf16; + using packed_type = bf16_2; + static __device__ inline constexpr bf16_2 pack(const bf16 &i) { return bf16_2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = half; + using packed_type = half_2; + static __device__ inline constexpr half_2 pack(const half &i) { return half_2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = half; + using packed_type = half_2; + static __device__ inline constexpr half_2 pack(const half &i) { return half_2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = float; + using packed_type = float2; + static __device__ inline constexpr float2 pack(const float &i) { return float2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = float; + using packed_type = float2; + static __device__ inline constexpr float2 pack(const float &i) { return float2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = double; + using packed_type = double2; + static __device__ inline constexpr double2 pack(const double &i) { return double2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = double; + using packed_type = double2; + static __device__ inline constexpr double2 pack(const double &i) { return double2{i, i}; } +}; +struct bool_2 { bool x, y; }; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = bool; + using packed_type = bool_2; + static __device__ inline constexpr bool_2 pack(const bool &i) { return bool_2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = bool; + using packed_type = bool_2; + static __device__ inline constexpr bool_2 pack(const bool &i) { return bool_2{i, i}; } +}; + +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = int8; + using packed_type = int8_4; + static __device__ inline constexpr int8_4 pack(const int8 &i) { return int8_4{i, i, i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 4; } + using unpacked_type = int8; + using packed_type = int8_4; + static __device__ inline constexpr int8_4 pack(const int8 &i) { return int8_4{i, i, i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = uint8; + using packed_type = uint8_4; + static __device__ inline constexpr uint8_4 pack(const uint8 &i) { return uint8_4{i, i, i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 4; } + using unpacked_type = uint8; + using packed_type = uint8_4; + static __device__ inline constexpr uint8_4 pack(const uint8 &i) { return uint8_4{i, i, i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = short; + using packed_type = short2; + static __device__ inline constexpr short2 pack(const short &i) { return short2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = short; + using packed_type = short2; + static __device__ inline constexpr short2 pack(const short &i) { return short2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = unsigned short; + using packed_type = ushort2; + static __device__ inline constexpr ushort2 pack(const unsigned short &i) { return ushort2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = unsigned short; + using packed_type = ushort2; + static __device__ inline constexpr ushort2 pack(const unsigned short &i) { return ushort2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = int; + using packed_type = int2; + static __device__ inline constexpr int2 pack(const int &i) { return int2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = int; + using packed_type = int2; + static __device__ inline constexpr int2 pack(const int &i) { return int2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = uint; + using packed_type = uint2; + static __device__ inline constexpr uint2 pack(const uint &i) { return uint2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = uint; + using packed_type = uint2; + static __device__ inline constexpr uint2 pack(const uint &i) { return uint2{i, i}; } +}; +struct int64_2 { int64_t x, y; }; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = int64_t; + using packed_type = int64_2; + static __device__ inline constexpr int64_2 pack(const int64_t &i) { return int64_2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = int64_t; + using packed_type = int64_2; + static __device__ inline constexpr int64_2 pack(const int64_t &i) { return int64_2{i, i}; } +}; +struct uint64_2 { uint64_t x, y; }; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = uint64_t; + using packed_type = uint64_2; + static __device__ inline constexpr uint64_2 pack(const uint64_t &i) { return uint64_2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = uint64_t; + using packed_type = uint64_2; + static __device__ inline constexpr uint64_2 pack(const uint64_t &i) { return uint64_2{i, i}; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 4; } +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 4; } +}; +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = fp8e4m3; + using packed_type = fp8e4m3_4; +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 4; } + using unpacked_type = fp8e4m3; + using packed_type = fp8e4m3_4; +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = fp8e5m2; + using packed_type = fp8e5m2_4; +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 4; } + using unpacked_type = fp8e5m2; + using packed_type = fp8e5m2_4; +}; +#endif +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template<> struct packing { + static __device__ inline constexpr int num() { return 1; } + using unpacked_type = fp8e8m0; + using packed_type = fp8e8m0_4; +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 4; } + using unpacked_type = fp8e8m0; + using packed_type = fp8e8m0_4; +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 2; } + using unpacked_type = fp4e2m1_2; + using packed_type = fp4e2m1_4; +}; +template<> struct packing { + static __device__ inline constexpr int num() { return 4; } + using unpacked_type = fp4e2m1_2; + using packed_type = fp4e2m1_4; +}; +#endif + + + + + + + +template struct convertor { + + + + + + + static __host__ __device__ inline T convert(const U & u) { + return (T)u; + } +}; +template<> struct convertor { + static __host__ __device__ inline float convert(const bf16 & u) { + return __bfloat162float(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline bf16 convert(const float & u) { + return __float2bfloat16_rn(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float2 convert(const bf16_2 & u) { + return __bfloat1622float2(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline bf16_2 convert(const float2 & u) { + return __float22bfloat162_rn(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float convert(const half & u) { + return __half2float(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline half convert(const float & u) { + return __float2half(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float2 convert(const half_2 & u) { + return __half22float2(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline half_2 convert(const float2 & u) { + return __float22half2_rn(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline bf16 convert(const half & u) { + return __float2bfloat16_rn(__half2float(u)); + } +}; +template<> struct convertor { + static __host__ __device__ inline half convert(const bf16 & u) { + return __float2half(__bfloat162float(u)); + } +}; +template<> struct convertor { + static __host__ __device__ inline bf16_2 convert(const half_2 & u) { + return __float22bfloat162_rn(__half22float2(u)); + } +}; +template<> struct convertor { + static __host__ __device__ inline half_2 convert(const bf16_2 & u) { + return __float22half2_rn(__bfloat1622float2(u)); + } +}; +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + +template<> struct convertor { + static __host__ __device__ inline fp8e8m0_4 convert(const float4& u) { + return __nv_fp8x4_e8m0(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float4 convert(const fp8e8m0_4& u) { + __nv_fp8_e8m0 *vals = reinterpret_cast<__nv_fp8_e8m0*>(const_cast<__nv_fp8x4_e8m0*>(&u)); + return make_float4(float(vals[0]), float(vals[1]), float(vals[2]), float(vals[3])); + } +}; +template<> struct convertor { + static __host__ __device__ inline fp8e8m0_2 convert(const float2& u) { + return __nv_fp8x2_e8m0(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float2 convert(const fp8e8m0_2& u) { + __nv_fp8_e8m0 *vals = reinterpret_cast<__nv_fp8_e8m0*>(const_cast<__nv_fp8x2_e8m0*>(&u)); + return make_float2(float(vals[0]), float(vals[1])); + } +}; +template<> struct convertor { + static __host__ __device__ inline fp8e8m0 convert(const float & u) { + return __nv_fp8_e8m0(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float convert(const fp8e8m0 & u) { + return float(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline bf16_2 convert(const fp8e8m0_4 & u) { + float4 f4 = convertor::convert(u); + float2 f2 = make_float2(f4.x, f4.y); + return __float22bfloat162_rn(f2); + } +}; +template<> struct convertor { + static __host__ __device__ inline fp8e8m0_4 convert(const bf16_2 & u) { + float2 f2 = __bfloat1622float2(u); + float4 f4 = make_float4(f2.x, f2.y, 0.0f, 0.0f); + return __nv_fp8x4_e8m0(f4); + } +}; + +template<> struct convertor { + static __device__ inline fp4e2m1 convert(const float & u) { + __nv_fp4_storage_t storage = __nv_cvt_float_to_fp4(u, __NV_E2M1, cudaRoundNearest); + return std::bit_cast(storage); + } +}; +template<> struct convertor { + static __device__ inline float convert(const fp4e2m1 & u) { + __nv_fp4_storage_t storage = std::bit_cast<__nv_fp4_storage_t>(u); + __half_raw hr = __nv_cvt_fp4_to_halfraw(storage, __NV_E2M1); + return __half2float(*reinterpret_cast<__half*>(&hr)); + } +}; +template<> struct convertor { + static __device__ inline fp4e2m1_2 convert(const float2 & u) { + __nv_fp4x2_storage_t storage = __nv_cvt_float2_to_fp4x2(u, __NV_E2M1, cudaRoundNearest); + return std::bit_cast(storage); + } +}; +template<> struct convertor { + static __device__ inline float2 convert(const fp4e2m1_2 & u) { + __nv_fp4x2_storage_t storage = std::bit_cast<__nv_fp4x2_storage_t>(u); + __half2_raw hr2 = __nv_cvt_fp4x2_to_halfraw2(storage, __NV_E2M1); + return __half22float2(*reinterpret_cast<__half2*>(&hr2)); + } +}; +#endif +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + +template<> struct convertor { + static __host__ __device__ inline fp8e4m3_4 convert(const float4& u) { + return __nv_fp8x4_e4m3(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float4 convert(const fp8e4m3_4& u) { + __nv_fp8_e4m3 *vals = reinterpret_cast<__nv_fp8_e4m3*>(const_cast<__nv_fp8x4_e4m3*>(&u)); + return make_float4(float(vals[0]), float(vals[1]), float(vals[2]), float(vals[3])); + } +}; +template<> struct convertor { + static __host__ __device__ inline fp8e4m3_2 convert(const float2& u) { + return __nv_fp8x2_e4m3(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float2 convert(const fp8e4m3_2& u) { + __nv_fp8_e4m3 *vals = reinterpret_cast<__nv_fp8_e4m3*>(const_cast<__nv_fp8x2_e4m3*>(&u)); + return make_float2(float(vals[0]), float(vals[1])); + } +}; +template<> struct convertor { + static __host__ __device__ inline fp8e4m3 convert(const float & u) { + return __nv_fp8_e4m3(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float convert(const fp8e4m3 & u) { + return float(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline bf16_2 convert(const fp8e4m3_4 & u) { + float4 f4 = convertor::convert(u); + float2 f2 = make_float2(f4.x, f4.y); + return __float22bfloat162_rn(f2); + } +}; +template<> struct convertor { + static __host__ __device__ inline fp8e4m3_4 convert(const bf16_2 & u) { + float2 f2 = __bfloat1622float2(u); + float4 f4 = make_float4(f2.x, f2.y, 0.0f, 0.0f); + return __nv_fp8x4_e4m3(f4); + } +}; + +template<> struct convertor { + static __host__ __device__ inline fp8e5m2_4 convert(const float4& u) { + return __nv_fp8x4_e5m2(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float4 convert(const fp8e5m2_4& u) { + __nv_fp8_e5m2 *vals = reinterpret_cast<__nv_fp8_e5m2*>(const_cast<__nv_fp8x4_e5m2*>(&u)); + return make_float4(float(vals[0]), float(vals[1]), float(vals[2]), float(vals[3])); + } +}; +template<> struct convertor { + static __host__ __device__ inline fp8e5m2_2 convert(const float2& u) { + return __nv_fp8x2_e5m2(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float2 convert(const fp8e5m2_2& u) { + __nv_fp8_e5m2 *vals = reinterpret_cast<__nv_fp8_e5m2*>(const_cast<__nv_fp8x2_e5m2*>(&u)); + return make_float2(float(vals[0]), float(vals[1])); + } +}; +template<> struct convertor { + static __host__ __device__ inline fp8e5m2 convert(const float & u) { + return __nv_fp8_e5m2(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float convert(const fp8e5m2 & u) { + return float(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline bf16_2 convert(const fp8e5m2_4 & u) { + float4 f4 = convertor::convert(u); + float2 f2 = make_float2(f4.x, f4.y); + return __float22bfloat162_rn(f2); + } +}; +template<> struct convertor { + static __host__ __device__ inline fp8e5m2_4 convert(const bf16_2 & u) { + float2 f2 = __bfloat1622float2(u); + float4 f4 = make_float4(f2.x, f2.y, 0.0f, 0.0f); + return __nv_fp8x4_e5m2(f4); + } +}; +#endif +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/common/common.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/common/common.cuh new file mode 100644 index 0000000..855bb69 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/common/common.cuh @@ -0,0 +1,15 @@ + + + + + + +#pragma once + +#include "base_types.cuh" +#include "base_ops.cuh" +#include "util.cuh" + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +#include "multimem.cuh" +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/common/multimem.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/common/multimem.cuh new file mode 100644 index 0000000..ac5f226 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/common/multimem.cuh @@ -0,0 +1,407 @@ + + + + + +#pragma once + +#include "base_types.cuh" + +namespace kittens { + +enum class reduce_op { + ADD = 0, + MIN = 1, + MAX = 2 +}; + +enum class memory_model { + WEAK = 0, + STRONG = 1 +}; + +template +struct multimem; + +template <> +struct multimem { + template + __device__ static inline void ld_reduce(int &dst, const int *src) { + if constexpr (Op == reduce_op::ADD) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.add.s32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.add.s32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MIN) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.min.s32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.min.s32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MAX) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.max.s32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.max.s32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } + } + } + template + __device__ static inline void st(int *dst, const int &src) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.st.weak.global.s32 [%0], %1;" + :: "l"(dst), "r"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.st.release.sys.global.s32 [%0], %1;" + :: "l"(dst), "r"(src) : "memory"); + } + } + template + __device__ static inline void red(int *dst, const int &src) { + if constexpr (Op == reduce_op::ADD) { + asm volatile("multimem.red.release.sys.global.add.s32 [%0], %1;" + : : "l"(dst), "r"(src) : "memory"); + } else if constexpr (Op == reduce_op::MIN) { + asm volatile("multimem.red.release.sys.global.min.s32 [%0], %1;" + : : "l"(dst), "r"(src) : "memory"); + } else if constexpr (Op == reduce_op::MAX) { + asm volatile("multimem.red.release.sys.global.max.s32 [%0], %1;" + : : "l"(dst), "r"(src) : "memory"); + } + } +}; + +template <> +struct multimem { + template + __device__ static inline void ld_reduce(uint &dst, const uint *src) { + if constexpr (Op == reduce_op::ADD) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.add.u32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.add.u32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MIN) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.min.u32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.min.u32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MAX) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.max.u32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.max.u32 %0, [%1];" + : "=r"(dst) : "l"(src) : "memory"); + } + } + } + template + __device__ static inline void st(uint *dst, const uint &src) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.st.weak.global.u32 [%0], %1;" + :: "l"(dst), "r"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.st.release.sys.global.u32 [%0], %1;" + :: "l"(dst), "r"(src) : "memory"); + } + } + template + __device__ static inline void red(uint *dst, const uint &src) { + if constexpr (Op == reduce_op::ADD) { + asm volatile("multimem.red.release.sys.global.add.u32 [%0], %1;" + : : "l"(dst), "r"(src) : "memory"); + } else if constexpr (Op == reduce_op::MIN) { + asm volatile("multimem.red.release.sys.global.min.u32 [%0], %1;" + : : "l"(dst), "r"(src) : "memory"); + } else if constexpr (Op == reduce_op::MAX) { + asm volatile("multimem.red.release.sys.global.max.u32 [%0], %1;" + : : "l"(dst), "r"(src) : "memory"); + } + } +}; + +template <> +struct multimem { + template + __device__ static inline void ld_reduce(float &dst, const float *src) { + static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for f32 ld_reduce operations"); + if constexpr (Op == reduce_op::ADD) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.add.f32 %0, [%1];" + : "=f"(dst) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.add.f32 %0, [%1];" + : "=f"(dst) : "l"(src) : "memory"); + } + } + } + template + __device__ static inline void st(float *dst, const float &src) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.st.weak.global.f32 [%0], %1;" + :: "l"(dst), "f"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.st.release.sys.global.f32 [%0], %1;" + :: "l"(dst), "f"(src) : "memory"); + } + } + template + __device__ static inline void red(float *dst, const float &src) { + static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for f32 red operations"); + if constexpr (Op == reduce_op::ADD) { + asm volatile("multimem.red.release.sys.global.add.f32 [%0], %1;" + : : "l"(dst), "f"(src) : "memory"); + } + } +}; + + +template <> +struct multimem { + template + __device__ static inline void ld_reduce(float2 &dst, const float2 *src) { + static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for f32 ld_reduce operations"); + if constexpr (Op == reduce_op::ADD) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.add.v2.f32 {%0, %1}, [%2];" + : "=f"(dst.x), "=f"(dst.y) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.add.v2.f32 {%0, %1}, [%2];" + : "=f"(dst.x), "=f"(dst.y) : "l"(src) : "memory"); + } + } + } + template + __device__ static inline void st(float2 *dst, const float2 &src) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.st.weak.global.v2.f32 [%0], {%1, %2};" + :: "l"(dst), "f"(src.x), "f"(src.y) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.st.release.sys.global.v2.f32 [%0], {%1, %2};" + :: "l"(dst), "f"(src.x), "f"(src.y) : "memory"); + } + } + template + __device__ static inline void red(float2 *dst, const float2 &src) { + static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for f32 red operations"); + if constexpr (Op == reduce_op::ADD) { + asm volatile("multimem.red.release.sys.global.add.v2.f32 [%0], {%1, %2};" + : : "l"(dst), "f"(src.x), "f"(src.y) : "memory"); + } + } +}; + +template <> +struct multimem { + template + __device__ static inline void ld_reduce(bf16 &dst, const bf16 *src) { + if constexpr (Op == reduce_op::ADD) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.add.acc::f32.bf16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.add.acc::f32.bf16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MIN) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.min.bf16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.min.bf16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MAX) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.max.bf16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.max.bf16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } + } + template + __device__ static inline void st(bf16 *dst, const bf16 &src) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.st.weak.global.bf16 [%0], %1;" + :: "l"(dst), "h"(*reinterpret_cast(&src)) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.st.release.sys.global.bf16 [%0], %1;" + :: "l"(dst), "h"(*reinterpret_cast(&src)) : "memory"); + } + } + template + __device__ static inline void red(bf16 *dst, const bf16 &src) { + static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for bf16 red operations"); + if constexpr (Op == reduce_op::ADD) { + asm volatile("multimem.red.release.sys.global.add.bf16 [%0], %1;" + : : "l"(dst), "h"(*reinterpret_cast(&src)) : "memory"); + } + } +}; + +template <> +struct multimem { + template + __device__ static inline void ld_reduce(bf16_2 &dst, const bf16_2 *src) { + if constexpr (Op == reduce_op::ADD) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.add.acc::f32.bf16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.add.acc::f32.bf16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MIN) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.min.bf16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.min.bf16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MAX) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.max.bf16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.max.bf16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } + } + template + __device__ static inline void st(bf16_2 *dst, const bf16_2 &src) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.st.weak.global.bf16x2 [%0], %1;" + :: "l"(dst), "r"(*reinterpret_cast(&src)) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.st.release.sys.global.bf16x2 [%0], %1;" + :: "l"(dst), "r"(*reinterpret_cast(&src)) : "memory"); + } + } + template + __device__ static inline void red(bf16_2 *dst, const bf16_2 &src) { + static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for bf16_2 red operations"); + if constexpr (Op == reduce_op::ADD) { + asm volatile("multimem.red.release.sys.global.add.bf16x2 [%0], %1;" + : : "l"(dst), "r"(*reinterpret_cast(&src)) : "memory"); + } + } +}; + +template <> +struct multimem { + template + __device__ static inline void ld_reduce(half &dst, const half *src) { + if constexpr (Op == reduce_op::ADD) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.add.acc::f32.f16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.add.acc::f32.f16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MIN) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.min.f16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.min.f16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MAX) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.max.f16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.max.f16 %0, [%1];" + : "=h"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } + } + template + __device__ static inline void st(half *dst, const half &src) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.st.weak.global.f16 [%0], %1;" + :: "l"(dst), "h"(*reinterpret_cast(&src)) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.st.release.sys.global.f16 [%0], %1;" + :: "l"(dst), "h"(*reinterpret_cast(&src)) : "memory"); + } + } + template + __device__ static inline void red(half *dst, const half &src) { + static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for f16 red operations"); + if constexpr (Op == reduce_op::ADD) { + asm volatile("multimem.red.release.sys.global.add.f16 [%0], %1;" + : : "l"(dst), "h"(*reinterpret_cast(&src)) : "memory"); + } + } +}; + +template <> +struct multimem { + template + __device__ static inline void ld_reduce(half_2 &dst, const half_2 *src) { + if constexpr (Op == reduce_op::ADD) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.add.acc::f32.f16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.add.acc::f32.f16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MIN) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.min.f16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.min.f16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } else if constexpr (Op == reduce_op::MAX) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.ld_reduce.weak.global.max.f16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.ld_reduce.acquire.sys.global.max.f16x2 %0, [%1];" + : "=r"(*reinterpret_cast(&dst)) : "l"(src) : "memory"); + } + } + } + template + __device__ static inline void st(half_2 *dst, const half_2 &src) { + if constexpr (M == memory_model::WEAK) { + asm volatile("multimem.st.weak.global.f16x2 [%0], %1;" + :: "l"(dst), "r"(*reinterpret_cast(&src)) : "memory"); + } else if constexpr (M == memory_model::STRONG) { + asm volatile("multimem.st.release.sys.global.f16x2 [%0], %1;" + :: "l"(dst), "r"(*reinterpret_cast(&src)) : "memory"); + } + } + template + __device__ static inline void red(half_2 *dst, const half_2 &src) { + static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for f16_2 red operations"); + if constexpr (Op == reduce_op::ADD) { + asm volatile("multimem.red.release.sys.global.add.f16x2 [%0], %1;" + : : "l"(dst), "r"(*reinterpret_cast(&src)) : "memory"); + } + } +}; + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/common/util.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/common/util.cuh new file mode 100644 index 0000000..ac666b3 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/common/util.cuh @@ -0,0 +1,551 @@ + + + + + +#pragma once + + + + + + +namespace kittens { + + + + + + +constexpr int BASE_TILE_DIM = 16; +template constexpr int TILE_COL_DIM = sizeof(T) == 1 ? BASE_TILE_DIM * 2 : BASE_TILE_DIM; +template constexpr int TILE_ROW_DIM = BASE_TILE_DIM; + + + +template constexpr int TILE_ELEMENTS{TILE_COL_DIM*TILE_ROW_DIM}; + + + +constexpr int WARP_THREADS{32}; + + + +constexpr int WARPGROUP_THREADS{128}; + + + + +constexpr int WARPGROUP_WARPS{4}; + + + + + +__device__ __forceinline__ int warpid() { return threadIdx.x >> 5; } + + + + +__device__ __forceinline__ int warpgroupid() { return threadIdx.x >> 7; } + + + + +__device__ __forceinline__ int laneid() { return threadIdx.x & 0x1f; } + + + + +__device__ __forceinline__ int smid() { + uint32_t r; + asm volatile("mov.u32 %0, %%smid;" : "=r"(r)); + return static_cast(r); +} + +#if defined(KITTENS_SM80) +constexpr int MAX_SHARED_MEMORY = 163 * 1024; +#elif defined(KITTENS_SM90) || defined(KITTENS_SM10X) +constexpr int MAX_SHARED_MEMORY = 227 * 1024; +#elif defined(KITTENS_SM120) +constexpr int MAX_SHARED_MEMORY = 99 * 1024; +#endif + +struct transpose { + static constexpr int N = 0; + static constexpr int T = 1; +}; +struct axis { + static constexpr int ROW = 0; + static constexpr int COL = 1; +}; + + + + + + + + + + + +namespace ducks { + + + + +struct default_type {}; + + +#define typeof(A) typename std::remove_const::type>::type + +} + + + + + + +static constexpr uint32_t MASK_ALL = 0xFFFFFFFF; + + + + + + + + + +template +__device__ static inline T packed_shfl_down_sync(uint32_t mask, const T &f, int delta) { + return __shfl_down_sync(mask, f, delta); +} +template<> +__device__ inline float2 packed_shfl_down_sync(uint32_t mask, const float2 &f, int delta) { + float2 r; + r.x = __shfl_down_sync(mask, f.x, delta); + r.y = __shfl_down_sync(mask, f.y, delta); + return r; +} + + + + + + + + +template +__device__ static inline T packed_shfl_xor_sync(uint32_t mask, const T &f, int lane_mask) { + return __shfl_xor_sync(mask, f, lane_mask); +} +template<> +__device__ inline float2 packed_shfl_xor_sync(uint32_t mask, const float2 &f, int lane_mask) { + float2 r; + r.x = __shfl_xor_sync(mask, f.x, lane_mask); + r.y = __shfl_xor_sync(mask, f.y, lane_mask); + return r; +} + + + + + + + + +template +__device__ static inline T packed_shfl_sync(uint32_t mask, const T &f, int src) { + return __shfl_sync(mask, f, src); +} +template<> +__device__ inline float2 packed_shfl_sync(uint32_t mask, const float2 &f, int src) { + float2 r; + r.x = __shfl_sync(mask, f.x, src); + r.y = __shfl_sync(mask, f.y, src); + return r; +} + + + + + + + + + + + + + + + + + + + + + + + + +#if defined(__CUDACC__) +#define KITTENS_ALIGN_AS(n) __align__(n) +#else +#define KITTENS_ALIGN_AS(n) alignas(n) +#endif + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +#define KITTENS_DEFAULT_ALIGN KITTENS_ALIGN_AS(128) +#else +#define KITTENS_DEFAULT_ALIGN KITTENS_ALIGN_AS(16) +#endif + + + + + + + + +template +__host__ __device__ constexpr T cdiv(T a, T b) { + return (a + b - 1) / b; +} + + + + +struct KITTENS_DEFAULT_ALIGN alignment_dummy { int dummy; }; + + + + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template +#else +template +#endif +struct shared_allocator { + int *ptr; + +private: + + template + struct variadic_array; + template + struct variadic_array { + using type = typename variadic_array::type[first_dim]; + }; + template + struct variadic_array { + using type = A; + }; + template + using variadic_array_t = typename variadic_array::type; + + template + __device__ inline void align_ptr() { + if constexpr (alignment > 0) { + uint64_t p = reinterpret_cast(ptr); + if(p % alignment != 0) { + ptr = (int*)(p + (alignment-(p%alignment))); + } + } + } + +public: + + + + + __device__ shared_allocator(int *_ptr): ptr(_ptr) {} + + + + + + + template + __device__ inline variadic_array_t& allocate() { + + align_ptr(); + using at = variadic_array_t; + at*p = reinterpret_cast(ptr); + ptr += sizeof(at)/sizeof(int); + return *p; + } + + + + + + + + template + __device__ inline variadic_array_t& allocate() { + + align_ptr(); + using at = variadic_array_t; + at*p = reinterpret_cast(ptr); + ptr += sizeof(at)/sizeof(int); + return *p; + } +}; + +#if (defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120)) + + + +using tma_allocator = shared_allocator<1024>; +using tma_swizzle_allocator = tma_allocator; + + + + +__device__ static inline int3 clusterIdx() { + int3 cluster_idx; + asm volatile("mov.u32 %0, %clusterid.x;\n" : "=r"(cluster_idx.x)); + asm volatile("mov.u32 %0, %clusterid.y;\n" : "=r"(cluster_idx.y)); + asm volatile("mov.u32 %0, %clusterid.z;\n" : "=r"(cluster_idx.z)); + return cluster_idx; +} + + +__device__ static inline int3 clusterid() { + int3 clusterid; + asm volatile("mov.u32 %0, %clusterid.x;\n" : "=r"(clusterid.x)); + asm volatile("mov.u32 %0, %clusterid.y;\n" : "=r"(clusterid.y)); + asm volatile("mov.u32 %0, %clusterid.z;\n" : "=r"(clusterid.z)); + return clusterid; +} + + +__device__ static inline int3 nclusterid() { + int3 nclusterid; + asm volatile("mov.u32 %0, %nclusterid.x;\n" : "=r"(nclusterid.x)); + asm volatile("mov.u32 %0, %nclusterid.y;\n" : "=r"(nclusterid.y)); + asm volatile("mov.u32 %0, %nclusterid.z;\n" : "=r"(nclusterid.z)); + return nclusterid; +} + + +__device__ static inline int cluster_ctarank() { + uint32_t ctarank; + asm volatile("mov.u32 %0, %cluster_ctarank;\n" : "=r"(ctarank)); + return ctarank; +} + + +__device__ static inline int cluster_nctarank() { + uint32_t nctarank; + asm volatile("mov.u32 %0, %cluster_nctarank;\n" : "=r"(nctarank)); + return nctarank; +} + +#endif + + + + + + + + + + + + + + + + + + +template +__device__ static inline int2 get_swizzled_2d_idx(const int num_rows, const int num_cols, const int linear_idx) { + static_assert(SUPERGROUP_SIZE > 0, "SUPERGROUP_SIZE must be greater than 0"); + if constexpr (ROW_MAJOR) { + const int supergroup_numel = num_rows*SUPERGROUP_SIZE; + const int supergroup_idx = linear_idx/supergroup_numel; + const int supersection_cols = (num_cols/SUPERGROUP_SIZE)*SUPERGROUP_SIZE; + const int supersection_numel = num_rows*supersection_cols; + const int finalsection_cols = num_cols-supersection_cols; + int row_idx, col_idx; + if (linear_idx < supersection_numel) { + row_idx = (linear_idx%supergroup_numel)/SUPERGROUP_SIZE; + col_idx = supergroup_idx*SUPERGROUP_SIZE + linear_idx%SUPERGROUP_SIZE; + } else { + const int remainder_task_id = linear_idx - supersection_numel; + row_idx = remainder_task_id/finalsection_cols; + col_idx = supersection_cols + remainder_task_id%finalsection_cols; + } + return { (supergroup_idx&1) ? num_rows-row_idx-1 : row_idx, col_idx }; + } else { + const int supergroup_numel = num_cols*SUPERGROUP_SIZE; + const int supergroup_idx = linear_idx/supergroup_numel; + const int supersection_rows = (num_rows/SUPERGROUP_SIZE)*SUPERGROUP_SIZE; + const int supersection_numel = num_cols*supersection_rows; + const int finalsection_rows = num_rows-supersection_rows; + int row_idx, col_idx; + if (linear_idx < supersection_numel) { + row_idx = supergroup_idx*SUPERGROUP_SIZE + linear_idx%SUPERGROUP_SIZE; + col_idx = (linear_idx%supergroup_numel)/SUPERGROUP_SIZE; + } else { + const int remainder_task_id = linear_idx - supersection_numel; + row_idx = supersection_rows + remainder_task_id%finalsection_rows; + col_idx = remainder_task_id/finalsection_rows; + } + return { row_idx, (supergroup_idx&1) ? num_cols-col_idx-1 : col_idx }; + } +} + +template +__device__ static inline int get_phasebit(uint32_t bitfield, int ring_id) { + if constexpr (half == 0) + return (bitfield >> (ring_id)) & 0b1; + else if constexpr (half == 1) + return (bitfield >> (ring_id + 16)) & 0b1; + else + asm volatile ("brkpt;\n"); + return -1; +} + +template +__device__ static inline void update_phasebit(uint32_t &bitfield, int ring_id) { + if constexpr (half == 0) + bitfield ^= (1 << (ring_id)); + else if constexpr (half == 1) + bitfield ^= (1 << (ring_id + 16)); + else + asm volatile ("brkpt;\n"); +} + +template __device__ static inline int ring_advance(int ring, int distance=1) { return (ring + distance) % N; } +template __device__ static inline int ring_retreat(int ring, int distance=1) { return (ring + 16*N - distance) % N; } + + + +#ifndef KITTENS_NO_HOST + + +#define CUCHECK(cmd) do { \ + CUresult err = cmd; \ + if (err != CUDA_SUCCESS) { \ + const char *errStr; \ + cuGetErrorString(err, &errStr); \ + fprintf(stderr, "Failed: CUDA error %s:%d '%s'\n", \ + __FILE__, __LINE__, errStr); \ + exit(EXIT_FAILURE); \ + } \ +} while(0) + + +#define CUDACHECK(cmd) do { \ + cudaError_t err = cmd; \ + if (err != cudaSuccess) { \ + fprintf(stderr, "Failed: CUDA error %s:%d '%s'\n", \ + __FILE__, __LINE__, cudaGetErrorString(err)); \ + exit(EXIT_FAILURE); \ + } \ +} while(0) + + +#define CHECK_CUDA_ERROR(val) check((val), #val, __FILE__, __LINE__) +template __host__ void check( + T err, char const* const func, char const* const file, int const line +) { + if (err != cudaSuccess) { + std::cerr << "CUDA Runtime Error at: " << file << ":" << line + << std::endl; + std::cerr << cudaGetErrorString(err) << " " << func << std::endl; + std::exit(EXIT_FAILURE); + } +} + + + + + + + __host__ inline int num_sms(int device_id = -1) { + if (device_id < 0) CUDACHECK(cudaGetDevice(&device_id)); + int sm_count; + CUDACHECK(cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device_id)); + return sm_count; +} + +template +struct LaunchConfig { + static constexpr unsigned int num_attributes = CLUSTER ? (PDL ? 3 : 2) : 1; + + cudaLaunchAttribute attributes[num_attributes]; + cudaLaunchConfig_t config = {0}; + + __host__ inline LaunchConfig(dim3 grid, dim3 block, size_t dynamic_shared_memory, + cudaStream_t stream) noexcept requires(!CLUSTER) { + if constexpr (PDL) { + attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attributes[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attributes; + config.numAttrs = num_attributes; + } else { + config.numAttrs = 0; + } + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes= dynamic_shared_memory; + config.stream = stream; + } + + __host__ inline LaunchConfig(dim3 grid, dim3 block, size_t dynamic_shared_memory, + cudaStream_t stream, dim3 cluster_preferred, dim3 cluster_minimum) noexcept requires(CLUSTER) { + attributes[0].id = cudaLaunchAttributePreferredClusterDimension; + attributes[0].val.preferredClusterDim.x = cluster_preferred.x; + attributes[0].val.preferredClusterDim.y = cluster_preferred.y; + attributes[0].val.preferredClusterDim.z = cluster_preferred.z; + attributes[1].id = cudaLaunchAttributeClusterDimension; + attributes[1].val.clusterDim.x = cluster_minimum.x; + attributes[1].val.clusterDim.y = cluster_minimum.y; + attributes[1].val.clusterDim.z = cluster_minimum.z; + if constexpr (PDL) { + attributes[2].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attributes[2].val.programmaticStreamSerializationAllowed = 1; + } + config.attrs = attributes; + config.numAttrs = num_attributes; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes= dynamic_shared_memory; + config.stream = stream; + } + + __host__ inline LaunchConfig(dim3 grid, dim3 block, size_t dynamic_shared_memory, + cudaStream_t stream, dim3 cluster) noexcept requires(CLUSTER) + : LaunchConfig(grid, block, dynamic_shared_memory, stream, cluster, cluster) { } + + __host__ inline LaunchConfig(const LaunchConfig& other) noexcept : config(other.config) { + std::copy_n(other.attributes, num_attributes, attributes); + config.attrs = attributes; + } + + __host__ inline LaunchConfig& operator=(const LaunchConfig& other) noexcept { + if (this != &other) { + std::copy_n(other.attributes, num_attributes, attributes); + config = other.config; + config.attrs = attributes; + } + return *this; + } + + __host__ LaunchConfig(LaunchConfig&&) = delete; + __host__ LaunchConfig& operator=(LaunchConfig&&) = delete; + + __host__ inline operator cudaLaunchConfig_t*() noexcept { return &config; } + __host__ inline operator const cudaLaunchConfig_t*() const noexcept { return &config; } +}; + +#endif + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/kittens.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/kittens.cuh new file mode 100644 index 0000000..6578fff --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/kittens.cuh @@ -0,0 +1,73 @@ + + + + + +#pragma once + + +#if defined(KITTENS_AMPERE) +#define KITTENS_SM80 +#endif +#if defined(KITTENS_HOPPER) +#define KITTENS_SM90 +#endif +#if defined(KITTENS_BLACKWELL) +#define KITTENS_SM100 +#endif + + +#if defined(KITTENS_SM80) + defined(KITTENS_SM90) + defined(KITTENS_SM100) + defined(KITTENS_SM103) + defined(KITTENS_SM120) != 1 +#error "Define exactly one of: KITTENS_SM80, KITTENS_SM90, KITTENS_SM100, KITTENS_SM103, KITTENS_SM120" +#endif + + +#if defined(KITTENS_SM100) || defined(KITTENS_SM103) +#define KITTENS_SM10X +#endif + + +#include +#include +#include +#include +#include +#ifdef KITTENS_NO_HOST +namespace std { using namespace cuda::std; } +using uint = uint32_t; +#if defined(__CUDACC_VER_MAJOR__) && __CUDACC_VER_MAJOR__ >= 13 +struct alignas(128) CUtensorMap { char __opaque[128]; }; +#else +struct alignas(64) CUtensorMap { char __opaque[128]; }; +#endif +#endif + + +#include +#include +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +#include +#endif +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) +#include +#endif + + +#ifndef KITTENS_NO_HOST +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#endif + + +#include "common/common.cuh" +#include "types/types.cuh" +#include "ops/ops.cuh" diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/group.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/group.cuh new file mode 100644 index 0000000..bb48e8d --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/group.cuh @@ -0,0 +1,117 @@ + + + + + +#pragma once + +#include "../../common/common.cuh" +#include "../../types/types.cuh" +#include "../thread/thread.cuh" + +#define KITTENS_CHECK_WARP static_assert(GROUP_WARPS==1, "Warp (GROUP_WARPS=1) function called from a non-warp group."); + +#define KITTENS_CHECK_WARPGROUP static_assert(GROUP_WARPS==4, "Warpgroup (GROUP_WARPS=4) function called from a non-warpgroup group."); + + +#if defined(KITTENS_SM90) +#include "mma/base/base.cuh" +#endif + +namespace kittens { + + + +template +struct group { +static constexpr int GROUP_WARPS = _GROUP_WARPS; +static constexpr int GROUP_THREADS = GROUP_WARPS * kittens::WARP_THREADS; +__device__ static inline int laneid() { return threadIdx.x % GROUP_THREADS; } +__device__ static inline int warpid() { return laneid() / kittens::WARP_THREADS; } +__device__ static inline int groupid() { return threadIdx.x / GROUP_THREADS; } + +__device__ static inline void sync(int id) { + asm volatile("bar.sync %0, %1;\n" :: "r"(id), "n"(GROUP_THREADS)); +} +template __device__ static inline void sync() { + static_assert(GROUP_WARPS==1, "barrier-less sync() can only be called by a single warp!"); + asm volatile("bar.warp.sync %0;\n" :: "n"(MASK)); +} +__device__ static inline void arrive(int id) { + asm volatile("bar.arrive %0, %1;\n" :: "r"(id), "n"(GROUP_THREADS)); +} + +#include "memory/memory.cuh" +#include "shared/shared.cuh" +#include "register/register.cuh" +#include "mma/mma.cuh" +#include "util/util.cuh" + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + +template __device__ static inline void increase_registers() { + static_assert(n_reg % 8 == 0, "n_reg must be a multiple of 8"); + asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" :: "n"(n_reg)); +} +template __device__ static inline void decrease_registers() { + static_assert(n_reg % 8 == 0, "n_reg must be a multiple of 8"); + asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" :: "n"(n_reg)); +} +__device__ static inline void producer_registers() { decrease_registers<24>(); } +template __device__ static inline void consumer_registers() { increase_registers<480/NCWG - 8*(NCWG>3) - 224*(NCWG==1)>(); } + + + + + + +struct tma { +#include "memory/tile/tma.cuh" +#include "memory/vec/tma.cuh" +#include "util/tma.cuh" +struct cluster { +#include "memory/tile/tma_cluster.cuh" +#include "memory/vec/tma_cluster.cuh" +#include "util/tma_cluster.cuh" +}; +}; + +#endif + +}; + +namespace everyone { + + +__device__ static inline void sync(int id) { + asm volatile("bar.sync %0;\n" :: "r"(id)); +} + + +namespace tma { +namespace cluster { +__device__ static inline void arrive() { + asm volatile ("barrier.cluster.arrive.release;\n"); +} +__device__ static inline void arrive_aligned() { + asm volatile ("barrier.cluster.arrive.release.aligned;\n"); +} +__device__ static inline void wait() { + asm volatile ("barrier.cluster.wait.acquire;\n"); +} +__device__ static inline void wait_aligned() { + asm volatile ("barrier.cluster.wait.acquire.aligned;\n"); +} +__device__ static inline void sync() { + arrive_aligned(); + wait_aligned(); +} +} +} + +}; + +using warp = group<1>; +using warpgroup = group<4>; + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/memory.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/memory.cuh new file mode 100644 index 0000000..404f8af --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/memory.cuh @@ -0,0 +1,7 @@ + + + + + +#include "tile/tile.cuh" +#include "vec/vec.cuh" diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/complex_global_to_register.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/complex_global_to_register.cuh new file mode 100644 index 0000000..989b677 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/complex_global_to_register.cuh @@ -0,0 +1,42 @@ + + + + + + + + + + + + + + +template>> +__device__ inline static void load(CRT &dst, const CGL &src, const COORD &idx) { + load(dst.real, src.real, idx); + load(dst.imag, src.imag, idx); +} +template>> +__device__ inline static void load(CRT &dst, const CGL &src, const COORD &idx) { + load<2, CRT, CGL>(dst, src, idx); +} + + + + + + + + + + +template>> +__device__ inline static void store(CGL &dst, const CRT &src, const COORD &idx) { + store(dst.real, src.real, idx); + store(dst.imag, src.imag, idx); +} +template>> +__device__ inline static void store(CGL &dst, const CRT &src, const COORD &idx) { + store<2, CRT, CGL>(dst, src, idx); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/complex_global_to_shared.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/complex_global_to_shared.cuh new file mode 100644 index 0000000..28c2565 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/complex_global_to_shared.cuh @@ -0,0 +1,69 @@ + + + + + + + + + + + + + + + +template> +__device__ static inline void load(CST &dst, const CGL &src, const COORD &idx) { + load(dst.real, src.real, idx); + load(dst.imag, src.imag, idx); +} +template> +__device__ static inline void load(CST &dst, const CGL &src, const COORD &idx) { + load<2, false, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx); + load<2, false, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx); +} + + + + + + + + + + + +template> +__device__ static inline void store(CGL &dst, const CST &src, const COORD &idx) { + store(dst.real, src.real, idx); + store(dst.imag, src.imag, idx); +} +template> +__device__ static inline void store(CGL &dst, const CST &src, const COORD &idx) { + store<2, false, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx); + store<2, false, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx); +} + + + + + + + + + + + + + +template> +__device__ static inline void load_async(CST &dst, const CGL &src, const COORD &idx) { + load_async(dst.real, src.real, idx); + load_async(dst.imag, src.imag, idx); +} +template> +__device__ static inline void load_async(CST &dst, const CGL &src, const COORD &idx) { + load_async<2, false, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx); + load_async<2, false, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/complex_shared_to_register.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/complex_shared_to_register.cuh new file mode 100644 index 0000000..63588be --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/complex_shared_to_register.cuh @@ -0,0 +1,32 @@ + + + + + + + + + + + + + +template +__device__ inline static void load(CRT &dst, const CST &src) { + load(dst.real, src.real); + load(dst.imag, src.imag); +} + + + + + + + + + +template +__device__ inline static void store(CST &dst, const CRT &src) { + store(dst.real, src.real); + store(dst.imag, src.imag); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/global_to_register.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/global_to_register.cuh new file mode 100644 index 0000000..573c31a --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/global_to_register.cuh @@ -0,0 +1,199 @@ + + + + + + + + + + + + + +template>> +__device__ inline static void load(RT &dst, const GL &src, const COORD &idx) { + using T2 = RT::dtype; + using U = typename GL::dtype; + + U *src_ptr = (U*)&src[(idx.template unit_coord())]; + const int row_stride = src.template stride(); + using U2 = base_types::packing::packed_type; + constexpr int pack_size = base_types::packing::num(); + constexpr int half_col = RT::tile_size_col / 2; + int warp_laneid = threadIdx.x % WARP_THREADS; + int local_warpid; + if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4)); + else local_warpid = warpid(); + const int row_offset = RT::rows*local_warpid; + #pragma unroll + for(int i = 0; i < RT::height; i++) { + int row = row_offset + i*RT::tile_size_row + (warp_laneid / 4); + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + pack_size*(warp_laneid % 4); + dst.tiles[i][j].data[0] = base_types::convertor::convert(*(U2*)(&src_ptr[(row+0)*row_stride + (col+0)])); + dst.tiles[i][j].data[2] = base_types::convertor::convert(*(U2*)(&src_ptr[(row+0)*row_stride + (col+half_col)])); + } + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + pack_size*(warp_laneid % 4); + dst.tiles[i][j].data[1] = base_types::convertor::convert(*(U2*)(&src_ptr[(row+8)*row_stride + (col+0)])); + dst.tiles[i][j].data[3] = base_types::convertor::convert(*(U2*)(&src_ptr[(row+8)*row_stride + (col+half_col)])); + } + } +} + + + + + + + + +template>> +__device__ inline static void load(RT &dst, const GL &src, const COORD &idx) { + using T = typename RT::T; + using U = typename GL::dtype; + + #if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert(!std::is_same_v && !std::is_same_v, "Unsupported type for load/store"); + #endif + + U *src_ptr = (U*)&src[(idx.template unit_coord())]; + const int row_stride = src.template stride(); + int warp_laneid = threadIdx.x % WARP_THREADS; + int local_warpid; + if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4)); + else local_warpid = warpid(); + const int row_offset = RT::rows*local_warpid; + #pragma unroll + for(int i = 0; i < RT::height; i++) { + int row = row_offset + i*RT::tile_size_row + 2*(warp_laneid % 4); + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + (warp_laneid / 4); + dst.tiles[i][j].data[0].x = base_types::convertor::convert(src_ptr[(row+0)*row_stride + (col+0)]); + dst.tiles[i][j].data[1].x = base_types::convertor::convert(src_ptr[(row+0)*row_stride + (col+8)]); + } + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + (warp_laneid / 4); + dst.tiles[i][j].data[0].y = base_types::convertor::convert(src_ptr[(row+1)*row_stride + (col+0)]); + dst.tiles[i][j].data[1].y = base_types::convertor::convert(src_ptr[(row+1)*row_stride + (col+8)]); + } + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + (warp_laneid / 4); + dst.tiles[i][j].data[2].x = base_types::convertor::convert(src_ptr[(row+8)*row_stride + (col+0)]); + dst.tiles[i][j].data[3].x = base_types::convertor::convert(src_ptr[(row+8)*row_stride + (col+8)]); + } + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + (warp_laneid / 4); + dst.tiles[i][j].data[2].y = base_types::convertor::convert(src_ptr[(row+9)*row_stride + (col+0)]); + dst.tiles[i][j].data[3].y = base_types::convertor::convert(src_ptr[(row+9)*row_stride + (col+8)]); + } + } +} +template>> +__device__ inline static void load(RT &dst, const GL &src, const COORD &idx) { + load<2>(dst, src, idx); +} + + + + + + + + +template>> +__device__ inline static void store(const GL &dst, const RT &src, const COORD &idx) { + using T2 = RT::dtype; + using U = typename GL::dtype; + + U *dst_ptr = (U*)&dst[(idx.template unit_coord())]; + const int row_stride = dst.template stride(); + using U2 = base_types::packing::packed_type; + constexpr int pack_size = base_types::packing::num(); + constexpr int half_col = RT::tile_size_col / 2; + int warp_laneid = threadIdx.x % WARP_THREADS; + int local_warpid; + if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4)); + else local_warpid = warpid(); + const int row_offset = RT::rows*local_warpid; + #pragma unroll + for(int i = 0; i < RT::height; i++) { + int row = row_offset + i*RT::tile_size_row + (warp_laneid / 4); + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + pack_size*(warp_laneid % 4); + *(U2*)(&dst_ptr[(row+0)*row_stride + (col+0)]) = base_types::convertor::convert(src.tiles[i][j].data[0]); + *(U2*)(&dst_ptr[(row+0)*row_stride + (col+half_col)]) = base_types::convertor::convert(src.tiles[i][j].data[2]); + } + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + pack_size*(warp_laneid % 4); + *(U2*)(&dst_ptr[(row+8)*row_stride + (col+0)]) = base_types::convertor::convert(src.tiles[i][j].data[1]); + *(U2*)(&dst_ptr[(row+8)*row_stride + (col+half_col)]) = base_types::convertor::convert(src.tiles[i][j].data[3]); + } + } +} + + + + + + + + +template>> +__device__ inline static void store(const GL &dst, const RT &src, const COORD &idx) { + using T = base_types::packing::unpacked_type; + using U = typename GL::dtype; + + #if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert(!std::is_same_v && !std::is_same_v, "Unsupported type for load/store"); + #endif + + U *dst_ptr = (U*)&dst[(idx.template unit_coord())]; + const int row_stride = dst.template stride(); + int warp_laneid = threadIdx.x % WARP_THREADS; + int local_warpid; + if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4)); + else local_warpid = warpid(); + const int row_offset = RT::rows*local_warpid; + #pragma unroll + for(int i = 0; i < RT::height; i++) { + int row = row_offset + i*RT::tile_size_row + 2*(warp_laneid % 4); + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + (warp_laneid / 4); + dst_ptr[(row+0)*row_stride + (col+0)] = base_types::convertor::convert(src.tiles[i][j].data[0].x); + dst_ptr[(row+0)*row_stride + (col+8)] = base_types::convertor::convert(src.tiles[i][j].data[1].x); + } + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + (warp_laneid / 4); + dst_ptr[(row+1)*row_stride + (col+0)] = base_types::convertor::convert(src.tiles[i][j].data[0].y); + dst_ptr[(row+1)*row_stride + (col+8)] = base_types::convertor::convert(src.tiles[i][j].data[1].y); + } + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + (warp_laneid / 4); + dst_ptr[(row+8)*row_stride + (col+0)] = base_types::convertor::convert(src.tiles[i][j].data[2].x); + dst_ptr[(row+8)*row_stride + (col+8)] = base_types::convertor::convert(src.tiles[i][j].data[3].x); + } + #pragma unroll + for(int j = 0; j < RT::width; j++) { + int col = j*RT::tile_size_col + (warp_laneid / 4); + dst_ptr[(row+9)*row_stride + (col+0)] = base_types::convertor::convert(src.tiles[i][j].data[2].y); + dst_ptr[(row+9)*row_stride + (col+8)] = base_types::convertor::convert(src.tiles[i][j].data[3].y); + } + } +} +template>> +__device__ inline static void store(const GL &dst, const RT &src, const COORD &idx) { + store<2>(dst, src, idx); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/global_to_shared.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/global_to_shared.cuh new file mode 100644 index 0000000..4f2ca78 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/global_to_shared.cuh @@ -0,0 +1,167 @@ + + + + + + + + + + + + + + +template> +__device__ static inline void load(ST &dst, const GL &src, const COORD &idx) { + using T = typename ST::dtype; + const int row_stride = src.template stride(); + + constexpr int elem_per_memcpy = sizeof(float4)/sizeof(typename ST::dtype); + constexpr int memcpy_per_row = ST::cols / elem_per_memcpy; + constexpr int total_calls = (ST::rows*ST::cols + GROUP_THREADS*elem_per_memcpy-1) / (GROUP_THREADS*elem_per_memcpy); + + coord<> unit_coord = idx.template unit_coord(); + typename GL::dtype *src_ptr = (typename GL::dtype*)&src[unit_coord]; + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst.data[0])); + int laneid = threadIdx.x % GROUP_THREADS; + + #pragma unroll + for(int i = 0; i < total_calls; i++) { + + int load_idx = i * GROUP_THREADS + laneid; + + int row = load_idx / memcpy_per_row; + int col = (load_idx*elem_per_memcpy) % ST::cols; + + if constexpr (assume_aligned) { + float4 tmp; + move::ldg(tmp, (float4*)&src_ptr[row*row_stride + col]); + move::sts(dst.idx(dst_ptr, {row, col}), tmp); + } + else { + if (row + unit_coord.template dim() < src.template shape()) { + float4 tmp; + move::ldg(tmp, (float4*)&src_ptr[row*row_stride + col]); + move::sts(dst.idx(dst_ptr, {row, col}), tmp); + } + else { + float4 zeros = {0.f,0.f,0.f,0.f}; + move::sts(dst.idx(dst_ptr, {row, col}), zeros); + } + } + } +} +template> +__device__ static inline void load(ST &dst, const GL &src, const COORD &idx) { + load<2, false, ST, GL, COORD>(dst, src, idx); +} + + + + + + + + + +template> +__device__ static inline void store(const GL &dst, const ST &src, const COORD &idx) { + using T = typename ST::dtype; + const int row_stride = dst.template stride(); + + constexpr int elem_per_memcpy = sizeof(float4)/sizeof(typename ST::dtype); + constexpr int memcpy_per_row = ST::cols / elem_per_memcpy; + constexpr int total_calls = (ST::rows*ST::cols + GROUP_THREADS*elem_per_memcpy-1) / (GROUP_THREADS*elem_per_memcpy); + + coord<> unit_coord = idx.template unit_coord(); + typename GL::dtype *dst_ptr = (typename GL::dtype*)&dst[unit_coord]; + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src.data[0])); + int laneid = threadIdx.x % GROUP_THREADS; + + #pragma unroll + for(int i = 0; i < total_calls; i++) { + + int load_idx = i * GROUP_THREADS + laneid; + + int row = load_idx / memcpy_per_row; + int col = (load_idx*elem_per_memcpy) % ST::cols; + + if constexpr (assume_aligned) { + float4 tmp; + move::lds(tmp, src.idx(src_ptr, {row, col})); + move::stg((float4*)&dst_ptr[row*row_stride + col], tmp); + } + else { + if (row + unit_coord.template dim() < dst.template shape()) { + float4 tmp; + move::lds(tmp, src.idx(src_ptr, {row, col})); + move::stg((float4*)&dst_ptr[row*row_stride + col], tmp); + } + } + } +} +template> +__device__ static inline void store(const GL &dst, const ST &src, const COORD &idx) { + store<2, false, ST, GL, COORD>(dst, src, idx); +} + + + + + + + + + + +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx) { + using T = typename ST::dtype; + const int row_stride = src.template stride(); + + constexpr int elem_per_memcpy = sizeof(float4)/sizeof(typename ST::dtype); + constexpr int memcpy_per_row = ST::cols / elem_per_memcpy; + constexpr int total_calls = (ST::rows*ST::cols + GROUP_THREADS*elem_per_memcpy-1) / (GROUP_THREADS*elem_per_memcpy); + + coord<> unit_coord = idx.template unit_coord(); + typename GL::dtype *src_ptr = (typename GL::dtype*)&src[unit_coord]; + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst.data[0])); + int laneid = threadIdx.x % GROUP_THREADS; + + #pragma unroll + for(int i = 0; i < total_calls; i++) { + + int load_idx = i * GROUP_THREADS + laneid; + + int row = load_idx / memcpy_per_row; + int col = (load_idx*elem_per_memcpy) % ST::cols; + + if constexpr (assume_aligned) { + asm volatile( + "cp.async.cg.shared.global.L2::128B [%0], [%1], 16;\n" + :: "r"(dst.idx(dst_ptr, {row, col})), "l"(&src_ptr[row*row_stride + col]) + : "memory" + ); + } + else { + if (row + unit_coord.template dim() < src.template shape()) { + asm volatile( + "cp.async.cg.shared.global.L2::128B [%0], [%1], 16;\n" + :: "r"(dst.idx(dst_ptr, {row, col})), "l"(&src_ptr[row*row_stride + col]) + : "memory" + ); + } + else { + + float4 zeros = {0.f,0.f,0.f,0.f}; + move::sts(dst.idx(dst_ptr, {row, col}), zeros); + } + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); +} +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx) { + load_async<2, false, ST, GL, COORD>(dst, src, idx); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/parallel_global_to_global.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/parallel_global_to_global.cuh new file mode 100644 index 0000000..1ce7a5c --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/parallel_global_to_global.cuh @@ -0,0 +1,51 @@ + + + + + + + +template +__device__ inline static void all_reduce(PGL &pgl, const coord &idx) { + static_assert(std::is_same_v, "Currently only bf16 is supported."); + static_assert(TILE_COLS > 0 && TILE_COLS % (WARP_THREADS * 2) == 0, "TILE_COLS must be a multiple of 64 for best performance."); + + const int warps_per_row = TILE_COLS / (WARP_THREADS * 2); + const int row_base = idx.r * TILE_ROWS; + const int col_base = idx.c * TILE_COLS; + const int warp_laneid = threadIdx.x % WARP_THREADS; + + for (int i = warpid(); i < TILE_ROWS * warps_per_row; i += GROUP_WARPS) { + int row_idx = row_base + i / warps_per_row; + int col_idx = col_base + (i % warps_per_row) * WARP_THREADS * 2 + warp_laneid * 2; + bf16_2 tmp; + bf16_2 *ptr = reinterpret_cast(pgl.mc_ptr_at({idx.b, idx.d, row_idx, col_idx})); + multimem::ld_reduce(tmp, ptr); + multimem::st(ptr, tmp); + } +} + +template +__device__ inline static void reduce(GL &dst, const coord &dst_idx, PGL &src, const coord &src_idx) { + static_assert(std::is_same_v, "Currently only bf16 is supported."); + static_assert(TILE_COLS > 0 && TILE_COLS % (WARP_THREADS * 2) == 0, "TILE_COLS must be a multiple of 64 for best performance."); + + const int warps_per_row = TILE_COLS / (WARP_THREADS * 2); + const int src_row_base = src_idx.r * TILE_ROWS; + const int src_col_base = src_idx.c * TILE_COLS; + const int dst_row_base = dst_idx.r * TILE_ROWS; + const int dst_col_base = dst_idx.c * TILE_COLS; + const int warp_laneid = threadIdx.x % WARP_THREADS; + + for (int i = warpid(); i < TILE_ROWS * warps_per_row; i += GROUP_WARPS) { + int src_row_idx = src_row_base + i / warps_per_row; + int src_col_idx = src_col_base + (i % warps_per_row) * WARP_THREADS * 2 + warp_laneid * 2; + int dst_row_idx = dst_row_base + i / warps_per_row; + int dst_col_idx = dst_col_base + (i % warps_per_row) * WARP_THREADS * 2 + warp_laneid * 2; + (void)dst_row_idx; + (void)dst_col_idx; + bf16_2 tmp; + multimem::ld_reduce(tmp, reinterpret_cast(src.mc_ptr_at({src_idx.b, src_idx.d, src_row_idx, src_col_idx}))); + move::stg(reinterpret_cast(&dst[{dst_idx.b, dst_idx.d, dst_row_idx, dst_col_idx}]), tmp); + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/shared_to_register.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/shared_to_register.cuh new file mode 100644 index 0000000..eb43d9e --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/shared_to_register.cuh @@ -0,0 +1,321 @@ + + + + + + + + + + + + + +template +__device__ inline static void load(RT &dst, const ST &src) { + constexpr int warp_height = RT::height; + static_assert(ST::rows/RT::rows == GROUP_WARPS, "Group load / store requires tile height to be a multiple of the RT height."); + static_assert(ST::rows%RT::rows == 0, "Group load / store requires tile height to be a multiple of the RT height."); + static_assert(ST::cols==RT::cols, "Group load / store requires tile widths to match."); + int local_warpid; + if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4)); + else local_warpid = warpid(); + using T2 = RT::dtype; + using U = ST::dtype; + using T = base_types::packing::unpacked_type; + using U2 = base_types::packing::packed_type; + int warp_laneid = ::kittens::laneid(); + + + uint32_t shared_addr = static_cast(__cvta_generic_to_shared(&src.data[0])); + + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + if constexpr (sizeof(typename ST::dtype) == 2) { + + U2 tmp[4]; + int row = (local_warpid*warp_height + i)*RT::tile_size_row + (warp_laneid % 16); + int col = j*RT::tile_size_col + (warp_laneid / 16) * 8; + if constexpr (std::is_same_v) { + move::ldsm4(tmp[0], tmp[1], tmp[2], tmp[3], src.idx(shared_addr, {row, col})); + } + else { + move::ldsm4t(tmp[0], tmp[2], tmp[1], tmp[3], src.idx(shared_addr, {row, col})); + } + dst.tiles[i][j].data[0] = base_types::convertor::convert(tmp[0]); + dst.tiles[i][j].data[1] = base_types::convertor::convert(tmp[1]); + dst.tiles[i][j].data[2] = base_types::convertor::convert(tmp[2]); + dst.tiles[i][j].data[3] = base_types::convertor::convert(tmp[3]); + } + else if constexpr (std::is_same_v && sizeof(typename ST::dtype) == 1) { + + int warp_group_16 = (warp_laneid / 16); + int lane_in_16 = warp_laneid % 16; + int row = (local_warpid*warp_height + i)*RT::tile_size_row + (lane_in_16 % 16); + int col = j*RT::tile_size_col + warp_group_16 * 16; + + U2 tmp[4]; + if constexpr (std::is_same_v) { + move::ldsm4(tmp[0], tmp[1], tmp[2], tmp[3], src.idx(shared_addr, {row, col})); + } + else { + move::ldsm4t(tmp[0], tmp[2], tmp[1], tmp[3], src.idx(shared_addr, {row, col})); + } + dst.tiles[i][j].data[0] = base_types::convertor::convert(tmp[0]); + dst.tiles[i][j].data[1] = base_types::convertor::convert(tmp[1]); + dst.tiles[i][j].data[2] = base_types::convertor::convert(tmp[2]); + dst.tiles[i][j].data[3] = base_types::convertor::convert(tmp[3]); + } + else if constexpr (std::is_same_v && sizeof(typename ST::dtype) == 4) { + + int row = (local_warpid*warp_height + i)*RT::tile_size_row + (warp_laneid / 4); + int col = j*RT::tile_size_col + 2*(warp_laneid % 4); + if constexpr (ST::rows != ST::underlying_rows || ST::cols != ST::underlying_cols) { + row += src.row_offset; + col += src.col_offset; + } + int blit = sizeof(typename ST::dtype) * ((warp_laneid%4) / 2); + U2 tmp[4]; + static constexpr int swizzle_repeat = ST::swizzle_bytes * 8; + static constexpr int subtile_cols = ST::swizzle_bytes / sizeof(U); + const int outer_idx = col/subtile_cols; + const uint32_t addr_1 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+0)*subtile_cols + col%subtile_cols); + const uint32_t addr_2 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+8)*subtile_cols + col%subtile_cols); + const int swizzle_1 = blit ^ ((addr_1 % swizzle_repeat) >> 7) << 4; + const int swizzle_2 = blit ^ ((addr_2 % swizzle_repeat) >> 7) << 4; + move::lds(tmp[0].x, (addr_1+ 0)^swizzle_1); + move::lds(tmp[0].y, (addr_1+ 4)^swizzle_1); + move::lds(tmp[2].x, (addr_1+32)^swizzle_1); + move::lds(tmp[2].y, (addr_1+36)^swizzle_1); + move::lds(tmp[1].x, (addr_2+ 0)^swizzle_2); + move::lds(tmp[1].y, (addr_2+ 4)^swizzle_2); + move::lds(tmp[3].x, (addr_2+32)^swizzle_2); + move::lds(tmp[3].y, (addr_2+36)^swizzle_2); + dst.tiles[i][j].data[0] = base_types::convertor::convert(tmp[0]); + dst.tiles[i][j].data[1] = base_types::convertor::convert(tmp[1]); + dst.tiles[i][j].data[2] = base_types::convertor::convert(tmp[2]); + dst.tiles[i][j].data[3] = base_types::convertor::convert(tmp[3]); + if(blit) { + #pragma unroll + for(int k = 0; k < 4; k++) { + dst.tiles[i][j].data[k] = T2{dst.tiles[i][j].data[k].y, dst.tiles[i][j].data[k].x}; + } + } + } + else { + + int row = (local_warpid*warp_height + i)*RT::tile_size_row + 2*(warp_laneid % 4); + int col = j*RT::tile_size_col + (warp_laneid / 4); + U2 tmp[4]; + move::lds(tmp[0].x, src.idx(shared_addr, {row+0, col+0})); + move::lds(tmp[0].y, src.idx(shared_addr, {row+1, col+0})); + move::lds(tmp[1].x, src.idx(shared_addr, {row+0, col+8})); + move::lds(tmp[1].y, src.idx(shared_addr, {row+1, col+8})); + move::lds(tmp[2].x, src.idx(shared_addr, {row+8, col+0})); + move::lds(tmp[2].y, src.idx(shared_addr, {row+9, col+0})); + move::lds(tmp[3].x, src.idx(shared_addr, {row+8, col+8})); + move::lds(tmp[3].y, src.idx(shared_addr, {row+9, col+8})); + dst.tiles[i][j].data[0] = base_types::convertor::convert(tmp[0]); + dst.tiles[i][j].data[1] = base_types::convertor::convert(tmp[1]); + dst.tiles[i][j].data[2] = base_types::convertor::convert(tmp[2]); + dst.tiles[i][j].data[3] = base_types::convertor::convert(tmp[3]); + } + } + } +} + + + + + + + + + + +template +__device__ inline static void store(ST &dst, const RT &src) { + constexpr int warp_height = RT::height; + static_assert(ST::rows/RT::rows == GROUP_WARPS, "Group load / store requires tile height to be a multiple of the RT height."); + static_assert(ST::rows%RT::rows == 0, "Group load / store requires tile height to be a multiple of the RT height."); + static_assert(ST::cols==RT::cols, "Group load / store requires tile widths to match."); + int local_warpid; + if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4)); + else local_warpid = warpid(); + using T2 = RT::dtype; + using U = ST::dtype; + using T = base_types::packing::unpacked_type; + using U2 = base_types::packing::packed_type; + int warp_laneid = ::kittens::laneid(); + + + uint32_t shared_addr = static_cast(__cvta_generic_to_shared(&dst.data[0])); + + #pragma unroll + for(int i = 0; i < warp_height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + if constexpr (sizeof(typename ST::dtype) == 2) { + + U2 tmp[4]; + tmp[0] = base_types::convertor::convert(src.tiles[i][j].data[0]); + tmp[1] = base_types::convertor::convert(src.tiles[i][j].data[1]); + tmp[2] = base_types::convertor::convert(src.tiles[i][j].data[2]); + tmp[3] = base_types::convertor::convert(src.tiles[i][j].data[3]); +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + int row = (local_warpid*warp_height + i)*RT::tile_size_row + (warp_laneid % 16); + int col = j*RT::tile_size_col + (warp_laneid / 16) * 8; + if constexpr (std::is_same_v) { + move::stsm4(dst.idx(shared_addr, {row, col}), tmp[0], tmp[1], tmp[2], tmp[3]); + } + else { + move::stsm4t(dst.idx(shared_addr, {row, col}), tmp[0], tmp[2], tmp[1], tmp[3]); + } +#else + if constexpr (std::is_same_v) { + int row = (local_warpid*warp_height + i)*RT::tile_size_row + (warp_laneid / 4); + int col = j*RT::tile_size_col + 2*(warp_laneid % 4); + move::sts(dst.idx(shared_addr, {row+0, col+0}), tmp[0]); + move::sts(dst.idx(shared_addr, {row+8, col+0}), tmp[1]); + move::sts(dst.idx(shared_addr, {row+0, col+8}), tmp[2]); + move::sts(dst.idx(shared_addr, {row+8, col+8}), tmp[3]); + } + else { + int row = (local_warpid*warp_height + i)*RT::tile_size_row + 2*(warp_laneid % 4); + int col = j*RT::tile_size_col + (warp_laneid / 4); + move::sts(dst.idx(shared_addr, {row+0, col+0}), tmp[0].x); + move::sts(dst.idx(shared_addr, {row+1, col+0}), tmp[0].y); + move::sts(dst.idx(shared_addr, {row+0, col+8}), tmp[1].x); + move::sts(dst.idx(shared_addr, {row+1, col+8}), tmp[1].y); + move::sts(dst.idx(shared_addr, {row+8, col+0}), tmp[2].x); + move::sts(dst.idx(shared_addr, {row+9, col+0}), tmp[2].y); + move::sts(dst.idx(shared_addr, {row+8, col+8}), tmp[3].x); + move::sts(dst.idx(shared_addr, {row+9, col+8}), tmp[3].y); + } +#endif + } + else if constexpr (std::is_same_v && sizeof(typename ST::dtype) == 1) { + + + int warp_group_16 = (warp_laneid / 16); + int lane_in_16 = warp_laneid % 16; + int row = (local_warpid*warp_height + i)*RT::tile_size_row + (lane_in_16 % 16); + int col = j*RT::tile_size_col + warp_group_16 * 16; + + U2 tmp[4]; + tmp[0] = base_types::convertor::convert(src.tiles[i][j].data[0]); + tmp[1] = base_types::convertor::convert(src.tiles[i][j].data[1]); + tmp[2] = base_types::convertor::convert(src.tiles[i][j].data[2]); + tmp[3] = base_types::convertor::convert(src.tiles[i][j].data[3]); + if constexpr (std::is_same_v) { + move::stsm4(dst.idx(shared_addr, {row, col}), tmp[0], tmp[1], tmp[2], tmp[3]); + } + else { + move::stsm4t(dst.idx(shared_addr, {row, col}), tmp[0], tmp[2], tmp[1], tmp[3]); + } + } + else if constexpr (std::is_same_v && sizeof(typename ST::dtype) == 4) { + + int row = (local_warpid*warp_height + i)*RT::tile_size_row + (warp_laneid / 4); + int col = j*RT::tile_size_col + 2*(warp_laneid % 4); + if constexpr (ST::rows != ST::underlying_rows || ST::cols != ST::underlying_cols) { + row += dst.row_offset; + col += dst.col_offset; + } + int blit = sizeof(typename ST::dtype) * ((warp_laneid%4) / 2); + T2 reg_tmp[4]; + if(blit) { + #pragma unroll + for(int k = 0; k < 4; k++) { + reg_tmp[k] = T2{src.tiles[i][j].data[k].y, src.tiles[i][j].data[k].x}; + } + } + else { + #pragma unroll + for(int k = 0; k < 4; k++) { + reg_tmp[k] = src.tiles[i][j].data[k]; + } + } + U2 tmp[4]; + tmp[0] = base_types::convertor::convert(reg_tmp[0]); + tmp[1] = base_types::convertor::convert(reg_tmp[1]); + tmp[2] = base_types::convertor::convert(reg_tmp[2]); + tmp[3] = base_types::convertor::convert(reg_tmp[3]); + static constexpr int swizzle_repeat = ST::swizzle_bytes * 8; + static constexpr int subtile_cols = ST::swizzle_bytes / sizeof(U); + const int outer_idx = col/subtile_cols; + const uint32_t addr_1 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+0)*subtile_cols + col%subtile_cols); + const uint32_t addr_2 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+8)*subtile_cols + col%subtile_cols); + const int swizzle_1 = blit ^ ((addr_1 % swizzle_repeat) >> 7) << 4; + const int swizzle_2 = blit ^ ((addr_2 % swizzle_repeat) >> 7) << 4; + move::sts((addr_1+ 0)^swizzle_1, tmp[0].x); + move::sts((addr_1+ 4)^swizzle_1, tmp[0].y); + move::sts((addr_1+32)^swizzle_1, tmp[2].x); + move::sts((addr_1+36)^swizzle_1, tmp[2].y); + move::sts((addr_2+ 0)^swizzle_2, tmp[1].x); + move::sts((addr_2+ 4)^swizzle_2, tmp[1].y); + move::sts((addr_2+32)^swizzle_2, tmp[3].x); + move::sts((addr_2+36)^swizzle_2, tmp[3].y); + } + else { + + int row = (local_warpid*warp_height + i)*RT::tile_size_row + 2*(warp_laneid % 4); + int col = j*RT::tile_size_col + (warp_laneid / 4); + U2 tmp[4]; + tmp[0] = base_types::convertor::convert(src.tiles[i][j].data[0]); + tmp[1] = base_types::convertor::convert(src.tiles[i][j].data[1]); + tmp[2] = base_types::convertor::convert(src.tiles[i][j].data[2]); + tmp[3] = base_types::convertor::convert(src.tiles[i][j].data[3]); + move::sts(dst.idx(shared_addr, {row+0, col+0}), tmp[0].x); + move::sts(dst.idx(shared_addr, {row+1, col+0}), tmp[0].y); + move::sts(dst.idx(shared_addr, {row+0, col+8}), tmp[1].x); + move::sts(dst.idx(shared_addr, {row+1, col+8}), tmp[1].y); + move::sts(dst.idx(shared_addr, {row+8, col+0}), tmp[2].x); + move::sts(dst.idx(shared_addr, {row+9, col+0}), tmp[2].y); + move::sts(dst.idx(shared_addr, {row+8, col+8}), tmp[3].x); + move::sts(dst.idx(shared_addr, {row+9, col+8}), tmp[3].y); + } + } + } +} + + + +template +__device__ inline static auto load(RV &dst, const ST &src, int2 row_col) { + KITTENS_CHECK_WARP; + static_assert(ST::cols>=RV::length, "Shared tile must be at least as wide as the vector."); + using T = RV::T; + using U = ST::T; + int warp_laneid = ::kittens::laneid(); + + + uint32_t shared_addr = static_cast(__cvta_generic_to_shared(&src.data[0])); + + #pragma unroll + for(int col = warp_laneid; col < RV::length; col+=WARP_THREADS) { + U tmp; + move::lds(tmp, src.idx(shared_addr, {row_col.x, row_col.y + col})); + dst.data[col/WARP_THREADS][0] = base_types::convertor::convert(tmp); + } +} + +template +__device__ inline static auto store(ST &dst, const RV &src, int2 row_col) { + KITTENS_CHECK_WARP; + static_assert(ST::cols>=RV::length, "Shared tile must be at least as wide as the vector."); + using T = RV::T; + using U = ST::T; + int warp_laneid = ::kittens::laneid(); + + + uint32_t shared_addr = static_cast(__cvta_generic_to_shared(&dst.data[0])); + + #pragma unroll + for(int col = warp_laneid; col < RV::length; col+=WARP_THREADS) { + U tmp = base_types::convertor::convert(src.data[col/WARP_THREADS][0]); + move::sts(dst.idx(shared_addr, {row_col.x, row_col.y + col}), tmp); + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/tensor_to_register.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/tensor_to_register.cuh new file mode 100644 index 0000000..980c1c5 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/tensor_to_register.cuh @@ -0,0 +1,493 @@ + + + + + + + + + + + + + +template +__device__ inline static void load_async(RT &dst, const TM &src) { + if constexpr (GROUP_WARPS == 1) { + static_assert(RT::rows == TM::rows, "register tile and tensor tile must match rows"); + static_assert(RT::cols == TM::cols, "register tile and tensor tile must match cols"); + + using T2 = RT::dtype; + using U = typename TM::dtype; + using U2 = base_types::packing::packed_type; + + if constexpr (sizeof(typename TM::dtype) == 1) { + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + asm volatile( + "tcgen05.ld.sync.aligned.16x128b.x2.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(*(uint32_t*) &dst.tiles[i][j].data[0]), + "=r"(*(uint32_t*) &dst.tiles[i][j].data[1]), + "=r"(*(uint32_t*) &dst.tiles[i][j].data[2]), + "=r"(*(uint32_t*) &dst.tiles[i][j].data[3]) + : "r"(src.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))) + ); + } + } + } else if constexpr (sizeof(typename TM::dtype) == 2) { + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + asm volatile( + "tcgen05.ld.sync.aligned.16x128b.x2.pack::16b.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(*(uint32_t*) &dst.tiles[i][j].data[0]), + "=r"(*(uint32_t*) &dst.tiles[i][j].data[1]), + "=r"(*(uint32_t*) &dst.tiles[i][j].data[2]), + "=r"(*(uint32_t*) &dst.tiles[i][j].data[3]) + : "r"(src.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)) + ); + } + } + } + else if constexpr (sizeof(typename TM::dtype) == 4) { + if constexpr (std::is_same_v) { + #pragma unroll + for(int i = 0; i < RT::height; i++) { + if constexpr (RT::width%4 == 0) { + #pragma unroll + for(int j = 0; j < RT::width; j+=4) { + U2 data[16]; + 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];\n" + : "=r"(data[0].x), "=r"(data[0].y), + "=r"(data[1].x), "=r"(data[1].y), + "=r"(data[2].x), "=r"(data[2].y), + "=r"(data[3].x), "=r"(data[3].y), + "=r"(data[4].x), "=r"(data[4].y), + "=r"(data[5].x), "=r"(data[5].y), + "=r"(data[6].x), "=r"(data[6].y), + "=r"(data[7].x), "=r"(data[7].y), + "=r"(data[8].x), "=r"(data[8].y), + "=r"(data[9].x), "=r"(data[9].y), + "=r"(data[10].x), "=r"(data[10].y), + "=r"(data[11].x), "=r"(data[11].y), + "=r"(data[12].x), "=r"(data[12].y), + "=r"(data[13].x), "=r"(data[13].y), + "=r"(data[14].x), "=r"(data[14].y), + "=r"(data[15].x), "=r"(data[15].y) + : "r"(src.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))) + ); + #pragma unroll + for(int k = 0; k < 4; k++) { + dst.tiles[i][j+0].data[k] = base_types::convertor::convert(data[k]); + dst.tiles[i][j+1].data[k] = base_types::convertor::convert(data[k+4]); + dst.tiles[i][j+2].data[k] = base_types::convertor::convert(data[k+8]); + dst.tiles[i][j+3].data[k] = base_types::convertor::convert(data[k+12]); + } + } + } + else if constexpr (RT::width%2 == 0) { + #pragma unroll + for(int j = 0; j < RT::width; j+=2) { + U2 data[8]; + 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];\n" + : "=r"(data[0].x), "=r"(data[0].y), + "=r"(data[1].x), "=r"(data[1].y), + "=r"(data[2].x), "=r"(data[2].y), + "=r"(data[3].x), "=r"(data[3].y), + "=r"(data[4].x), "=r"(data[4].y), + "=r"(data[5].x), "=r"(data[5].y), + "=r"(data[6].x), "=r"(data[6].y), + "=r"(data[7].x), "=r"(data[7].y) + : "r"(src.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))) + ); + #pragma unroll + for(int k = 0; k < 4; k++) { + dst.tiles[i][j+0].data[k] = base_types::convertor::convert(data[k]); + dst.tiles[i][j+1].data[k] = base_types::convertor::convert(data[k+4]); + } + } + } + else { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + U2 data[4]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n" + : "=r"(data[0].x), "=r"(data[0].y), + "=r"(data[1].x), "=r"(data[1].y), + "=r"(data[2].x), "=r"(data[2].y), + "=r"(data[3].x), "=r"(data[3].y) + : "r"(src.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))) + ); + #pragma unroll + for(int k = 0; k < 4; k++) { + dst.tiles[i][j].data[k] = base_types::convertor::convert(data[k]); + } + } + } + } + } + else if constexpr (std::is_same_v) { + #pragma unroll + for(int i = 0; i < RT::height; i++) { + if constexpr (RT::width%4 == 0) { + #pragma unroll + for(int j = 0; j < RT::width; j+=4) { + U2 data[16]; + 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];\n" + : "=f"(data[0].x), "=f"(data[0].y), + "=f"(data[1].x), "=f"(data[1].y), + "=f"(data[2].x), "=f"(data[2].y), + "=f"(data[3].x), "=f"(data[3].y), + "=f"(data[4].x), "=f"(data[4].y), + "=f"(data[5].x), "=f"(data[5].y), + "=f"(data[6].x), "=f"(data[6].y), + "=f"(data[7].x), "=f"(data[7].y), + "=f"(data[8].x), "=f"(data[8].y), + "=f"(data[9].x), "=f"(data[9].y), + "=f"(data[10].x), "=f"(data[10].y), + "=f"(data[11].x), "=f"(data[11].y), + "=f"(data[12].x), "=f"(data[12].y), + "=f"(data[13].x), "=f"(data[13].y), + "=f"(data[14].x), "=f"(data[14].y), + "=f"(data[15].x), "=f"(data[15].y) + : "r"(src.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))) + ); + #pragma unroll + for(int k = 0; k < 4; k++) { + dst.tiles[i][j+0].data[k] = base_types::convertor::convert(data[k]); + dst.tiles[i][j+1].data[k] = base_types::convertor::convert(data[k+4]); + dst.tiles[i][j+2].data[k] = base_types::convertor::convert(data[k+8]); + dst.tiles[i][j+3].data[k] = base_types::convertor::convert(data[k+12]); + } + } + } + else if constexpr (RT::width%2 == 0) { + #pragma unroll + for(int j = 0; j < RT::width; j+=2) { + U2 data[8]; + 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];\n" + : "=f"(data[0].x), "=f"(data[0].y), + "=f"(data[1].x), "=f"(data[1].y), + "=f"(data[2].x), "=f"(data[2].y), + "=f"(data[3].x), "=f"(data[3].y), + "=f"(data[4].x), "=f"(data[4].y), + "=f"(data[5].x), "=f"(data[5].y), + "=f"(data[6].x), "=f"(data[6].y), + "=f"(data[7].x), "=f"(data[7].y) + : "r"(src.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))) + ); + #pragma unroll + for(int k = 0; k < 4; k++) { + dst.tiles[i][j+0].data[k] = base_types::convertor::convert(data[k]); + dst.tiles[i][j+1].data[k] = base_types::convertor::convert(data[k+4]); + } + } + } + else { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + U2 data[4]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n" + : "=f"(data[0].x), "=f"(data[0].y), + "=f"(data[1].x), "=f"(data[1].y), + "=f"(data[2].x), "=f"(data[2].y), + "=f"(data[3].x), "=f"(data[3].y) + : "r"(src.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))) + ); + #pragma unroll + for(int k = 0; k < 4; k++) { + dst.tiles[i][j].data[k] = base_types::convertor::convert(data[k]); + } + } + } + } + } + else { + static_assert(sizeof(U) == 999, "Unsupported 4-byte tensor memory type."); + } + } + } + else { + static_assert(GROUP_WARPS==4 || GROUP_WARPS==8); + constexpr int warp_rows = TM::rows/GROUP_WARPS; + static_assert(TM::cols==RT::cols); + static_assert(warp_rows==RT::rows); + if constexpr (GROUP_WARPS == 4) { + auto src_subtile = src.template subtile>(32*warpid(), 0); + ::kittens::group<1>::load_async(dst, src_subtile); + } + else { + auto src_subtile = src.template subtile>(32*(warpid()%4)+16*(warpid()/4), 0); + ::kittens::group<1>::load_async(dst, src_subtile); + } + } +} + + + + + + + + + + +template +__device__ inline static void store_async(TM &dst, const RT &src) { + if constexpr (GROUP_WARPS == 1) { + static_assert(RT::rows == TM::rows, "register tile and tensor tile must match rows"); + static_assert(RT::cols == TM::cols, "register tile and tensor tile must match cols"); + + using T2 = RT::dtype; + using T = base_types::packing::unpacked_type; + using U = TM::dtype; + using U2 = base_types::packing::packed_type; + + if constexpr (sizeof(typename TM::dtype) <= 2) { + #pragma unroll + for(int i = 0; i < RT::height; i++) { + if constexpr (RT::width%4 == 0) { + #pragma unroll + for(int j = 0; j < RT::width; j+=4) { + 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};\n" + :: "r"(dst.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))), + "r"(*(uint32_t*)&src.tiles[i][j+0].data[0]), + "r"(*(uint32_t*)&src.tiles[i][j+0].data[1]), + "r"(*(uint32_t*)&src.tiles[i][j+0].data[2]), + "r"(*(uint32_t*)&src.tiles[i][j+0].data[3]), + "r"(*(uint32_t*)&src.tiles[i][j+1].data[0]), + "r"(*(uint32_t*)&src.tiles[i][j+1].data[1]), + "r"(*(uint32_t*)&src.tiles[i][j+1].data[2]), + "r"(*(uint32_t*)&src.tiles[i][j+1].data[3]), + "r"(*(uint32_t*)&src.tiles[i][j+2].data[0]), + "r"(*(uint32_t*)&src.tiles[i][j+2].data[1]), + "r"(*(uint32_t*)&src.tiles[i][j+2].data[2]), + "r"(*(uint32_t*)&src.tiles[i][j+2].data[3]), + "r"(*(uint32_t*)&src.tiles[i][j+3].data[0]), + "r"(*(uint32_t*)&src.tiles[i][j+3].data[1]), + "r"(*(uint32_t*)&src.tiles[i][j+3].data[2]), + "r"(*(uint32_t*)&src.tiles[i][j+3].data[3]) + ); + } + } + else if constexpr (RT::width%2 == 0) { + #pragma unroll + for(int j = 0; j < RT::width; j+=2) { + asm volatile( + "tcgen05.st.sync.aligned.16x128b.x4.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};\n" + :: "r"(dst.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))), + "r"(*(uint32_t*)&src.tiles[i][j+0].data[0]), + "r"(*(uint32_t*)&src.tiles[i][j+0].data[1]), + "r"(*(uint32_t*)&src.tiles[i][j+0].data[2]), + "r"(*(uint32_t*)&src.tiles[i][j+0].data[3]), + "r"(*(uint32_t*)&src.tiles[i][j+1].data[0]), + "r"(*(uint32_t*)&src.tiles[i][j+1].data[1]), + "r"(*(uint32_t*)&src.tiles[i][j+1].data[2]), + "r"(*(uint32_t*)&src.tiles[i][j+1].data[3]) + ); + } + } + else { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + asm volatile( + "tcgen05.st.sync.aligned.16x128b.x2.b32 [%0], {%1, %2, %3, %4};\n" + :: "r"(dst.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))), + "r"(*(uint32_t*)&src.tiles[i][j].data[0]), + "r"(*(uint32_t*)&src.tiles[i][j].data[1]), + "r"(*(uint32_t*)&src.tiles[i][j].data[2]), + "r"(*(uint32_t*)&src.tiles[i][j].data[3]) + ); + } + } + } + } + else if constexpr (sizeof(typename TM::dtype) == 4) { + if constexpr (std::is_same_v) { + #pragma unroll + for(int i = 0; i < RT::height; i++) { + if constexpr(RT::width%4 == 0) { + #pragma unroll + for(int j = 0; j < RT::width; j+=4) { + U2 data[16]; + #pragma unroll + for(int k = 0; k < 4; k++) { + data[k] = base_types::convertor::convert(src.tiles[i][j].data[k]); + data[k+4] = base_types::convertor::convert(src.tiles[i][j+1].data[k]); + data[k+8] = base_types::convertor::convert(src.tiles[i][j+2].data[k]); + data[k+12] = base_types::convertor::convert(src.tiles[i][j+3].data[k]); + } + 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};\n" + :: "r"(dst.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))), + "r"(data[0].x), "r"(data[0].y), + "r"(data[1].x), "r"(data[1].y), + "r"(data[2].x), "r"(data[2].y), + "r"(data[3].x), "r"(data[3].y), + "r"(data[4].x), "r"(data[4].y), + "r"(data[5].x), "r"(data[5].y), + "r"(data[6].x), "r"(data[6].y), + "r"(data[7].x), "r"(data[7].y), + "r"(data[8].x), "r"(data[8].y), + "r"(data[9].x), "r"(data[9].y), + "r"(data[10].x), "r"(data[10].y), + "r"(data[11].x), "r"(data[11].y), + "r"(data[12].x), "r"(data[12].y), + "r"(data[13].x), "r"(data[13].y), + "r"(data[14].x), "r"(data[14].y), + "r"(data[15].x), "r"(data[15].y) + ); + } + } + else if constexpr(RT::width%2 == 0) { + #pragma unroll + for(int j = 0; j < RT::width; j+=2) { + U2 data[8]; + #pragma unroll + for(int k = 0; k < 4; k++) { + data[k] = base_types::convertor::convert(src.tiles[i][j].data[k]); + data[k+4] = base_types::convertor::convert(src.tiles[i][j+1].data[k]); + } + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x4.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};\n" + :: "r"(dst.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))), + "r"(data[0].x), "r"(data[0].y), + "r"(data[1].x), "r"(data[1].y), + "r"(data[2].x), "r"(data[2].y), + "r"(data[3].x), "r"(data[3].y), + "r"(data[4].x), "r"(data[4].y), + "r"(data[5].x), "r"(data[5].y), + "r"(data[6].x), "r"(data[6].y), + "r"(data[7].x), "r"(data[7].y) + ); + } + } + else { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + U2 data[4]; + #pragma unroll + for(int k = 0; k < 4; k++) { + data[k] = base_types::convertor::convert(src.tiles[i][j].data[k]); + } + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x2.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};\n" + :: "r"(dst.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))), + "r"(data[0].x), "r"(data[0].y), + "r"(data[1].x), "r"(data[1].y), + "r"(data[2].x), "r"(data[2].y), + "r"(data[3].x), "r"(data[3].y) + ); + } + } + } + } + else if constexpr (std::is_same_v) { + #pragma unroll + for(int i = 0; i < RT::height; i++) { + if constexpr(RT::width%4 == 0) { + #pragma unroll + for(int j = 0; j < RT::width; j+=4) { + U2 data[16]; + #pragma unroll + for(int k = 0; k < 4; k++) { + data[k] = base_types::convertor::convert(src.tiles[i][j].data[k]); + data[k+4] = base_types::convertor::convert(src.tiles[i][j+1].data[k]); + data[k+8] = base_types::convertor::convert(src.tiles[i][j+2].data[k]); + data[k+12] = base_types::convertor::convert(src.tiles[i][j+3].data[k]); + } + 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};\n" + :: "r"(dst.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))), + "f"(data[0].x), "f"(data[0].y), + "f"(data[1].x), "f"(data[1].y), + "f"(data[2].x), "f"(data[2].y), + "f"(data[3].x), "f"(data[3].y), + "f"(data[4].x), "f"(data[4].y), + "f"(data[5].x), "f"(data[5].y), + "f"(data[6].x), "f"(data[6].y), + "f"(data[7].x), "f"(data[7].y), + "f"(data[8].x), "f"(data[8].y), + "f"(data[9].x), "f"(data[9].y), + "f"(data[10].x), "f"(data[10].y), + "f"(data[11].x), "f"(data[11].y), + "f"(data[12].x), "f"(data[12].y), + "f"(data[13].x), "f"(data[13].y), + "f"(data[14].x), "f"(data[14].y), + "f"(data[15].x), "f"(data[15].y) + ); + } + } + else if constexpr(RT::width%2 == 0) { + #pragma unroll + for(int j = 0; j < RT::width; j+=2) { + U2 data[8]; + #pragma unroll + for(int k = 0; k < 4; k++) { + data[k] = base_types::convertor::convert(src.tiles[i][j].data[k]); + data[k+4] = base_types::convertor::convert(src.tiles[i][j+1].data[k]); + } + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x4.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};\n" + :: "r"(dst.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))), + "f"(data[0].x), "f"(data[0].y), + "f"(data[1].x), "f"(data[1].y), + "f"(data[2].x), "f"(data[2].y), + "f"(data[3].x), "f"(data[3].y), + "f"(data[4].x), "f"(data[4].y), + "f"(data[5].x), "f"(data[5].y), + "f"(data[6].x), "f"(data[6].y), + "f"(data[7].x), "f"(data[7].y) + ); + } + } + else { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + U2 data[4]; + #pragma unroll + for(int k = 0; k < 4; k++) { + data[k] = base_types::convertor::convert(src.tiles[i][j].data[k]); + } + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x2.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};\n" + :: "r"(dst.addr + ((i * RT::tile_size_row) << 16) + (j * RT::tile_size_col)/(4/(uint32_t)sizeof(U))), + "f"(data[0].x), "f"(data[0].y), + "f"(data[1].x), "f"(data[1].y), + "f"(data[2].x), "f"(data[2].y), + "f"(data[3].x), "f"(data[3].y) + ); + } + } + } + } + else { + static_assert(sizeof(U) == 999, "Unsupported 4-byte tensor memory type."); + } + } + } + else { + static_assert(GROUP_WARPS==4 || GROUP_WARPS==8); + constexpr int warp_rows = TM::rows/GROUP_WARPS; + static_assert(TM::cols==RT::cols); + static_assert(warp_rows==RT::rows); + if constexpr (GROUP_WARPS == 4) { + auto dst_subtile = dst.template subtile>(32*warpid(), 0); + ::kittens::group<1>::store_async(dst_subtile, src); + } + else { + auto dst_subtile = dst.template subtile>(32*(warpid()%4)+16*(warpid()/4), 0); + ::kittens::group<1>::store_async(dst_subtile, src); + } + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/tile.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/tile.cuh new file mode 100644 index 0000000..b67fee5 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/tile.cuh @@ -0,0 +1,20 @@ + + + + + +#include "global_to_register.cuh" +#include "global_to_shared.cuh" +#include "shared_to_register.cuh" +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +#include "parallel_global_to_global.cuh" +#endif +#ifdef KITTENS_SM10X +#include "tensor_to_register.cuh" +#endif + +#include "complex_shared_to_register.cuh" +#include "complex_global_to_register.cuh" +#include "complex_global_to_shared.cuh" + + diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/tma.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/tma.cuh new file mode 100644 index 0000000..4f0a208 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/tma.cuh @@ -0,0 +1,134 @@ + + + + + +template> +__device__ static inline void prefetch(ST &dst, const GL &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::prefetch(dst, src, idx); + } +} +template> +__device__ static inline void prefetch(ST &dst, const GL &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::prefetch(dst, src, idx); + } +} + +template> +__device__ static inline void store_async(const GL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_async(dst, src, idx); + } +} +template> +__device__ static inline void store_async(const GL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_async(dst, src, idx); + } +} + +template> +__device__ static inline void store_async(const PGL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_async(dst, src, idx); + } +} +template> +__device__ static inline void store_async(const PGL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_async(dst, src, idx); + } +} + +template> +__device__ static inline void store_add_async(const GL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_add_async(dst, src, idx); + } +} +template> +__device__ static inline void store_add_async(const GL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_add_async(dst, src, idx); + } +} + +template> +__device__ static inline void store_add_async(const PGL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_add_async(dst, src, idx); + } +} +template> +__device__ static inline void store_add_async(const PGL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_add_async(dst, src, idx); + } +} + +template> +__device__ static inline void store_min_async(const GL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_min_async(dst, src, idx); + } +} +template> +__device__ static inline void store_min_async(const GL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_min_async(dst, src, idx); + } +} + +template> +__device__ static inline void store_min_async(const PGL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_min_async(dst, src, idx); + } +} +template> +__device__ static inline void store_min_async(const PGL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_min_async(dst, src, idx); + } +} + +template> +__device__ static inline void store_max_async(const GL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_max_async(dst, src, idx); + } +} +template> +__device__ static inline void store_max_async(const GL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_max_async(dst, src, idx); + } +} + +template> +__device__ static inline void store_max_async(const PGL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_max_async(dst, src, idx); + } +} +template> +__device__ static inline void store_max_async(const PGL &dst, const ST &src, const COORD &idx) { + if(laneid() == 0) { + ::kittens::tma::store_max_async(dst, src, idx); + } +} + +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar) { + if(laneid() == 0) { + ::kittens::tma::load_async(dst, src, idx, bar); + } +} +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar) { + if(laneid() == 0) { + ::kittens::tma::load_async(dst, src, idx, bar); + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/tma_cluster.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/tma_cluster.cuh new file mode 100644 index 0000000..6982df8 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/tma_cluster.cuh @@ -0,0 +1,32 @@ + + + + + +#ifdef KITTENS_SM10X +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) { + if(laneid() == 0) { + ::kittens::tma::cluster::load_async(dst, src, idx, bar, cluster_mask, dst_mbar_cta); + } +} +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) { + if(laneid() == 0) { + ::kittens::tma::cluster::load_async(dst, src, idx, bar, cluster_mask, dst_mbar_cta); + } +} +#else +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask) { + if(laneid() == 0) { + ::kittens::tma::cluster::load_async(dst, src, idx, bar, cluster_mask); + } +} +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask) { + if(laneid() == 0) { + ::kittens::tma::cluster::load_async(dst, src, idx, bar, cluster_mask); + } +} +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/global_to_register.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/global_to_register.cuh new file mode 100644 index 0000000..1540cd8 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/global_to_register.cuh @@ -0,0 +1,138 @@ + + + + + + + + + + + + + +template +__device__ inline static void load(RV &dst, const GL &src, const coord> &idx) { + if constexpr (GROUP_WARPS == 1) { + using T2 = RV::dtype; + using U = typename GL::dtype; + using U2 = base_types::packing::packed_type; + using T = base_types::packing::unpacked_type; + + U *src_ptr = (U*)&src[(idx.template unit_coord<-1, 3>())]; + int laneid = ::kittens::laneid(); + + if constexpr (std::is_same_v) { + #pragma unroll + for(auto w = 0; w < (RV::outer_dim+3)/4; w++) { + int idx = w*64 + (laneid/4)*8 + 2*(laneid%4); + int o_dim = w*4 + (laneid/4) / 2; + int i_dim = (laneid/4) % 2; + + if(idx < RV::outer_dim*16) + dst[o_dim][i_dim] = base_types::convertor::convert(*(U2*)&src_ptr[idx]); + } + + #pragma unroll + for(auto w = 0; w < RV::outer_dim; w++) { + int leader = 8*(w%4) + (laneid%4); + dst[w][0] = packed_shfl_sync(MASK_ALL, dst[w][0], leader); + dst[w][1] = packed_shfl_sync(MASK_ALL, dst[w][1], leader+4); + } + } + else if constexpr (std::is_same_v) { + + + #pragma unroll + for(auto w = 0; w < (RV::outer_dim+1)/2; w++) { + int idx = w*32 + (laneid%4)*8 + (laneid/4); + int o_dim = w*2 + (laneid%4) / 2; + + if(idx < RV::outer_dim*16) { + T tmp = base_types::convertor::convert(src_ptr[idx]); + if(laneid%2==0) dst[o_dim][0].x = tmp; + else dst[o_dim][0].y = tmp; + } + } + + #pragma unroll + for(auto w = 0; w < RV::outer_dim; w++) { + int leader = (laneid/4)*4 + 2*(w%2); + dst[w][0].x = __shfl_sync(MASK_ALL, dst[w][0].x, leader); + dst[w][0].y = __shfl_sync(MASK_ALL, dst[w][0].y, leader+1); + } + } + else if constexpr (std::is_same_v) { + #pragma unroll + for(auto w = 0; w < RV::outer_dim; w++) { + if(w < RV::outer_dim-1 || RV::length%32 == 0 || laneid<16) { + dst[w][0] = base_types::convertor::convert(src_ptr[w*32 + laneid]); + } + } + } + } + else { + + ::kittens::group<1>::load(dst, src, coord(idx.b, idx.d, idx.r, idx.c*GROUP_WARPS+warpid())); + } +} + + + + + + + + +template +__device__ inline static void store(GL &dst, const RV &src, const coord> &idx) { + if constexpr (GROUP_WARPS == 1) { + using T2 = RV::dtype; + using U = typename GL::dtype; + using U2 = base_types::packing::packed_type; + using T = base_types::packing::unpacked_type; + + U *dst_ptr = (U*)&dst[(idx.template unit_coord<-1, 3>())]; + int laneid = ::kittens::laneid(); + + if constexpr (std::is_same_v) { + #pragma unroll + for(auto w = 0; w < (RV::outer_dim+3)/4; w++) { + int idx = w*64 + (laneid/4)*8 + 2*(laneid%4); + int o_dim = w*4 + (laneid/4) / 2; + int i_dim = (laneid/4) % 2; + + if(idx < RV::outer_dim*16) + *(U2*)&dst_ptr[idx] = base_types::convertor::convert(src[o_dim][i_dim]); + } + } + else if constexpr (std::is_same_v) { + + + #pragma unroll + for(auto w = 0; w < (RV::outer_dim+1)/2; w++) { + int idx = w*32 + (laneid%4)*8 + (laneid/4); + int o_dim = w*2 + (laneid%4) / 2; + + if(idx < RV::outer_dim*16) { + U tmp; + if(laneid%2==0) tmp = base_types::convertor::convert(src[o_dim][0].x); + else tmp = base_types::convertor::convert(src[o_dim][0].y); + dst_ptr[idx] = tmp; + } + } + } + else if constexpr (std::is_same_v) { + #pragma unroll + for(auto w = 0; w < RV::outer_dim; w++) { + if(w < RV::outer_dim-1 || RV::length%32 == 0 || laneid<16) { + dst_ptr[w*32 + laneid] = base_types::convertor::convert(src[w][0]); + } + } + } + } + else { + + ::kittens::group<1>::store(dst, src, coord(idx.b, idx.d, idx.r, idx.c*GROUP_WARPS+warpid())); + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/global_to_shared.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/global_to_shared.cuh new file mode 100644 index 0000000..cf265b6 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/global_to_shared.cuh @@ -0,0 +1,113 @@ + + + + + + + + + + + + + + + + +template> +__device__ static inline void load(SV &dst, const GL &src, const COORD &idx) { + constexpr uint32_t elem_per_transfer = sizeof(float4) / sizeof(typename SV::dtype); + constexpr uint32_t total_calls = SV::length / elem_per_transfer; + typename GL::dtype *src_ptr = (typename GL::dtype*)&src[(idx.template unit_coord<-1, 3>())]; + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst.data[0])); + #pragma unroll + for(uint32_t i = threadIdx.x%GROUP_THREADS; i < total_calls; i+=GROUP_THREADS) { + if(i * elem_per_transfer < SV::length) { + float4 tmp; + move::ldg(tmp, (float4*)&src_ptr[i*elem_per_transfer]); + move::sts(dst_ptr + sizeof(typename SV::dtype)*i*elem_per_transfer, tmp); + } + } +} + + + + + + + + + + + + +template> +__device__ static inline void store(GL &dst, const SV &src, const COORD &idx) { + constexpr uint32_t elem_per_transfer = sizeof(float4) / sizeof(typename SV::dtype); + constexpr uint32_t total_calls = SV::length / elem_per_transfer; + typename GL::dtype *dst_ptr = (typename GL::dtype*)&dst[(idx.template unit_coord<-1, 3>())]; + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src.data[0])); + #pragma unroll + for(uint32_t i = threadIdx.x%GROUP_THREADS; i < total_calls; i+=GROUP_THREADS) { + if(i * elem_per_transfer < SV::length) { + float4 tmp; + move::lds(tmp, src_ptr + sizeof(typename SV::dtype)*i*elem_per_transfer); + move::stg((float4*)&dst_ptr[i*elem_per_transfer], tmp); + } + } +} + + + + + + + + + +template> +__device__ static inline void load_async(SV &dst, const GL &src, const COORD &idx) { + constexpr uint32_t elem_per_transfer = sizeof(float4) / sizeof(typename SV::dtype); + constexpr uint32_t total_calls = SV::length / elem_per_transfer; + typename GL::dtype *src_ptr = (typename GL::dtype*)&src[(idx.template unit_coord<-1, 3>())]; + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst.data[0])); + #pragma unroll + for(uint32_t i = threadIdx.x%GROUP_THREADS; i < total_calls; i+=GROUP_THREADS) { + if(i * elem_per_transfer < SV::length) { + asm volatile( + "cp.async.cg.shared.global.L2::128B [%0], [%1], 16;\n" + :: "r"(dst_ptr + (uint32_t)sizeof(typename SV::dtype)*i*elem_per_transfer), "l"((uint64_t)&src_ptr[i*elem_per_transfer]) + : "memory" + ); + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); +} + + + + + + + + + + +template> +__device__ static inline void load_async_small(SV &dst, const GL &src, const COORD &idx) { + constexpr uint32_t elem_per_transfer = sizeof(float) / sizeof(typename SV::dtype); + constexpr uint32_t total_calls = SV::length / elem_per_transfer; + typename GL::dtype *src_ptr = (typename GL::dtype*)&src[(idx.template unit_coord<-1, 3>())]; + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst.data[0])); + #pragma unroll + for(uint32_t i = threadIdx.x%GROUP_THREADS; i < total_calls; i+=GROUP_THREADS) { + if(i * elem_per_transfer < SV::length) { + asm volatile( + "cp.async.ca.shared.global.L2::128B [%0], [%1], 4;\n" + :: "r"(dst_ptr + (uint32_t)sizeof(typename SV::dtype)*i*elem_per_transfer), "l"((uint64_t)&src_ptr[i*elem_per_transfer]) + : "memory" + ); + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/shared_to_register.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/shared_to_register.cuh new file mode 100644 index 0000000..1a279ae --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/shared_to_register.cuh @@ -0,0 +1,159 @@ + + + + + + + + + + + + + +template +__device__ inline static void load(RV &dst, const SV &src) { + using T2 = RV::dtype; + using U = SV::dtype; + using U2 = base_types::packing::packed_type; + using T = base_types::packing::unpacked_type; + if constexpr (GROUP_WARPS == 1) { + static_assert(SV::length == RV::length); + + int laneid = ::kittens::laneid(); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src.data[0])); + + __syncwarp(); + if constexpr (std::is_same_v) { + #pragma unroll + for(auto w = 0; w < (RV::outer_dim+3)/4; w++) { + int idx = w*64 + (laneid/4)*8 + 2*(laneid%4); + int o_dim = w*4 + (laneid/4) / 2; + int i_dim = (laneid/4) % 2; + + if(idx < RV::outer_dim*16) { + U2 tmp; + move::lds(tmp, src_ptr + sizeof(typename SV::dtype)*idx); + dst[o_dim][i_dim] = base_types::convertor::convert(tmp); + } + } + __syncwarp(); + + #pragma unroll + for(auto w = 0; w < RV::outer_dim; w++) { + int leader = 8*(w%4) + (laneid%4); + dst[w][0] = packed_shfl_sync(MASK_ALL, dst[w][0], leader); + dst[w][1] = packed_shfl_sync(MASK_ALL, dst[w][1], leader+4); + } + } + else if constexpr (std::is_same_v) { + + + #pragma unroll + for(auto w = 0; w < (RV::outer_dim+1)/2; w++) { + int idx = w*32 + (laneid%4)*8 + (laneid/4); + int o_dim = w*2 + (laneid%4) / 2; + + if(idx < RV::outer_dim*16) { + U tmp; + move::lds(tmp, src_ptr + sizeof(typename SV::dtype)*idx); + if(laneid%2==0) dst[o_dim][0].x = base_types::convertor::convert(tmp); + else dst[o_dim][0].y = base_types::convertor::convert(tmp); + } + } + __syncwarp(); + + #pragma unroll + for(auto w = 0; w < RV::outer_dim; w++) { + int leader = (laneid/4)*4 + 2*(w%2); + dst[w][0].x = __shfl_sync(MASK_ALL, dst[w][0].x, leader); + dst[w][0].y = __shfl_sync(MASK_ALL, dst[w][0].y, leader+1); + } + } + else if constexpr (std::is_same_v) { + #pragma unroll + for(auto w = 0; w < RV::outer_dim; w++) { + if(w < RV::outer_dim-1 || RV::length%32 == 0 || laneid<16) { + U tmp; + move::lds(tmp, src_ptr + sizeof(typename SV::dtype)*(w*32 + laneid)); + dst[w][0] = base_types::convertor::convert(tmp); + } + } + } + } + else { + static_assert(SV::length == RV::length*GROUP_WARPS); + auto &_src = src.template subvec(warpid()); + + ::kittens::group<1>::load(dst, _src); + } +} + + + + + + + + + +template +__device__ inline static void store(SV &dst, const RV &src) { + using T2 = RV::dtype; + using U = SV::dtype; + using U2 = base_types::packing::packed_type; + using T = base_types::packing::unpacked_type; + + if constexpr (GROUP_WARPS == 1) { + static_assert(SV::length == RV::length); + + int laneid = ::kittens::laneid(); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst.data[0])); + + __syncwarp(); + if constexpr (std::is_same_v) { + #pragma unroll + for(auto w = 0; w < (RV::outer_dim+3)/4; w++) { + int idx = w*64 + (laneid/4)*8 + 2*(laneid%4); + int o_dim = w*4 + (laneid/4) / 2; + int i_dim = (laneid/4) % 2; + + if(idx < RV::outer_dim*16) { + U2 tmp = base_types::convertor::convert(src[o_dim][i_dim]); + move::sts(dst_ptr + sizeof(typename SV::dtype)*idx, tmp); + } + } + } + else if constexpr (std::is_same_v) { + + + #pragma unroll + for(auto w = 0; w < (RV::outer_dim+1)/2; w++) { + int idx = w*32 + (laneid%4)*8 + (laneid/4); + int o_dim = w*2 + (laneid%4) / 2; + + if(idx < RV::outer_dim*16) { + U tmp; + if(laneid%2==0) tmp = base_types::convertor::convert(src[o_dim][0].x); + else tmp = base_types::convertor::convert(src[o_dim][0].y); + move::sts(dst_ptr + sizeof(typename SV::dtype)*idx, tmp); + } + } + } + else if constexpr (std::is_same_v) { + #pragma unroll + for(auto w = 0; w < RV::outer_dim; w++) { + if(w < RV::outer_dim-1 || RV::length%32 == 0 || laneid<16) { + U tmp = base_types::convertor::convert(src[w][0]); + move::sts(dst_ptr + sizeof(typename SV::dtype)*(w*32 + laneid), tmp); + } + } + } + } + else { + static_assert(SV::length == RV::length*GROUP_WARPS); + auto &_dst = dst.template subvec(warpid()); + + ::kittens::group<1>::store(_dst, src); + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/tma.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/tma.cuh new file mode 100644 index 0000000..2f9ddaa --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/tma.cuh @@ -0,0 +1,221 @@ + + + + + + + + + + + + + + + +template> +__device__ static inline void prefetch(SV &dst, const GL &src, const COORD &idx) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(src.template get_tma()); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + ::kittens::detail::tma::vec_prefetch_tma_internal(tma_ptr, tma_coord); + } +} +__KITTENS_TMA_DEFINE_DEFAULT_LOAD_CACHE_VEC__(prefetch) + + + + + + + + + + + + + + +template> +__device__ static inline void store_async(const GL &dst, const SV &src, const COORD &idx) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + store_commit_group(); +} +__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_async) + +template> +__device__ static inline void store_async(const PGL &dst, const SV &src, const COORD &idx) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + store_commit_group(); +} +__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_async) + + + + + + + + + + + + +template> +__device__ static inline void store_add_async(const GL &dst, const SV &src, const COORD &idx) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_add_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + store_commit_group(); +} +__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_add_async) + +template> +__device__ static inline void store_add_async(const PGL &dst, const SV &src, const COORD &idx) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_add_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + store_commit_group(); +} +__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_add_async) + + + + + + + + + + + + +template> +__device__ static inline void store_min_async(const GL &dst, const SV &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_min_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + store_commit_group(); +} +__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_min_async) + +template> +__device__ static inline void store_min_async(const PGL &dst, const SV &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_min_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + store_commit_group(); +} +__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_min_async) + + + + + + + + + + + +template> +__device__ static inline void store_max_async(const GL &dst, const SV &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_max_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + store_commit_group(); +} +__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_max_async) + +template> +__device__ static inline void store_max_async(const PGL &dst, const SV &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_max_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + store_commit_group(); +} +__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_max_async) + + + + + + + + + + + + +template> +__device__ static inline void load_async(SV &dst, const GL &src, const COORD &idx, semaphore& bar) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(src.template get_tma()); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(&bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst)); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t dst_i_ptr = dst_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_load_async_tma_internal(tma_ptr, dst_i_ptr, mbar_ptr, tma_coord); + } +} +__KITTENS_TMA_DEFINE_SEMAPHORE_CACHE_VEC__(load_async) diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/tma_cluster.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/tma_cluster.cuh new file mode 100644 index 0000000..3f45de7 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/tma_cluster.cuh @@ -0,0 +1,31 @@ + + + + + + + + + + + + + + + + + + template> + __device__ static inline void load_async(SV &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(src.template get_tma()); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(&bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst)); + for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2; i += WARP_THREADS) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t dst_i_ptr = dst_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::cluster::vec_load_async_tma_internal(tma_ptr, dst_i_ptr, mbar_ptr, tma_coord, cluster_mask, dst_mbar_cta); + } + } + __KITTENS_TMA_DEFINE_CLUSTER_SEMAPHORE_CACHE_VEC__(load_async) diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/vec.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/vec.cuh new file mode 100644 index 0000000..bf995b5 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/vec.cuh @@ -0,0 +1,10 @@ + + + + + +#include "shared_to_register.cuh" +#include "global_to_register.cuh" +#include "global_to_shared.cuh" + + diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x112.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x112.impl new file mode 100644 index 0000000..bf205aa --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x112.impl @@ -0,0 +1,334 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %61, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n112k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55}, " \ + "{%56, %57, %58, %59}, " \ + "%60, " \ + "p, 1, %63, %62;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %61, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n112k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55}, " \ + "{%56, %57, %58, %59}, " \ + "%60, " \ + "p, 1, %63, %62;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %33, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n112k16.f16.f16.f16 " \ + "{%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, " \ + "p, 1, %35, %34;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %58, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n112k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55}, " \ + "%56, " \ + "%57, " \ + "p, 1, %61, %59, %60;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %58, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n112k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55}, " \ + "%56, " \ + "%57, " \ + "p, 1, %61, %59, %60;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %30, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n112k16.f16.f16.f16 " \ + "{%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, " \ + "p, 1, %33, %31, %32;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x128.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x128.impl new file mode 100644 index 0000000..5690b2c --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x128.impl @@ -0,0 +1,1029 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %69, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "{%64, %65, %66, %67}, " \ + "%68, " \ + "p, 1, %71, %70;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y), + "+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y), + "+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y), + "+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y), + "+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %69, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "{%64, %65, %66, %67}, " \ + "%68, " \ + "p, 1, %71, %70;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y), + "+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y), + "+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y), + "+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y), + "+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %37, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "{%32, %33, %34, %35}, " \ + "%36, " \ + "p, 1, %39, %38;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %69, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "{%64, %65, %66, %67}, " \ + "%68, " \ + "p, 1, %70;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y), + "+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y), + "+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y), + "+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y), + "+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %69, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "{%64, %65, %66, %67}, " \ + "%68, " \ + "p, 1, %70;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y), + "+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y), + "+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y), + "+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y), + "+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %37, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.f16.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "{%32, %33, %34, %35}, " \ + "%36, " \ + "p, 1, %38;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %37, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.f16.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "{%32, %33, %34, %35}, " \ + "%36, " \ + "p, 1, %38;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %69, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "{%64, %65, %66, %67}, " \ + "%68, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %69, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "{%64, %65, %66, %67}, " \ + "%68, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %66, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "%64, " \ + "%65, " \ + "p, 1, %69, %67, %68;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y), + "+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y), + "+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y), + "+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y), + "+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %66, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "%64, " \ + "%65, " \ + "p, 1, %69, %67, %68;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y), + "+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y), + "+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y), + "+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y), + "+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %34, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "%32, " \ + "%33, " \ + "p, 1, %37, %35, %36;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %66, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "%64, " \ + "%65, " \ + "p, 1, %67;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y), + "+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y), + "+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y), + "+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y), + "+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %66, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "%64, " \ + "%65, " \ + "p, 1, %67;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y), + "+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y), + "+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y), + "+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y), + "+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y), + "+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y), + "+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y), + "+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y), + "+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %34, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.f16.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "%32, " \ + "%33, " \ + "p, 1, %35;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %34, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.f16.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "%32, " \ + "%33, " \ + "p, 1, %35;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][7].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %66, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "%64, " \ + "%65, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %66, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n128k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "%64, " \ + "%65, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x144.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x144.impl new file mode 100644 index 0000000..7f6756a --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x144.impl @@ -0,0 +1,382 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %77, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n144k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71}, " \ + "{%72, %73, %74, %75}, " \ + "%76, " \ + "p, 1, %79, %78;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %77, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n144k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71}, " \ + "{%72, %73, %74, %75}, " \ + "%76, " \ + "p, 1, %79, %78;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %41, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n144k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35}, " \ + "{%36, %37, %38, %39}, " \ + "%40, " \ + "p, 1, %43, %42;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %74, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n144k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71}, " \ + "%72, " \ + "%73, " \ + "p, 1, %77, %75, %76;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %74, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n144k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71}, " \ + "%72, " \ + "%73, " \ + "p, 1, %77, %75, %76;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %38, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n144k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35}, " \ + "%36, " \ + "%37, " \ + "p, 1, %41, %39, %40;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x16.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x16.impl new file mode 100644 index 0000000..34cf64f --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x16.impl @@ -0,0 +1,190 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %13, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n16k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7}, " \ + "{%8, %9, %10, %11}, " \ + "%12, " \ + "p, 1, %15, %14;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %13, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7}, " \ + "{%8, %9, %10, %11}, " \ + "%12, " \ + "p, 1, %15, %14;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %9, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n16k16.f16.f16.f16 " \ + "{%0, %1, %2, %3}, " \ + "{%4, %5, %6, %7}, " \ + "%8, " \ + "p, 1, %11, %10;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %10, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n16k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7}, " \ + "%8, " \ + "%9, " \ + "p, 1, %13, %11, %12;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %10, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7}, " \ + "%8, " \ + "%9, " \ + "p, 1, %13, %11, %12;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %6, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n16k16.f16.f16.f16 " \ + "{%0, %1, %2, %3}, " \ + "%4, " \ + "%5, " \ + "p, 1, %9, %7, %8;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x160.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x160.impl new file mode 100644 index 0000000..cf9d4cd --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x160.impl @@ -0,0 +1,914 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %85, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "{%80, %81, %82, %83}, " \ + "%84, " \ + "p, 1, %87, %86;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %85, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "{%80, %81, %82, %83}, " \ + "%84, " \ + "p, 1, %87, %86;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %45, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39}, " \ + "{%40, %41, %42, %43}, " \ + "%44, " \ + "p, 1, %47, %46;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %85, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "{%80, %81, %82, %83}, " \ + "%84, " \ + "p, 1, %86;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %85, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "{%80, %81, %82, %83}, " \ + "%84, " \ + "p, 1, %86;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %85, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "{%80, %81, %82, %83}, " \ + "%84, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %85, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "{%80, %81, %82, %83}, " \ + "%84, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %82, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "%80, " \ + "%81, " \ + "p, 1, %85, %83, %84;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %82, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "%80, " \ + "%81, " \ + "p, 1, %85, %83, %84;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %42, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39}, " \ + "%40, " \ + "%41, " \ + "p, 1, %45, %43, %44;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %82, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "%80, " \ + "%81, " \ + "p, 1, %83;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %82, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "%80, " \ + "%81, " \ + "p, 1, %83;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %82, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "%80, " \ + "%81, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %82, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n160k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \ + "%80, " \ + "%81, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x176.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x176.impl new file mode 100644 index 0000000..79fadbd --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x176.impl @@ -0,0 +1,430 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %93, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n176k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87}, " \ + "{%88, %89, %90, %91}, " \ + "%92, " \ + "p, 1, %95, %94;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %93, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n176k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87}, " \ + "{%88, %89, %90, %91}, " \ + "%92, " \ + "p, 1, %95, %94;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %49, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n176k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43}, " \ + "{%44, %45, %46, %47}, " \ + "%48, " \ + "p, 1, %51, %50;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %90, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n176k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87}, " \ + "%88, " \ + "%89, " \ + "p, 1, %93, %91, %92;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %90, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n176k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87}, " \ + "%88, " \ + "%89, " \ + "p, 1, %93, %91, %92;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %46, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n176k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43}, " \ + "%44, " \ + "%45, " \ + "p, 1, %49, %47, %48;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x192.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x192.impl new file mode 100644 index 0000000..d6c7068 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x192.impl @@ -0,0 +1,954 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %101, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "{%96, %97, %98, %99}, " \ + "%100, " \ + "p, 1, %103, %102;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %101, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "{%96, %97, %98, %99}, " \ + "%100, " \ + "p, 1, %103, %102;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %53, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "{%48, %49, %50, %51}, " \ + "%52, " \ + "p, 1, %55, %54;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %101, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "{%96, %97, %98, %99}, " \ + "%100, " \ + "p, 1, %102;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %101, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "{%96, %97, %98, %99}, " \ + "%100, " \ + "p, 1, %102;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %101, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "{%96, %97, %98, %99}, " \ + "%100, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %101, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "{%96, %97, %98, %99}, " \ + "%100, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %98, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "%96, " \ + "%97, " \ + "p, 1, %101, %99, %100;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %98, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "%96, " \ + "%97, " \ + "p, 1, %101, %99, %100;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %50, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "%48, " \ + "%49, " \ + "p, 1, %53, %51, %52;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %98, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "%96, " \ + "%97, " \ + "p, 1, %99;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %98, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "%96, " \ + "%97, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %98, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n192k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \ + "%96, " \ + "%97, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x208.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x208.impl new file mode 100644 index 0000000..06caa1e --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x208.impl @@ -0,0 +1,478 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %109, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n208k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103}, " \ + "{%104, %105, %106, %107}, " \ + "%108, " \ + "p, 1, %111, %110;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %109, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n208k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103}, " \ + "{%104, %105, %106, %107}, " \ + "%108, " \ + "p, 1, %111, %110;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %57, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n208k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51}, " \ + "{%52, %53, %54, %55}, " \ + "%56, " \ + "p, 1, %59, %58;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %106, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n208k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103}, " \ + "%104, " \ + "%105, " \ + "p, 1, %109, %107, %108;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %106, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n208k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103}, " \ + "%104, " \ + "%105, " \ + "p, 1, %109, %107, %108;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %54, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n208k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51}, " \ + "%52, " \ + "%53, " \ + "p, 1, %57, %55, %56;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x224.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x224.impl new file mode 100644 index 0000000..dc342d1 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x224.impl @@ -0,0 +1,1138 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %117, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "{%112, %113, %114, %115}, " \ + "%116, " \ + "p, 1, %119, %118;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %117, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "{%112, %113, %114, %115}, " \ + "%116, " \ + "p, 1, %119, %118;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %61, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55}, " \ + "{%56, %57, %58, %59}, " \ + "%60, " \ + "p, 1, %63, %62;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %117, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "{%112, %113, %114, %115}, " \ + "%116, " \ + "p, 1, %118;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %117, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "{%112, %113, %114, %115}, " \ + "%116, " \ + "p, 1, %118;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %117, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "{%112, %113, %114, %115}, " \ + "%116, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y), + "+r"(dst.tiles[0][12].data[0].x), "+r"(dst.tiles[0][12].data[0].y), + "+r"(dst.tiles[0][12].data[1].x), "+r"(dst.tiles[0][12].data[1].y), + "+r"(dst.tiles[0][12].data[2].x), "+r"(dst.tiles[0][12].data[2].y), + "+r"(dst.tiles[0][12].data[3].x), "+r"(dst.tiles[0][12].data[3].y), + "+r"(dst.tiles[0][13].data[0].x), "+r"(dst.tiles[0][13].data[0].y), + "+r"(dst.tiles[0][13].data[1].x), "+r"(dst.tiles[0][13].data[1].y), + "+r"(dst.tiles[0][13].data[2].x), "+r"(dst.tiles[0][13].data[2].y), + "+r"(dst.tiles[0][13].data[3].x), "+r"(dst.tiles[0][13].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %117, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "{%112, %113, %114, %115}, " \ + "%116, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y), + "+r"(dst.tiles[0][12].data[0].x), "+r"(dst.tiles[0][12].data[0].y), + "+r"(dst.tiles[0][12].data[1].x), "+r"(dst.tiles[0][12].data[1].y), + "+r"(dst.tiles[0][12].data[2].x), "+r"(dst.tiles[0][12].data[2].y), + "+r"(dst.tiles[0][12].data[3].x), "+r"(dst.tiles[0][12].data[3].y), + "+r"(dst.tiles[0][13].data[0].x), "+r"(dst.tiles[0][13].data[0].y), + "+r"(dst.tiles[0][13].data[1].x), "+r"(dst.tiles[0][13].data[1].y), + "+r"(dst.tiles[0][13].data[2].x), "+r"(dst.tiles[0][13].data[2].y), + "+r"(dst.tiles[0][13].data[3].x), "+r"(dst.tiles[0][13].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %114, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "%112, " \ + "%113, " \ + "p, 1, %117, %115, %116;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %114, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "%112, " \ + "%113, " \ + "p, 1, %117, %115, %116;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %58, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55}, " \ + "%56, " \ + "%57, " \ + "p, 1, %61, %59, %60;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %114, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "%112, " \ + "%113, " \ + "p, 1, %117, %115, %116;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %114, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "%112, " \ + "%113, " \ + "p, 1, %117, %115, %116;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %114, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "%112, " \ + "%113, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y), + "+r"(dst.tiles[0][12].data[0].x), "+r"(dst.tiles[0][12].data[0].y), + "+r"(dst.tiles[0][12].data[1].x), "+r"(dst.tiles[0][12].data[1].y), + "+r"(dst.tiles[0][12].data[2].x), "+r"(dst.tiles[0][12].data[2].y), + "+r"(dst.tiles[0][12].data[3].x), "+r"(dst.tiles[0][12].data[3].y), + "+r"(dst.tiles[0][13].data[0].x), "+r"(dst.tiles[0][13].data[0].y), + "+r"(dst.tiles[0][13].data[1].x), "+r"(dst.tiles[0][13].data[1].y), + "+r"(dst.tiles[0][13].data[2].x), "+r"(dst.tiles[0][13].data[2].y), + "+r"(dst.tiles[0][13].data[3].x), "+r"(dst.tiles[0][13].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %114, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n224k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \ + "%112, " \ + "%113, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y), + "+r"(dst.tiles[0][12].data[0].x), "+r"(dst.tiles[0][12].data[0].y), + "+r"(dst.tiles[0][12].data[1].x), "+r"(dst.tiles[0][12].data[1].y), + "+r"(dst.tiles[0][12].data[2].x), "+r"(dst.tiles[0][12].data[2].y), + "+r"(dst.tiles[0][12].data[3].x), "+r"(dst.tiles[0][12].data[3].y), + "+r"(dst.tiles[0][13].data[0].x), "+r"(dst.tiles[0][13].data[0].y), + "+r"(dst.tiles[0][13].data[1].x), "+r"(dst.tiles[0][13].data[1].y), + "+r"(dst.tiles[0][13].data[2].x), "+r"(dst.tiles[0][13].data[2].y), + "+r"(dst.tiles[0][13].data[3].x), "+r"(dst.tiles[0][13].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x240.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x240.impl new file mode 100644 index 0000000..aaf5336 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x240.impl @@ -0,0 +1,526 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %125, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n240k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119}, " \ + "{%120, %121, %122, %123}, " \ + "%124, " \ + "p, 1, %127, %126;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %125, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n240k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119}, " \ + "{%120, %121, %122, %123}, " \ + "%124, " \ + "p, 1, %127, %126;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %65, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n240k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59}, " \ + "{%60, %61, %62, %63}, " \ + "%64, " \ + "p, 1, %67, %66;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %122, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n240k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119}, " \ + "%120, " \ + "%121, " \ + "p, 1, %125, %123, %124;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %122, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n240k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119}, " \ + "%120, " \ + "%121, " \ + "p, 1, %125, %123, %124;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %62, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n240k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59}, " \ + "%60, " \ + "%61, " \ + "p, 1, %65, %63, %64;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x256.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x256.impl new file mode 100644 index 0000000..9c2e06e --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x256.impl @@ -0,0 +1,1604 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %133, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "{%128, %129, %130, %131}, " \ + "%132, " \ + "p, 1, %135, %134;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y), + "+f"(dst.tiles[0][15].data[0].x), "+f"(dst.tiles[0][15].data[0].y), + "+f"(dst.tiles[0][15].data[1].x), "+f"(dst.tiles[0][15].data[1].y), + "+f"(dst.tiles[0][15].data[2].x), "+f"(dst.tiles[0][15].data[2].y), + "+f"(dst.tiles[0][15].data[3].x), "+f"(dst.tiles[0][15].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %133, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "{%128, %129, %130, %131}, " \ + "%132, " \ + "p, 1, %135, %134;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y), + "+f"(dst.tiles[0][15].data[0].x), "+f"(dst.tiles[0][15].data[0].y), + "+f"(dst.tiles[0][15].data[1].x), "+f"(dst.tiles[0][15].data[1].y), + "+f"(dst.tiles[0][15].data[2].x), "+f"(dst.tiles[0][15].data[2].y), + "+f"(dst.tiles[0][15].data[3].x), "+f"(dst.tiles[0][15].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %69, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "{%64, %65, %66, %67}, " \ + "%68, " \ + "p, 1, %71, %70;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %133, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "{%128, %129, %130, %131}, " \ + "%132, " \ + "p, 1, %134;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y), + "+f"(dst.tiles[0][15].data[0].x), "+f"(dst.tiles[0][15].data[0].y), + "+f"(dst.tiles[0][15].data[1].x), "+f"(dst.tiles[0][15].data[1].y), + "+f"(dst.tiles[0][15].data[2].x), "+f"(dst.tiles[0][15].data[2].y), + "+f"(dst.tiles[0][15].data[3].x), "+f"(dst.tiles[0][15].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %133, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "{%128, %129, %130, %131}, " \ + "%132, " \ + "p, 1, %134;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y), + "+f"(dst.tiles[0][15].data[0].x), "+f"(dst.tiles[0][15].data[0].y), + "+f"(dst.tiles[0][15].data[1].x), "+f"(dst.tiles[0][15].data[1].y), + "+f"(dst.tiles[0][15].data[2].x), "+f"(dst.tiles[0][15].data[2].y), + "+f"(dst.tiles[0][15].data[3].x), "+f"(dst.tiles[0][15].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %69, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.f16.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "{%64, %65, %66, %67}, " \ + "%68, " \ + "p, 1, %70;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %69, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.f16.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "{%64, %65, %66, %67}, " \ + "%68, " \ + "p, 1, %70;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %133, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "{%128, %129, %130, %131}, " \ + "%132, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y), + "+r"(dst.tiles[0][12].data[0].x), "+r"(dst.tiles[0][12].data[0].y), + "+r"(dst.tiles[0][12].data[1].x), "+r"(dst.tiles[0][12].data[1].y), + "+r"(dst.tiles[0][12].data[2].x), "+r"(dst.tiles[0][12].data[2].y), + "+r"(dst.tiles[0][12].data[3].x), "+r"(dst.tiles[0][12].data[3].y), + "+r"(dst.tiles[0][13].data[0].x), "+r"(dst.tiles[0][13].data[0].y), + "+r"(dst.tiles[0][13].data[1].x), "+r"(dst.tiles[0][13].data[1].y), + "+r"(dst.tiles[0][13].data[2].x), "+r"(dst.tiles[0][13].data[2].y), + "+r"(dst.tiles[0][13].data[3].x), "+r"(dst.tiles[0][13].data[3].y), + "+r"(dst.tiles[0][14].data[0].x), "+r"(dst.tiles[0][14].data[0].y), + "+r"(dst.tiles[0][14].data[1].x), "+r"(dst.tiles[0][14].data[1].y), + "+r"(dst.tiles[0][14].data[2].x), "+r"(dst.tiles[0][14].data[2].y), + "+r"(dst.tiles[0][14].data[3].x), "+r"(dst.tiles[0][14].data[3].y), + "+r"(dst.tiles[0][15].data[0].x), "+r"(dst.tiles[0][15].data[0].y), + "+r"(dst.tiles[0][15].data[1].x), "+r"(dst.tiles[0][15].data[1].y), + "+r"(dst.tiles[0][15].data[2].x), "+r"(dst.tiles[0][15].data[2].y), + "+r"(dst.tiles[0][15].data[3].x), "+r"(dst.tiles[0][15].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %133, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "{%128, %129, %130, %131}, " \ + "%132, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y), + "+r"(dst.tiles[0][12].data[0].x), "+r"(dst.tiles[0][12].data[0].y), + "+r"(dst.tiles[0][12].data[1].x), "+r"(dst.tiles[0][12].data[1].y), + "+r"(dst.tiles[0][12].data[2].x), "+r"(dst.tiles[0][12].data[2].y), + "+r"(dst.tiles[0][12].data[3].x), "+r"(dst.tiles[0][12].data[3].y), + "+r"(dst.tiles[0][13].data[0].x), "+r"(dst.tiles[0][13].data[0].y), + "+r"(dst.tiles[0][13].data[1].x), "+r"(dst.tiles[0][13].data[1].y), + "+r"(dst.tiles[0][13].data[2].x), "+r"(dst.tiles[0][13].data[2].y), + "+r"(dst.tiles[0][13].data[3].x), "+r"(dst.tiles[0][13].data[3].y), + "+r"(dst.tiles[0][14].data[0].x), "+r"(dst.tiles[0][14].data[0].y), + "+r"(dst.tiles[0][14].data[1].x), "+r"(dst.tiles[0][14].data[1].y), + "+r"(dst.tiles[0][14].data[2].x), "+r"(dst.tiles[0][14].data[2].y), + "+r"(dst.tiles[0][14].data[3].x), "+r"(dst.tiles[0][14].data[3].y), + "+r"(dst.tiles[0][15].data[0].x), "+r"(dst.tiles[0][15].data[0].y), + "+r"(dst.tiles[0][15].data[1].x), "+r"(dst.tiles[0][15].data[1].y), + "+r"(dst.tiles[0][15].data[2].x), "+r"(dst.tiles[0][15].data[2].y), + "+r"(dst.tiles[0][15].data[3].x), "+r"(dst.tiles[0][15].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %130, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "%128, " \ + "%129, " \ + "p, 1, %133, %131, %132;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y), + "+f"(dst.tiles[0][15].data[0].x), "+f"(dst.tiles[0][15].data[0].y), + "+f"(dst.tiles[0][15].data[1].x), "+f"(dst.tiles[0][15].data[1].y), + "+f"(dst.tiles[0][15].data[2].x), "+f"(dst.tiles[0][15].data[2].y), + "+f"(dst.tiles[0][15].data[3].x), "+f"(dst.tiles[0][15].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %130, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "%128, " \ + "%129, " \ + "p, 1, %133, %131, %132;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y), + "+f"(dst.tiles[0][15].data[0].x), "+f"(dst.tiles[0][15].data[0].y), + "+f"(dst.tiles[0][15].data[1].x), "+f"(dst.tiles[0][15].data[1].y), + "+f"(dst.tiles[0][15].data[2].x), "+f"(dst.tiles[0][15].data[2].y), + "+f"(dst.tiles[0][15].data[3].x), "+f"(dst.tiles[0][15].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %66, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "%64, " \ + "%65, " \ + "p, 1, %69, %67, %68;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %130, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "%128, " \ + "%129, " \ + "p, 1, %131;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y), + "+f"(dst.tiles[0][15].data[0].x), "+f"(dst.tiles[0][15].data[0].y), + "+f"(dst.tiles[0][15].data[1].x), "+f"(dst.tiles[0][15].data[1].y), + "+f"(dst.tiles[0][15].data[2].x), "+f"(dst.tiles[0][15].data[2].y), + "+f"(dst.tiles[0][15].data[3].x), "+f"(dst.tiles[0][15].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %130, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.f32.e5m.e5m " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "%128, " \ + "%129, " \ + "p, 1, %131;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y), + "+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y), + "+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y), + "+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y), + "+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y), + "+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y), + "+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y), + "+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y), + "+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y), + "+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y), + "+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y), + "+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y), + "+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y), + "+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y), + "+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y), + "+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y), + "+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y), + "+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y), + "+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y), + "+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y), + "+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y), + "+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y), + "+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y), + "+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y), + "+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y), + "+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y), + "+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y), + "+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y), + "+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y), + "+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y), + "+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y), + "+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y), + "+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y), + "+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y), + "+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y), + "+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y), + "+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y), + "+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y), + "+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y), + "+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y), + "+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y), + "+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y), + "+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y), + "+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y), + "+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y), + "+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y), + "+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y), + "+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y), + "+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y), + "+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y), + "+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y), + "+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y), + "+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y), + "+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y), + "+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y), + "+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y), + "+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y), + "+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y), + "+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y), + "+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y), + "+f"(dst.tiles[0][15].data[0].x), "+f"(dst.tiles[0][15].data[0].y), + "+f"(dst.tiles[0][15].data[1].x), "+f"(dst.tiles[0][15].data[1].y), + "+f"(dst.tiles[0][15].data[2].x), "+f"(dst.tiles[0][15].data[2].y), + "+f"(dst.tiles[0][15].data[3].x), "+f"(dst.tiles[0][15].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %66, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.f16.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "%64, " \ + "%65, " \ + "p, 1, %67;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %66, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.f16.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \ + "%64, " \ + "%65, " \ + "p, 1, %67;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][10].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][11].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][12].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][13].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][14].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][15].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %130, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "%128, " \ + "%129, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y), + "+r"(dst.tiles[0][12].data[0].x), "+r"(dst.tiles[0][12].data[0].y), + "+r"(dst.tiles[0][12].data[1].x), "+r"(dst.tiles[0][12].data[1].y), + "+r"(dst.tiles[0][12].data[2].x), "+r"(dst.tiles[0][12].data[2].y), + "+r"(dst.tiles[0][12].data[3].x), "+r"(dst.tiles[0][12].data[3].y), + "+r"(dst.tiles[0][13].data[0].x), "+r"(dst.tiles[0][13].data[0].y), + "+r"(dst.tiles[0][13].data[1].x), "+r"(dst.tiles[0][13].data[1].y), + "+r"(dst.tiles[0][13].data[2].x), "+r"(dst.tiles[0][13].data[2].y), + "+r"(dst.tiles[0][13].data[3].x), "+r"(dst.tiles[0][13].data[3].y), + "+r"(dst.tiles[0][14].data[0].x), "+r"(dst.tiles[0][14].data[0].y), + "+r"(dst.tiles[0][14].data[1].x), "+r"(dst.tiles[0][14].data[1].y), + "+r"(dst.tiles[0][14].data[2].x), "+r"(dst.tiles[0][14].data[2].y), + "+r"(dst.tiles[0][14].data[3].x), "+r"(dst.tiles[0][14].data[3].y), + "+r"(dst.tiles[0][15].data[0].x), "+r"(dst.tiles[0][15].data[0].y), + "+r"(dst.tiles[0][15].data[1].x), "+r"(dst.tiles[0][15].data[1].y), + "+r"(dst.tiles[0][15].data[2].x), "+r"(dst.tiles[0][15].data[2].y), + "+r"(dst.tiles[0][15].data[3].x), "+r"(dst.tiles[0][15].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %130, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n256k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, " \ + "%128, " \ + "%129, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y), + "+r"(dst.tiles[0][6].data[0].x), "+r"(dst.tiles[0][6].data[0].y), + "+r"(dst.tiles[0][6].data[1].x), "+r"(dst.tiles[0][6].data[1].y), + "+r"(dst.tiles[0][6].data[2].x), "+r"(dst.tiles[0][6].data[2].y), + "+r"(dst.tiles[0][6].data[3].x), "+r"(dst.tiles[0][6].data[3].y), + "+r"(dst.tiles[0][7].data[0].x), "+r"(dst.tiles[0][7].data[0].y), + "+r"(dst.tiles[0][7].data[1].x), "+r"(dst.tiles[0][7].data[1].y), + "+r"(dst.tiles[0][7].data[2].x), "+r"(dst.tiles[0][7].data[2].y), + "+r"(dst.tiles[0][7].data[3].x), "+r"(dst.tiles[0][7].data[3].y), + "+r"(dst.tiles[0][8].data[0].x), "+r"(dst.tiles[0][8].data[0].y), + "+r"(dst.tiles[0][8].data[1].x), "+r"(dst.tiles[0][8].data[1].y), + "+r"(dst.tiles[0][8].data[2].x), "+r"(dst.tiles[0][8].data[2].y), + "+r"(dst.tiles[0][8].data[3].x), "+r"(dst.tiles[0][8].data[3].y), + "+r"(dst.tiles[0][9].data[0].x), "+r"(dst.tiles[0][9].data[0].y), + "+r"(dst.tiles[0][9].data[1].x), "+r"(dst.tiles[0][9].data[1].y), + "+r"(dst.tiles[0][9].data[2].x), "+r"(dst.tiles[0][9].data[2].y), + "+r"(dst.tiles[0][9].data[3].x), "+r"(dst.tiles[0][9].data[3].y), + "+r"(dst.tiles[0][10].data[0].x), "+r"(dst.tiles[0][10].data[0].y), + "+r"(dst.tiles[0][10].data[1].x), "+r"(dst.tiles[0][10].data[1].y), + "+r"(dst.tiles[0][10].data[2].x), "+r"(dst.tiles[0][10].data[2].y), + "+r"(dst.tiles[0][10].data[3].x), "+r"(dst.tiles[0][10].data[3].y), + "+r"(dst.tiles[0][11].data[0].x), "+r"(dst.tiles[0][11].data[0].y), + "+r"(dst.tiles[0][11].data[1].x), "+r"(dst.tiles[0][11].data[1].y), + "+r"(dst.tiles[0][11].data[2].x), "+r"(dst.tiles[0][11].data[2].y), + "+r"(dst.tiles[0][11].data[3].x), "+r"(dst.tiles[0][11].data[3].y), + "+r"(dst.tiles[0][12].data[0].x), "+r"(dst.tiles[0][12].data[0].y), + "+r"(dst.tiles[0][12].data[1].x), "+r"(dst.tiles[0][12].data[1].y), + "+r"(dst.tiles[0][12].data[2].x), "+r"(dst.tiles[0][12].data[2].y), + "+r"(dst.tiles[0][12].data[3].x), "+r"(dst.tiles[0][12].data[3].y), + "+r"(dst.tiles[0][13].data[0].x), "+r"(dst.tiles[0][13].data[0].y), + "+r"(dst.tiles[0][13].data[1].x), "+r"(dst.tiles[0][13].data[1].y), + "+r"(dst.tiles[0][13].data[2].x), "+r"(dst.tiles[0][13].data[2].y), + "+r"(dst.tiles[0][13].data[3].x), "+r"(dst.tiles[0][13].data[3].y), + "+r"(dst.tiles[0][14].data[0].x), "+r"(dst.tiles[0][14].data[0].y), + "+r"(dst.tiles[0][14].data[1].x), "+r"(dst.tiles[0][14].data[1].y), + "+r"(dst.tiles[0][14].data[2].x), "+r"(dst.tiles[0][14].data[2].y), + "+r"(dst.tiles[0][14].data[3].x), "+r"(dst.tiles[0][14].data[3].y), + "+r"(dst.tiles[0][15].data[0].x), "+r"(dst.tiles[0][15].data[0].y), + "+r"(dst.tiles[0][15].data[1].x), "+r"(dst.tiles[0][15].data[1].y), + "+r"(dst.tiles[0][15].data[2].x), "+r"(dst.tiles[0][15].data[2].y), + "+r"(dst.tiles[0][15].data[3].x), "+r"(dst.tiles[0][15].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x32.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x32.impl new file mode 100644 index 0000000..205773f --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x32.impl @@ -0,0 +1,564 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %21, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "{%16, %17, %18, %19}, " \ + "%20, " \ + "p, 1, %23, %22;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %21, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "{%16, %17, %18, %19}, " \ + "%20, " \ + "p, 1, %23, %22;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %13, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7}, " \ + "{%8, %9, %10, %11}, " \ + "%12, " \ + "p, 1, %15, %14;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %21, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "{%16, %17, %18, %19}, " \ + "%20, " \ + "p, 1, %22;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %21, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "{%16, %17, %18, %19}, " \ + "%20, " \ + "p, 1, %22;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %13, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.f16.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7}, " \ + "{%8, %9, %10, %11}, " \ + "%12, " \ + "p, 1, %14;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %21, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "{%16, %17, %18, %19}, " \ + "%20, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %21, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "{%16, %17, %18, %19}, " \ + "%20, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %18, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "%16, " \ + "%17, " \ + "p, 1, %21, %19, %20;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %18, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "%16, " \ + "%17, " \ + "p, 1, %21, %19, %20;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %10, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7}, " \ + "%8, " \ + "%9, " \ + "p, 1, %13, %11, %12;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %18, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "%16, " \ + "%17, " \ + "p, 1, %19;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %18, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "%16, " \ + "%17, " \ + "p, 1, %19;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %10, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.f16.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7}, " \ + "%8, " \ + "%9, " \ + "p, 1, %11;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %10, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.f16.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7}, " \ + "%8, " \ + "%9, " \ + "p, 1, %11;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %18, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "%16, " \ + "%17, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %18, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n32k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "%16, " \ + "%17, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x48.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x48.impl new file mode 100644 index 0000000..702ee70 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x48.impl @@ -0,0 +1,238 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %29, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n48k16.f32.bf16.bf16 " \ + "{%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, " \ + "p, 1, %31, %30;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %29, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n48k16.f32.f16.f16 " \ + "{%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, " \ + "p, 1, %31, %30;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %17, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n48k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11}, " \ + "{%12, %13, %14, %15}, " \ + "%16, " \ + "p, 1, %19, %18;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %26, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n48k16.f32.bf16.bf16 " \ + "{%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, " \ + "p, 1, %29, %27, %28;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %26, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n48k16.f32.f16.f16 " \ + "{%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, " \ + "p, 1, %29, %27, %28;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %14, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n48k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11}, " \ + "%12, " \ + "%13, " \ + "p, 1, %17, %15, %16;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x64.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x64.impl new file mode 100644 index 0000000..ff7983d --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x64.impl @@ -0,0 +1,739 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %37, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "{%32, %33, %34, %35}, " \ + "%36, " \ + "p, 1, %39, %38;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %37, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "{%32, %33, %34, %35}, " \ + "%36, " \ + "p, 1, %39, %38;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %21, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "{%16, %17, %18, %19}, " \ + "%20, " \ + "p, 1, %23, %22;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %37, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "{%32, %33, %34, %35}, " \ + "%36, " \ + "p, 1, %38;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %37, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "{%32, %33, %34, %35}, " \ + "%36, " \ + "p, 1, %38;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %21, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.f16.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "{%16, %17, %18, %19}, " \ + "%20, " \ + "p, 1, %22;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %21, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.f16.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "{%16, %17, %18, %19}, " \ + "%20, " \ + "p, 1, %22;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %37, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "{%32, %33, %34, %35}, " \ + "%36, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %37, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "{%32, %33, %34, %35}, " \ + "%36, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %34, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "%32, " \ + "%33, " \ + "p, 1, %37, %35, %36;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %34, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "%32, " \ + "%33, " \ + "p, 1, %37, %35, %36;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %18, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "%16, " \ + "%17, " \ + "p, 1, %21, %19, %20;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %34, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "%32, " \ + "%33, " \ + "p, 1, %35;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %34, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "%32, " \ + "%33, " \ + "p, 1, %35;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %18, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.f16.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "%16, " \ + "%17, " \ + "p, 1, %19;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %18, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.f16.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \ + "%16, " \ + "%17, " \ + "p, 1, %19;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %34, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "%32, " \ + "%33, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %34, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n64k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \ + "%32, " \ + "%33, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x80.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x80.impl new file mode 100644 index 0000000..7a33667 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x80.impl @@ -0,0 +1,286 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %45, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n80k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39}, " \ + "{%40, %41, %42, %43}, " \ + "%44, " \ + "p, 1, %47, %46;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %45, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n80k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39}, " \ + "{%40, %41, %42, %43}, " \ + "%44, " \ + "p, 1, %47, %46;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %25, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n80k16.f16.f16.f16 " \ + "{%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, " \ + "p, 1, %27, %26;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %42, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n80k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39}, " \ + "%40, " \ + "%41, " \ + "p, 1, %45, %43, %44;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %42, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n80k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39}, " \ + "%40, " \ + "%41, " \ + "p, 1, %45, %43, %44;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %22, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n80k16.f16.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19}, " \ + "%20, " \ + "%21, " \ + "p, 1, %25, %23, %24;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x96.impl b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x96.impl new file mode 100644 index 0000000..b7f955f --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x96.impl @@ -0,0 +1,887 @@ +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt_base & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %53, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "{%48, %49, %50, %51}, " \ + "%52, " \ + "p, 1, %55, %54;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %53, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "{%48, %49, %50, %51}, " \ + "%52, " \ + "p, 1, %55, %54;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %29, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k16.f16.f16.f16 " \ + "{%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, " \ + "p, 1, %31, %30;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %53, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "{%48, %49, %50, %51}, " \ + "%52, " \ + "p, 1, %54;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %53, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "{%48, %49, %50, %51}, " \ + "%52, " \ + "p, 1, %54;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %29, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.f16.e4m3.e4m3 " \ + "{%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, " \ + "p, 1, %30;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %29, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.f16.e5m2.e5m2 " \ + "{%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, " \ + "p, 1, %30;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d), + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %53, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "{%48, %49, %50, %51}, " \ + "%52, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %53, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "{%48, %49, %50, %51}, " \ + "%52, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y) + + : "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]), + "r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]), + + "l"(b_st_desc), + "r"(scale_d) + ); + } + + } + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ) { + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid type combination for WGMMA." + ); + static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option"); + + if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %50, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k16.f32.bf16.bf16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "%48, " \ + "%49, " \ + "p, 1, %53, %51, %52;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %50, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k16.f32.f16.f16 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "%48, " \ + "%49, " \ + "p, 1, %53, %51, %52;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %26, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k16.f16.f16.f16 " \ + "{%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, " \ + "p, 1, %29, %27, %28;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + "n"(trans_a), + "n"(trans_b), + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %50, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.f32.e4m3.e4m3 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "%48, " \ + "%49, " \ + "p, 1, %51;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %50, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.f32.e5m2.e5m2 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "%48, " \ + "%49, " \ + "p, 1, %51;\n" \ + "}\n" + + + : "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y), + "+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y), + "+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y), + "+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y), + "+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y), + "+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y), + "+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y), + "+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y), + "+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y), + "+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y), + "+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y), + "+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y), + "+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y), + "+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y), + "+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y), + "+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y), + "+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y), + "+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y), + "+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y), + "+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y), + "+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y), + "+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y), + "+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y), + "+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %26, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.f16.e4m3.e4m3 " \ + "{%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, " \ + "p, 1, %27;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %26, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.f16.e5m2.e5m2 " \ + "{%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, " \ + "p, 1, %27;\n" \ + "}\n" + + + : "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][0].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][1].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][2].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][3].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][4].data[3]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[0]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[1]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[2]), + "+r"(*(uint32_t*)&dst.tiles[0][5].data[3]) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d), + + + "n"(scale_b) + ); + } + + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %50, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.satfinite.s32.s8.s8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "%48, " \ + "%49, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + + else if constexpr (std::is_same_v && std::is_same_v) { + asm volatile ( + "{\n" + ".reg .pred p;\n" \ + "setp.ne.b32 p, %50, 0;\n" \ + "wgmma.mma_async.sync.aligned.m64n96k32.satfinite.s32.u8.u8 " \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \ + "%48, " \ + "%49, " \ + "p;\n" \ + "}\n" + + : "+r"(dst.tiles[0][0].data[0].x), "+r"(dst.tiles[0][0].data[0].y), + "+r"(dst.tiles[0][0].data[1].x), "+r"(dst.tiles[0][0].data[1].y), + "+r"(dst.tiles[0][0].data[2].x), "+r"(dst.tiles[0][0].data[2].y), + "+r"(dst.tiles[0][0].data[3].x), "+r"(dst.tiles[0][0].data[3].y), + "+r"(dst.tiles[0][1].data[0].x), "+r"(dst.tiles[0][1].data[0].y), + "+r"(dst.tiles[0][1].data[1].x), "+r"(dst.tiles[0][1].data[1].y), + "+r"(dst.tiles[0][1].data[2].x), "+r"(dst.tiles[0][1].data[2].y), + "+r"(dst.tiles[0][1].data[3].x), "+r"(dst.tiles[0][1].data[3].y), + "+r"(dst.tiles[0][2].data[0].x), "+r"(dst.tiles[0][2].data[0].y), + "+r"(dst.tiles[0][2].data[1].x), "+r"(dst.tiles[0][2].data[1].y), + "+r"(dst.tiles[0][2].data[2].x), "+r"(dst.tiles[0][2].data[2].y), + "+r"(dst.tiles[0][2].data[3].x), "+r"(dst.tiles[0][2].data[3].y), + "+r"(dst.tiles[0][3].data[0].x), "+r"(dst.tiles[0][3].data[0].y), + "+r"(dst.tiles[0][3].data[1].x), "+r"(dst.tiles[0][3].data[1].y), + "+r"(dst.tiles[0][3].data[2].x), "+r"(dst.tiles[0][3].data[2].y), + "+r"(dst.tiles[0][3].data[3].x), "+r"(dst.tiles[0][3].data[3].y), + "+r"(dst.tiles[0][4].data[0].x), "+r"(dst.tiles[0][4].data[0].y), + "+r"(dst.tiles[0][4].data[1].x), "+r"(dst.tiles[0][4].data[1].y), + "+r"(dst.tiles[0][4].data[2].x), "+r"(dst.tiles[0][4].data[2].y), + "+r"(dst.tiles[0][4].data[3].x), "+r"(dst.tiles[0][4].data[3].y), + "+r"(dst.tiles[0][5].data[0].x), "+r"(dst.tiles[0][5].data[0].y), + "+r"(dst.tiles[0][5].data[1].x), "+r"(dst.tiles[0][5].data[1].y), + "+r"(dst.tiles[0][5].data[2].x), "+r"(dst.tiles[0][5].data[2].y), + "+r"(dst.tiles[0][5].data[3].x), "+r"(dst.tiles[0][5].data[3].y) + + : "l"(a_st_desc), + "l"(b_st_desc), + + "r"(scale_d) + ); + } + } +}; diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/base.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/base.cuh new file mode 100644 index 0000000..1da1375 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/base.cuh @@ -0,0 +1,47 @@ +#pragma once + +#include "../../../../common/common.cuh" +#include "../../../../types/types.cuh" + +namespace kittens { +namespace detail { +namespace wgmma { + + +template +struct base { + template __device__ static inline void rt_st( + rt &dst, + const rt & a_rt, + const uint64_t b_st_desc, + int scale_d = 1 + ); + template __device__ static inline void st_st( + rt &dst, + const uint64_t a_st_desc, + const uint64_t b_st_desc, + int scale_d = 1 + ); +}; + + +#include "64x16.impl" +#include "64x32.impl" +#include "64x48.impl" +#include "64x64.impl" +#include "64x80.impl" +#include "64x96.impl" +#include "64x112.impl" +#include "64x128.impl" +#include "64x144.impl" +#include "64x160.impl" +#include "64x176.impl" +#include "64x192.impl" +#include "64x208.impl" +#include "64x224.impl" +#include "64x240.impl" +#include "64x256.impl" + +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/mma.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/mma.cuh new file mode 100644 index 0000000..2010402 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/mma.cuh @@ -0,0 +1,17 @@ + + + + + + +#include "warp.cuh" + + +#if defined(KITTENS_SM90) +#include "warpgroup.cuh" +#endif + + +#ifdef KITTENS_SM10X +#include "tcgen05.cuh" +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/tcgen05.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/tcgen05.cuh new file mode 100644 index 0000000..0c5db67 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/tcgen05.cuh @@ -0,0 +1,236 @@ + + + + +template +__device__ static inline void mma(D &d, const A &a, const B &b, semaphore &sem) { + if(laneid() == 0) ::kittens::mma(d, a, b, sem); +} +template +__device__ static inline void mma2(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm2(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mma(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + if(laneid() == 0) ::kittens::mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mma2(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mm(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mm2(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma2(d, a, b, sa, sb, sem); +} + +template +__device__ static inline void mma_AB(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mma2_AB(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mma_ABt(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mma2_ABt(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mma_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mma2_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma2(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mma_AtB(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mma2_AtB(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mma_AtBt(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mma2_AtBt(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} + +template +__device__ static inline void mm_AB(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm2_AB(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mm_ABt(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm2_ABt(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mm_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mm2_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma2(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mm_AtB(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm2_AtB(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mm_AtBt(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm2_AtBt(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} + + + + +template +__device__ static inline void mma(D &d, const A &a, const B &b) { + if(laneid() == 0) ::kittens::mma(d, a, b); +} +template +__device__ static inline void mma2(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm2(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mma(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + if(laneid() == 0) ::kittens::mma(d, a, b, sa, sb); +} +template +__device__ static inline void mma2(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma(d, a, b, sa, sb); +} +template +__device__ static inline void mm(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma(d, a, b, sa, sb); +} +template +__device__ static inline void mm2(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma2(d, a, b, sa, sb); +} + +template +__device__ static inline void mma_AB(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mma2_AB(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mma_ABt(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mma2_ABt(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mma_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma(d, a, b, sa, sb); +} +template +__device__ static inline void mma2_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma2(d, a, b, sa, sb); +} +template +__device__ static inline void mma_AtB(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mma2_AtB(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mma_AtBt(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mma2_AtBt(D &d, const A &a, const B &b) { + mma2(d, a, b); +} + +template +__device__ static inline void mm_AB(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm2_AB(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mm_ABt(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm2_ABt(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mm_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma(d, a, b, sa, sb); +} +template +__device__ static inline void mm2_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma2(d, a, b, sa, sb); +} +template +__device__ static inline void mm_AtB(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm2_AtB(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mm_AtBt(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm2_AtBt(D &d, const A &a, const B &b) { + mma2(d, a, b); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/warp.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/warp.cuh new file mode 100644 index 0000000..a94c57f --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/warp.cuh @@ -0,0 +1,1057 @@ + + + + + + + + + + + + + + + + + + + + + + +__device__ static inline void hmma16816( float2 &d0, float2 &d1, + const bf16_2 &a0, const bf16_2 &a1, const bf16_2 &a2, const bf16_2 &a3, + const bf16_2 &b0, const bf16_2 &b1, + const float2 &c0, const float2 &c1 ) { + asm volatile( + + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " \ + "{%0, %1, %2, %3}, " \ + "{%4, %5, %6, %7}, " \ + "{%8, %9}, " \ + "{%10, %11, %12, %13};" + + + : "+f"(d0.x), "+f"(d0.y), + "+f"(d1.x), "+f"(d1.y) + + + : "r"(*(uint32_t*)(&a0)), "r"(*(uint32_t*)(&a1)), + "r"(*(uint32_t*)(&a2)), "r"(*(uint32_t*)(&a3)), + + + "r"(*(uint32_t*)(&b0)), "r"(*(uint32_t*)(&b1)), + + + "f"(c0.x), "f"(c0.y), + "f"(c1.x), "f"(c1.y) + ); +} + + + + + + + + + + + + + + + + + +__device__ static inline void hmma16816( float2 &d0, float2 &d1, + const half_2 &a0, const half_2 &a1, const half_2 &a2, const half_2 &a3, + const half_2 &b0, const half_2 &b1, + const float2 &c0, const float2 &c1 ) { + asm volatile( + + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " \ + "{%0, %1, %2, %3}, " \ + "{%4, %5, %6, %7}, " \ + "{%8, %9}, " \ + "{%10, %11, %12, %13};" + + + : "+f"(d0.x), "+f"(d0.y), + "+f"(d1.x), "+f"(d1.y) + + + : "r"(*(uint32_t*)(&a0)), "r"(*(uint32_t*)(&a1)), + "r"(*(uint32_t*)(&a2)), "r"(*(uint32_t*)(&a3)), + + + "r"(*(uint32_t*)(&b0)), "r"(*(uint32_t*)(&b1)), + + + "f"(c0.x), "f"(c0.y), + "f"(c1.x), "f"(c1.y) + ); +} + + + + + + + + + + + + + + + + + +__device__ static inline void hmma16816( half_2 &d0, half_2 &d1, + const half_2 &a0, const half_2 &a1, const half_2 &a2, const half_2 &a3, + const half_2 &b0, const half_2 &b1, + const half_2 &c0, const half_2 &c1 ) { + asm volatile( + + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " \ + "{%0, %1}, " \ + "{%2, %3, %4, %5}, " \ + "{%6, %7}, " \ + "{%8, %9};" + + + : "=r"(*(uint32_t*)(&d0)), "=r"(*(uint32_t*)(&d1)) + + + : "r"(*(uint32_t*)(&a0)), "r"(*(uint32_t*)(&a1)), + "r"(*(uint32_t*)(&a2)), "r"(*(uint32_t*)(&a3)), + + + "r"(*(uint32_t*)(&b0)), "r"(*(uint32_t*)(&b1)), + + + "r"(*(uint32_t*)(&c0)), "r"(*(uint32_t*)(&c1)) + ); +} + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + + + + + + + + + + + +__device__ static inline void hmma16816( float2 &d0, float2 &d1, + const fp8e4m3_4 &a0, const fp8e4m3_4 &a1, + const fp8e4m3_4 &a2, const fp8e4m3_4 &a3, + const fp8e4m3_4 &b0, const fp8e4m3_4 &b1, + const float2 &c0, const float2 &c1) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, %1, %2, %3}, " + "{%4, %5, %6, %7}, " + "{%8, %9}, " + "{%10, %11, %12, %13};" + + + : "+f"(d0.x), "+f"(d0.y), + "+f"(d1.x), "+f"(d1.y) + + + : "r"(*(uint32_t*)(&a0)), "r"(*(uint32_t*)(&a1)), + "r"(*(uint32_t*)(&a2)), "r"(*(uint32_t*)(&a3)), + + + "r"(*(uint32_t*)(&b0)), "r"(*(uint32_t*)(&b1)), + + + "f"(c0.x), "f"(c0.y), + "f"(c1.x), "f"(c1.y) + ); +} +#endif + + + + + + +__device__ static inline void imma16832(int2 &d0, int2 &d1, + const int8_4 &a0, const int8_4 &a1, const int8_4 &a2, const int8_4 &a3, + const int8_4 &b0, const int8_4 &b1, + const int2 &c0, const int2 &c1) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.satfinite.s32.s8.s8.s32 " + "{%0, %1, %2, %3}, " + "{%4, %5, %6, %7}, " + "{%8, %9}, " + "{%10, %11, %12, %13};" + : "=r"(d0.x), "=r"(d0.y), + "=r"(d1.x), "=r"(d1.y) + : "r"(*(uint32_t*)(&a0)), "r"(*(uint32_t*)(&a1)), + "r"(*(uint32_t*)(&a2)), "r"(*(uint32_t*)(&a3)), + "r"(*(uint32_t*)(&b0)), "r"(*(uint32_t*)(&b1)), + "r"(c0.x), "r"(c0.y), + "r"(c1.x), "r"(c1.y) + ); +} +__device__ static inline void imma16832(int2 &d0, int2 &d1, + const int8_4 &a0, const int8_4 &a1, const int8_4 &a2, const int8_4 &a3, + const uint8_4 &b0, const uint8_4 &b1, + const int2 &c0, const int2 &c1) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.satfinite.s32.s8.u8.s32 " + "{%0, %1, %2, %3}, " + "{%4, %5, %6, %7}, " + "{%8, %9}, " + "{%10, %11, %12, %13};" + : "=r"(d0.x), "=r"(d0.y), + "=r"(d1.x), "=r"(d1.y) + : "r"(*(uint32_t*)(&a0)), "r"(*(uint32_t*)(&a1)), + "r"(*(uint32_t*)(&a2)), "r"(*(uint32_t*)(&a3)), + "r"(*(uint32_t*)(&b0)), "r"(*(uint32_t*)(&b1)), + "r"(c0.x), "r"(c0.y), + "r"(c1.x), "r"(c1.y) + ); +} +__device__ static inline void imma16832(int2 &d0, int2 &d1, + const uint8_4 &a0, const uint8_4 &a1, const uint8_4 &a2, const uint8_4 &a3, + const int8_4 &b0, const int8_4 &b1, + const int2 &c0, const int2 &c1) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.satfinite.s32.u8.s8.s32 " + "{%0, %1, %2, %3}, " + "{%4, %5, %6, %7}, " + "{%8, %9}, " + "{%10, %11, %12, %13};" + : "=r"(d0.x), "=r"(d0.y), + "=r"(d1.x), "=r"(d1.y) + : "r"(*(uint32_t*)(&a0)), "r"(*(uint32_t*)(&a1)), + "r"(*(uint32_t*)(&a2)), "r"(*(uint32_t*)(&a3)), + "r"(*(uint32_t*)(&b0)), "r"(*(uint32_t*)(&b1)), + "r"(c0.x), "r"(c0.y), + "r"(c1.x), "r"(c1.y) + ); +} +__device__ static inline void imma16832(int2 &d0, int2 &d1, + const uint8_4 &a0, const uint8_4 &a1, const uint8_4 &a2, const uint8_4 &a3, + const uint8_4 &b0, const uint8_4 &b1, + const int2 &c0, const int2 &c1) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.satfinite.s32.u8.u8.s32 " + "{%0, %1, %2, %3}, " + "{%4, %5, %6, %7}, " + "{%8, %9}, " + "{%10, %11, %12, %13};" + : "=r"(d0.x), "=r"(d0.y), + "=r"(d1.x), "=r"(d1.y) + : "r"(*(uint32_t*)(&a0)), "r"(*(uint32_t*)(&a1)), + "r"(*(uint32_t*)(&a2)), "r"(*(uint32_t*)(&a3)), + "r"(*(uint32_t*)(&b0)), "r"(*(uint32_t*)(&b1)), + "r"(c0.x), "r"(c0.y), + "r"(c1.x), "r"(c1.y) + ); +} + + + + + + + + + + + + +__device__ static inline void mma_AB_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} + + + + + + + + + + + + +__device__ static inline void mma_AB_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + + + + +__device__ static inline void mma_AB_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} +#endif + + + + + + + + + + + +__device__ static inline void mma_AB_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} + + + + + + + + + + + +__device__ static inline void mma_ABt_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} + + + + + + + + + + + + +__device__ static inline void mma_ABt_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + + + + +__device__ static inline void mma_ABt_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} +#endif +template +__device__ static inline void mma_ABt_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + static_assert((std::is_same_v || std::is_same_v) && + (std::is_same_v || std::is_same_v)); + imma16832( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + imma16832( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} + + + + + + + + + + + + + +__device__ static inline void mma_AtB_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} + + + + + + + + + + + + +__device__ static inline void mma_AtB_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + + + + +__device__ static inline void mma_AtB_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} +#endif + + + + + + + + + + + + +__device__ static inline void mma_AtBt_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} + + + + + + + + + + + + +__device__ static inline void mma_AtBt_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + + + + +__device__ static inline void mma_AtBt_base(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c) { + hmma16816( + d.data[0], d.data[1], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[0], b.data[2], + c.data[0], c.data[1] + ); + hmma16816( + d.data[2], d.data[3], + a.data[0], a.data[1], a.data[2], a.data[3], + b.data[1], b.data[3], + c.data[2], c.data[3] + ); +} +#endif + + + + + + + + + + + + + + + +template +__device__ static inline void mma_AB(D &d, + const A &a, + const B &b, + const C &c) { + KITTENS_CHECK_WARP + static_assert(D::rows == A::rows && D::cols == B::cols); + static_assert(A::cols == B::rows); + static_assert(D::rows == C::rows && D::cols == C::cols); + #if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert( + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) + ); + #else + static_assert( + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) + ); + #endif + #pragma unroll + for(int n = 0; n < D::height; n++) { + #pragma unroll + for(int m = 0; m < D::width; m++) { + mma_AB_base( + d.tiles[n][m], + a.tiles[n][0], + b.tiles[0][m], + c.tiles[n][m] + ); + #pragma unroll + for(int k = 1; k < A::width; k++) { + mma_AB_base( + d.tiles[n][m], + a.tiles[n][k], + b.tiles[k][m], + d.tiles[n][m] + ); + } + } + } +} + + + + + + + + + + + + + + +template +__device__ static inline void mma_ABt(D &d, + const A &a, + const B &b, + const C &c) { + KITTENS_CHECK_WARP + static_assert(D::rows == A::rows && D::cols == B::rows); + static_assert(A::cols == B::cols); + static_assert(D::rows == C::rows && D::cols == C::cols); + #if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert( + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && + (std::is_same_v || std::is_same_v) && + (std::is_same_v || std::is_same_v) && + std::is_same_v) + ); + #else + static_assert( + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && + (std::is_same_v || std::is_same_v) && + (std::is_same_v || std::is_same_v) && + std::is_same_v) + ); + #endif + #pragma unroll + for(int n = 0; n < D::height; n++) { + #pragma unroll + for(int m = 0; m < D::width; m++) { + mma_ABt_base( + d.tiles[n][m], + a.tiles[n][0], + b.tiles[m][0], + c.tiles[n][m] + ); + #pragma unroll + for(int k = 1; k < A::width; k++) { + mma_ABt_base( + d.tiles[n][m], + a.tiles[n][k], + b.tiles[m][k], + d.tiles[n][m] + ); + } + } + } +} + + + + + + + + + + + + + + +template +__device__ static inline void mma_AtB(D &d, + const A &a, + const B &b, + const C &c) { + KITTENS_CHECK_WARP + static_assert(D::rows == A::cols && D::cols == B::cols); + static_assert(A::rows == B::rows); + static_assert(D::rows == C::rows && D::cols == C::cols); + #if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert( + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) + ); + #else + static_assert( + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) + ); + #endif + #pragma unroll + for(int n = 0; n < D::height; n++) { + #pragma unroll + for(int m = 0; m < D::width; m++) { + mma_AtB_base( + d.tiles[n][m], + a.tiles[0][n], + b.tiles[0][m], + c.tiles[n][m] + ); + #pragma unroll + for(int k = 1; k < A::height; k++) { + mma_AtB_base( + d.tiles[n][m], + a.tiles[k][n], + b.tiles[k][m], + d.tiles[n][m] + ); + } + } + } +} + + + + + + + + + + + + + + +template +__device__ static inline void mma_AtBt(D &d, + const A &a, + const B &b, + const C &c) { + KITTENS_CHECK_WARP + static_assert(D::rows == A::cols && D::cols == B::rows); + static_assert(A::rows == B::cols); + static_assert(D::rows == C::rows && D::cols == C::cols); + #if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert( + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) + ); + #else + static_assert( + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) + ); + #endif + #pragma unroll + for(int n = 0; n < D::height; n++) { + #pragma unroll + for(int m = 0; m < D::width; m++) { + mma_AtBt_base( + d.tiles[n][m], + a.tiles[0][n], + b.tiles[m][0], + c.tiles[n][m] + ); + #pragma unroll + for(int k = 1; k < A::height; k++) { + mma_AtBt_base( + d.tiles[n][m], + a.tiles[k][n], + b.tiles[m][k], + d.tiles[n][m] + ); + } + } + } +} + +template +__device__ static inline void mma(D &d, + const A &a, + const B &b, + const C &c) { + KITTENS_CHECK_WARP + if constexpr(trans_A == transpose::T) { + if constexpr(trans_B == transpose::T) { + mma_AtBt(d, a, b, c); + } else { + mma_AtB(d, a, b, c); + } + } else { + if constexpr(trans_B == transpose::T) { + mma_ABt(d, a, b, c); + } else { + mma_AB(d, a, b, c); + } + } +} +template +__device__ static inline C mma(const A &a, + const B &b, + const C &c) { + KITTENS_CHECK_WARP + C d; + if constexpr(trans_A == transpose::T) { + if constexpr(trans_B == transpose::T) { + mma_AtBt(d, a, b, c); + } else { + mma_AtB(d, a, b, c); + } + } else { + if constexpr(trans_B == transpose::T) { + mma_ABt(d, a, b, c); + } else { + mma_AB(d, a, b, c); + } + } + return d; +} + + + + + + + + + + + + + + + + + + + + + + + +template +__device__ static inline void mma_AB(crt_hf &d, + const crt_hf &a, + const crt_hf &b, + const crt_hf &c) { + KITTENS_CHECK_WARP + + + ::kittens::group<1>::copy(d.real, c.real); + ::kittens::group<1>::copy(d.imag, c.imag); + + + rt_hf tmp; + + constexpr half factor = std::bit_cast<__half>(uint16_t(0xFB80)); + ::kittens::group<1>::mul(tmp, a.imag, factor); + mma_AB(d.real, a.real, b.real, d.real); + mma_AB(d.real, tmp, b.imag, d.real); + + mma_AB(d.imag, a.real, b.imag, d.imag); + mma_AB(d.imag, a.imag, b.real, d.imag); +} + + + + + + + + + + + + + + +template +__device__ static inline void mma_AB(crt_fl &d, + const crt_bf &a, + const crt_bf &b, + const crt_fl &c) { + KITTENS_CHECK_WARP + + + ::kittens::group<1>::copy(d.real, c.real); + ::kittens::group<1>::copy(d.imag, c.imag); + + + kittens::rt_bf tmp; + + constexpr bf16 factor = std::bit_cast<__nv_bfloat16>(uint16_t(0xBF80)); + ::kittens::group<1>::mul(tmp, a.imag, factor); + mma_AB(d.real, a.real, b.real, d.real); + mma_AB(d.real, tmp, b.imag, d.real); + + mma_AB(d.imag, a.real, b.imag, d.imag); + mma_AB(d.imag, a.imag, b.real, d.imag); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/warpgroup.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/warpgroup.cuh new file mode 100644 index 0000000..5b5a0a3 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/warpgroup.cuh @@ -0,0 +1,1167 @@ + + + + + + + + + + + + + + + + + + + + + +template +__device__ static inline void mma_fence(D &dst) { + KITTENS_CHECK_WARPGROUP + #pragma unroll + for(int i = 0; i < D::height; i++) { + #pragma unroll + for(int j = 0; j < D::width; j++) { + #pragma unroll + for(int k = 0; k < D::packed_per_tile; k++) { + if constexpr(std::is_same_v) { + asm volatile("" : "+f"(dst.tiles[i][j].data[k].x) :: "memory"); + asm volatile("" : "+f"(dst.tiles[i][j].data[k].y) :: "memory"); + } else if constexpr(std::is_same_v) { + asm volatile("" : "+r"(dst.tiles[i][j].data[k].x) :: "memory"); + asm volatile("" : "+r"(dst.tiles[i][j].data[k].y) :: "memory"); + } else { + asm volatile("" : "+r"(*(uint32_t*)&dst.tiles[i][j].data[k]) :: "memory"); + } + } + } + } + asm volatile ("wgmma.fence.sync.aligned;\n" ::: "memory"); + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); +} +template +__device__ static inline void mma_fence(D &dst) { + KITTENS_CHECK_WARPGROUP + #pragma unroll + for(int i = 0; i < D::height; i++) { + #pragma unroll + for(int j = 0; j < D::width; j++) { + #pragma unroll + for(int k = 0; k < dst.real.packed_per_tile; k++) { + if constexpr(std::is_same_v) { + asm volatile("" : "+f"(dst.real.tiles[i][j].data[k].x) :: "memory"); + asm volatile("" : "+f"(dst.real.tiles[i][j].data[k].y) :: "memory"); + asm volatile("" : "+f"(dst.imag.tiles[i][j].data[k].x) :: "memory"); + asm volatile("" : "+f"(dst.imag.tiles[i][j].data[k].y) :: "memory"); + } else if constexpr(std::is_same_v) { + asm volatile("" : "+r"(dst.real.tiles[i][j].data[k].x) :: "memory"); + asm volatile("" : "+r"(dst.real.tiles[i][j].data[k].y) :: "memory"); + asm volatile("" : "+r"(dst.imag.tiles[i][j].data[k].x) :: "memory"); + asm volatile("" : "+r"(dst.imag.tiles[i][j].data[k].y) :: "memory"); + } else { + asm volatile("" : "+r"(*(uint32_t*)&dst.real.tiles[i][j].data[k]) :: "memory"); + asm volatile("" : "+r"(*(uint32_t*)&dst.imag.tiles[i][j].data[k]) :: "memory"); + } + } + } + } + asm volatile ("wgmma.fence.sync.aligned;\n" ::: "memory"); + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); +} +template +__device__ static inline void mma_fence() { + KITTENS_CHECK_WARPGROUP + asm volatile ("wgmma.fence.sync.aligned;\n" ::: "memory"); + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); +} + + + + +template +__device__ static inline void mma_commit_group() { + KITTENS_CHECK_WARPGROUP + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); +} + + + + + + + + +template +__device__ static inline void mma_async_wait() { + KITTENS_CHECK_WARPGROUP + asm volatile ("wgmma.wait_group.sync.aligned %0;" : : "n"(N) : "memory"); +} + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +template +__device__ static inline void mma_AB(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M_DIV_4 = A::height; + static_assert(D::height == M_DIV_4); + constexpr int N = B::cols / kittens::TILE_COL_DIM; + static_assert(D::width == N); + constexpr int K = A::width; + static_assert(B::rows/kittens::TILE_ROW_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + static_assert(!std::is_same_v && !std::is_same_v, "Transpose is not supported for integer WGMMA; use mma_ABt / mm_ABt."); + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + using base = kittens::detail::wgmma::base*N, 0, 1>; + kittens::st_descriptor, 1> b_desc(b); + + if constexpr (fence) { mma_fence(d); } + + + #pragma unroll + for(int m = 0; m < M_DIV_4; m++) { + rt, TILE_COL_DIM*N, ducks::rt_layout::row> &d_ref = group<1>::subtile_inplace>(d, m); + base::rt_st( + d_ref, + a.tiles[m][0], + b_desc.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::rt_st( + d_ref, + a.tiles[m][k], + b_desc.chunk_descriptor(k), + 1 + ); + } + } + mma_commit_group(); +} +template +__device__ static inline void mm_AB(D &d, + const A &a, + const B &b) { + mma_AB(d, a, b); +} + +template +__device__ static inline void mma_AB(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M = A::rows / kittens::TILE_ROW_DIM; + static_assert(M == 4); + static_assert(D::height == 1); + constexpr int N = B::cols / kittens::TILE_COL_DIM; + constexpr int K = A::cols / kittens::TILE_COL_DIM; + static_assert(B::rows / kittens::TILE_ROW_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + static_assert(!std::is_same_v && !std::is_same_v, "Transpose is not supported for integer WGMMA; use mma_ABt / mm_ABt."); + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + using base = kittens::detail::wgmma::base*N, 0, 1>; + kittens::st_descriptor, 0> a_desc(a); + kittens::st_descriptor, 1> b_desc(b); + + if constexpr (fence) { mma_fence(d); } + + + base::st_st( + d, + a_desc.chunk_descriptor(0), + b_desc.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d, + a_desc.chunk_descriptor(k), + b_desc.chunk_descriptor(k), + 1 + ); + } + mma_commit_group(); +} +template +__device__ static inline void mm_AB(D &d, + const A &a, + const B &b) { + mma_AB(d, a, b); +} + + + + + + + + + + + + + + + + +template +__device__ static inline void mma_ABt(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M_DIV_4 = A::height; + static_assert(D::height == M_DIV_4); + constexpr int N = B::rows / kittens::TILE_ROW_DIM; + constexpr int K = A::width; + static_assert(B::cols/kittens::TILE_COL_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + using base = kittens::detail::wgmma::base*N, 0, 0>; + kittens::st_descriptor, 0> b_desc(b); + + if constexpr (fence) { mma_fence(d); } + + + #pragma unroll + for(int m = 0; m < M_DIV_4; m++) { + rt, TILE_COL_DIM*N, ducks::rt_layout::row> &d_ref = group<1>::subtile_inplace>(d, m); + base::rt_st( + d_ref, + a.tiles[m][0], + b_desc.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::rt_st( + d_ref, + a.tiles[m][k], + b_desc.chunk_descriptor(k), + 1 + ); + } + } + mma_commit_group(); +} +template +__device__ static inline void mm_ABt(D &d, + const A &a, + const B &b) { + mma_ABt(d, a, b); +} + + + + + + + + + + + + + + + + +template +__device__ static inline void mma_ABt(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M = A::rows / kittens::TILE_ROW_DIM; + static_assert(M == 4); + static_assert(D::height == 1); + constexpr int N = B::rows / kittens::TILE_ROW_DIM; + constexpr int K = A::cols / kittens::TILE_COL_DIM; + static_assert(B::cols/kittens::TILE_COL_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + using base = kittens::detail::wgmma::base*N, 0, 0>; + kittens::st_descriptor, 0> a_desc(a); + kittens::st_descriptor, 0> b_desc(b); + + if constexpr (fence) { mma_fence(d); } + + + base::st_st( + d, + a_desc.chunk_descriptor(0), + b_desc.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d, + a_desc.chunk_descriptor(k), + b_desc.chunk_descriptor(k), + 1 + ); + } + mma_commit_group(); +} +template +__device__ static inline void mm_ABt(D &d, + const A &a, + const B &b) { + mma_ABt(d, a, b); +} + + + + + + + + + + + + + + + + +template +__device__ static inline void mma_AtB(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M = A::cols / kittens::TILE_COL_DIM; + static_assert(M == 4); + static_assert(D::height == 1); + constexpr int N = B::cols / kittens::TILE_COL_DIM; + constexpr int K = A::rows / kittens::TILE_ROW_DIM; + static_assert(B::rows/kittens::TILE_ROW_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + static_assert(!std::is_same_v && !std::is_same_v, "Transpose is not supported for integer WGMMA; use mma_ABt / mm_ABt."); + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + using base = kittens::detail::wgmma::base*N, 1, 1>; + kittens::st_descriptor, 1> a_desc(a); + kittens::st_descriptor, 1> b_desc(b); + + if constexpr (fence) { mma_fence(d); } + + + base::st_st( + d, + a_desc.chunk_descriptor(0), + b_desc.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d, + a_desc.chunk_descriptor(k), + b_desc.chunk_descriptor(k), + 1 + ); + } + mma_commit_group(); +} +template +__device__ static inline void mm_AtB(D &d, + const A &a, + const B &b) { + mma_AtB(d, a, b); +} + + + + + + + + + + + + +template +__device__ static inline void mma_AtBt(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M = A::cols / kittens::TILE_COL_DIM; + static_assert(M == 4); + static_assert(D::height == 1); + constexpr int N = B::rows / kittens::TILE_ROW_DIM; + constexpr int K = A::rows / kittens::TILE_ROW_DIM; + static_assert(B::cols/kittens::TILE_COL_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + static_assert(!std::is_same_v && !std::is_same_v, "Transpose is not supported for integer WGMMA; use mma_ABt / mm_ABt."); + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + using base = kittens::detail::wgmma::base*N, 1, 0>; + kittens::st_descriptor, 1> a_desc(a); + kittens::st_descriptor, 0> b_desc(b); + + if constexpr (fence) { mma_fence(d); } + + + base::st_st( + d, + a_desc.chunk_descriptor(0), + b_desc.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d, + a_desc.chunk_descriptor(k), + b_desc.chunk_descriptor(k), + 1 + ); + } + mma_commit_group(); +} +template +__device__ static inline void mm_AtBt(D &d, + const A &a, + const B &b) { + mma_AtBt(d, a, b); +} + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +template +__device__ static inline void mma_AB(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M_DIV_4 = A::height; + static_assert(D::height == M_DIV_4); + constexpr int N = B::cols / kittens::TILE_COL_DIM; + constexpr int K = A::width; + static_assert(B::rows/kittens::TILE_ROW_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + static_assert(!std::is_same_v && !std::is_same_v, "Transpose is not supported for integer WGMMA; use mma_ABt / mm_ABt."); + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + using base = kittens::detail::wgmma::base*N, 0, 1>; + kittens::st_descriptor, 1> b_desc_real(b.real); + kittens::st_descriptor, 1> b_desc_imag(b.imag); + + if constexpr (fence) { mma_fence(d); } + + + #pragma unroll + for(int m = 0; m < M_DIV_4; m++) { + rt, TILE_COL_DIM*N, ducks::rt_layout::row> &d_ref = group<1>::subtile_inplace>(d.real, m); + base::rt_st( + d_ref, + a.real.tiles[m][0], + b_desc_real.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::rt_st( + d_ref, + a.real.tiles[m][k], + b_desc_real.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::rt_st<-1>( + d_ref, + a.imag.tiles[m][k], + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + } + #pragma unroll + for(int m = 0; m < M_DIV_4; m++) { + rt, TILE_COL_DIM*N, ducks::rt_layout::row> &d_ref = group<1>::subtile_inplace>(d.imag, m); + base::rt_st( + d_ref, + a.real.tiles[m][0], + b_desc_imag.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::rt_st( + d_ref, + a.real.tiles[m][k], + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::rt_st( + d_ref, + a.imag.tiles[m][k], + b_desc_real.chunk_descriptor(k), + 1 + ); + } + } + mma_commit_group(); +} +template +__device__ static inline void mm_AB(D &d, + const A &a, + const B &b) { + mma_AB(d, a, b); +} + +template +__device__ static inline void mma_AB(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M = A::rows / kittens::TILE_ROW_DIM; + static_assert(M == 4); + static_assert(D::height == 1); + constexpr int N = B::cols / kittens::TILE_COL_DIM; + constexpr int K = A::cols / kittens::TILE_COL_DIM; + static_assert(B::rows/kittens::TILE_ROW_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + static_assert(!std::is_same_v && !std::is_same_v, "Transpose is not supported for integer WGMMA; use mma_ABt / mm_ABt."); + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + using base = kittens::detail::wgmma::base*N, 0, 1>; + kittens::st_descriptor, 0> a_desc_real(a.real); + kittens::st_descriptor, 0> a_desc_imag(a.imag); + kittens::st_descriptor, 1> b_desc_real(b.real); + kittens::st_descriptor, 1> b_desc_imag(b.imag); + + if constexpr (fence) { mma_fence(d); } + + + base::st_st( + d.real, + a_desc_real.chunk_descriptor(0), + b_desc_real.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d.real, + a_desc_real.chunk_descriptor(k), + b_desc_real.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::st_st<-1>( + d.real, + a_desc_imag.chunk_descriptor(k), + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + base::st_st( + d.imag, + a_desc_real.chunk_descriptor(0), + b_desc_imag.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d.imag, + a_desc_real.chunk_descriptor(k), + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::st_st( + d.imag, + a_desc_imag.chunk_descriptor(k), + b_desc_real.chunk_descriptor(k), + 1 + ); + } + mma_commit_group(); +} +template +__device__ static inline void mm_AB(D &d, + const A &a, + const B &b) { + mma_AB(d, a, b); +} + + + + + + + + + + + + + + + + +template +__device__ static inline void mma_ABt(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M_DIV_4 = A::height; + static_assert(D::height == M_DIV_4); + constexpr int N = B::rows / kittens::TILE_ROW_DIM; + constexpr int K = A::width; + static_assert(B::cols/kittens::TILE_COL_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + using base = kittens::detail::wgmma::base*N, 0, 0>; + kittens::st_descriptor, 0> b_desc_real(b.real); + kittens::st_descriptor, 0> b_desc_imag(b.imag); + + if constexpr (fence) { mma_fence(d); } + + + #pragma unroll + for(int m = 0; m < M_DIV_4; m++) { + rt, TILE_ROW_DIM*N, ducks::rt_layout::row> &d_ref = group<1>::subtile_inplace>(d.real, m); + base::rt_st( + d_ref, + a.real.tiles[m][0], + b_desc_real.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::rt_st( + d_ref, + a.real.tiles[m][k], + b_desc_real.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::rt_st<-1>( + d_ref, + a.imag.tiles[m][k], + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + } + #pragma unroll + for(int m = 0; m < M_DIV_4; m++) { + rt, TILE_ROW_DIM*N, ducks::rt_layout::row> &d_ref = group<1>::subtile_inplace>(d.imag, m); + base::rt_st( + d_ref, + a.real.tiles[m][0], + b_desc_imag.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::rt_st( + d_ref, + a.real.tiles[m][k], + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::rt_st( + d_ref, + a.imag.tiles[m][k], + b_desc_real.chunk_descriptor(k), + 1 + ); + } + } + mma_commit_group(); +} +template +__device__ static inline void mm_ABt(D &d, + const A &a, + const B &b) { + mma_ABt(d, a, b); +} + + + + + + + + + + + + + + + + +template +__device__ static inline void mma_ABt(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M = A::rows / kittens::TILE_ROW_DIM; + static_assert(M == 4); + static_assert(D::height == 1); + constexpr int N = B::rows / kittens::TILE_ROW_DIM; + constexpr int K = A::cols / kittens::TILE_COL_DIM; + static_assert(B::cols/kittens::TILE_COL_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + using base = kittens::detail::wgmma::base*N, 0, 0>; + kittens::st_descriptor, 0> a_desc_real(a.real); + kittens::st_descriptor, 0> a_desc_imag(a.imag); + kittens::st_descriptor, 0> b_desc_real(b.real); + kittens::st_descriptor, 0> b_desc_imag(b.imag); + + if constexpr (fence) { mma_fence(d); } + + + base::st_st( + d.real, + a_desc_real.chunk_descriptor(0), + b_desc_real.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d.real, + a_desc_real.chunk_descriptor(k), + b_desc_real.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::st_st<-1>( + d.real, + a_desc_imag.chunk_descriptor(k), + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + base::st_st( + d.imag, + a_desc_real.chunk_descriptor(0), + b_desc_imag.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d.imag, + a_desc_real.chunk_descriptor(k), + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::st_st( + d.imag, + a_desc_imag.chunk_descriptor(k), + b_desc_real.chunk_descriptor(k), + 1 + ); + } + mma_commit_group(); +} +template +__device__ static inline void mm_ABt(D &d, + const A &a, + const B &b) { + mma_ABt(d, a, b); +} + + + + + + + + + + + + + + + + +template +__device__ static inline void mma_AtB(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M = A::cols / kittens::TILE_COL_DIM; + static_assert(M == 4); + static_assert(D::height == 1); + constexpr int N = B::cols / kittens::TILE_COL_DIM; + constexpr int K = A::rows / kittens::TILE_ROW_DIM; + static_assert(B::rows/kittens::TILE_ROW_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + static_assert(!std::is_same_v && !std::is_same_v, "Transpose is not supported for integer WGMMA; use mma_ABt / mm_ABt."); + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + using base = kittens::detail::wgmma::base*N, 1, 1>; + kittens::st_descriptor, 1> a_desc_real(a.real); + kittens::st_descriptor, 1> a_desc_imag(a.imag); + kittens::st_descriptor, 1> b_desc_real(b.real); + kittens::st_descriptor, 1> b_desc_imag(b.imag); + + if constexpr (fence) { mma_fence(d); } + + + base::st_st( + d.real, + a_desc_real.chunk_descriptor(0), + b_desc_real.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d.real, + a_desc_real.chunk_descriptor(k), + b_desc_real.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::st_st<-1>( + d.real, + a_desc_imag.chunk_descriptor(k), + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + base::st_st( + d.imag, + a_desc_real.chunk_descriptor(0), + b_desc_imag.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d.imag, + a_desc_real.chunk_descriptor(k), + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::st_st( + d.imag, + a_desc_imag.chunk_descriptor(k), + b_desc_real.chunk_descriptor(k), + 1 + ); + } + mma_commit_group(); +} +template +__device__ static inline void mm_AtB(D &d, + const A &a, + const B &b) { + mma_AtB(d, a, b); +} + + + + + + + + + + + + +template +__device__ static inline void mma_AtBt(D &d, + const A &a, + const B &b) { + + KITTENS_CHECK_WARPGROUP + constexpr int M = A::cols / kittens::TILE_COL_DIM; + static_assert(M == 4); + static_assert(D::height == 1); + constexpr int N = B::rows / kittens::TILE_ROW_DIM; + constexpr int K = A::rows / kittens::TILE_ROW_DIM; + static_assert(B::cols/kittens::TILE_COL_DIM == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; + using T_D = D::T; + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + static_assert(!std::is_same_v && !std::is_same_v, "Transpose is not supported for integer WGMMA; use mma_ABt / mm_ABt."); + static_assert(!std::is_same_v && !std::is_same_v, "Currently unsupported type"); + using base = kittens::detail::wgmma::base*N, 1, 0>; + kittens::st_descriptor, 1> a_desc_real(a.real); + kittens::st_descriptor, 1> a_desc_imag(a.imag); + kittens::st_descriptor, 0> b_desc_real(b.real); + kittens::st_descriptor, 0> b_desc_imag(b.imag); + + if constexpr (fence) { mma_fence(d); } + + + base::st_st( + d.real, + a_desc_real.chunk_descriptor(0), + b_desc_real.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d.real, + a_desc_real.chunk_descriptor(k), + b_desc_real.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::st_st<-1>( + d.real, + a_desc_imag.chunk_descriptor(k), + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + base::st_st( + d.imag, + a_desc_real.chunk_descriptor(0), + b_desc_imag.chunk_descriptor(0), + accumulate + ); + #pragma unroll + for(int k = 1; k < K; k++) { + base::st_st( + d.imag, + a_desc_real.chunk_descriptor(k), + b_desc_imag.chunk_descriptor(k), + 1 + ); + } + #pragma unroll + for(int k = 0; k < K; k++) { + base::st_st( + d.imag, + a_desc_imag.chunk_descriptor(k), + b_desc_real.chunk_descriptor(k), + 1 + ); + } + mma_commit_group(); +} +template +__device__ static inline void mm_AtBt(D &d, + const A &a, + const B &b) { + mma_AtBt(d, a, b); +} + + + +template +__device__ static inline void mma(D &d, + const A &a, + const B &b) { + if constexpr(trans_A == transpose::T) { + if constexpr(trans_B == transpose::T) { + mma_AtBt(d, a, b); + } else { + mma_AtB(d, a, b); + } + } else { + if constexpr(trans_B == transpose::T) { + mma_ABt(d, a, b); + } else { + mma_AB(d, a, b); + } + } +} +template +__device__ static inline void mm(D &d, + const A &a, + const B &b) { + if constexpr(trans_A == transpose::T) { + if constexpr(trans_B == transpose::T) { + mm_AtBt(d, a, b); + } else { + mm_AtB(d, a, b); + } + } else { + if constexpr(trans_B == transpose::T) { + mm_ABt(d, a, b); + } else { + mm_AB(d, a, b); + } + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/register.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/register.cuh new file mode 100644 index 0000000..404f8af --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/register.cuh @@ -0,0 +1,7 @@ + + + + + +#include "tile/tile.cuh" +#include "vec/vec.cuh" diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/complex_conversions.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/complex_conversions.cuh new file mode 100644 index 0000000..c2bd5f1 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/complex_conversions.cuh @@ -0,0 +1,98 @@ + + + + + + + + + + + + + + + + + + + + +template +__device__ static inline void swap_layout(crt::type> &dst, const crt &src) { + swap_layout(dst.real, src.real); + swap_layout(dst.real, src.real); +} + + + + + + + + + + +template +__device__ static inline crt::type>& swap_layout_inplace(crt &tile) { + tile.real = swap_layout_inplace(tile.real); + tile.imag = swap_layout_inplace(tile.imag); + return tile; +} + + + + + + + + + + + + + + + + +template +__device__ static inline void transpose_sep(crt &dst, const crt &src) { + transpose_sep(dst.real, src.real); + transpose_sep(dst.imag, src.imag); +} + + + + + + + + + + +template +__device__ static inline crt& transpose_inplace(crt &tile) { + tile.real = transpose_inplace(tile.real); + tile.imag = transpose_inplace(tile.imag); + + return tile; +} + + + + + + + + + + + + + + +template +__device__ static inline void copy(crt &dst, const crt &src) { + copy(dst.real, src.real); + copy(dst.imag, src.imag); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/complex_maps.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/complex_maps.cuh new file mode 100644 index 0000000..55c3f36 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/complex_maps.cuh @@ -0,0 +1,137 @@ + + + + + + + + + + + +template +__device__ static inline void zero(T &dst) { + zero(dst.real); + zero(dst.imag); +} + + + + + + + +template +__device__ static inline void exp(T &dst, const T &src) { + using dtype = T::dtype; + dtype tmp; + + dtype rdst; + dtype idst; + + + exp(rdst, src.real); + copy(idst, rdst); + + cos(tmp, src.imag); + mul(rdst, rdst, tmp); + sin(tmp, src.imag); + mul(idst, idst, tmp); + + copy(dst.real, rdst); + copy(dst.imag, idst); +} + + + + + + + + +template +__device__ static inline void add(T &dst, const T &lhs, const T &rhs) { + add(dst.real, lhs.real, rhs.real); + add(dst.imag, lhs.imag, rhs.imag); +} + + + + + + + + + +template +__device__ static inline void sub(T &dst, const T &lhs, const T &rhs) { + sub(dst.real, lhs.real, rhs.real); + sub(dst.imag, lhs.imag, rhs.imag); +} + + + + + + + + +template +__device__ static inline void mul(T &dst, const T &lhs, const T &rhs) { + using dtype = T::component; + dtype tmp; + + dtype rdst; + dtype idst; + + + + mul(rdst, lhs.real, rhs.real); + mul(tmp, lhs.imag, rhs.imag); + sub(rdst, rdst, tmp); + + + mul(idst, lhs.imag, rhs.real); + mul(tmp, lhs.real, rhs.imag); + add(idst, idst, tmp); + + copy(dst.real, rdst); + copy(dst.imag, idst); +} + + + + + + + + +template +__device__ static inline void div(T &dst, const T &lhs, const T &rhs) { + using dtype = T::dtype; + dtype tmp; + dtype denom; + + dtype rdst; + dtype idst; + + + mul(tmp, rhs.real, rhs.real); + mul(denom, rhs.imag, rhs.imag); + add(denom, tmp, denom); + + mul(rdst, lhs.real, rhs.real); + mul(tmp, lhs.imag, rhs.imag); + add(rdst, rdst, tmp); + + mul(dst.imag, lhs.imag, rhs.real); + mul(tmp, lhs.real, rhs.imag); + sub(idst, idst, tmp); + + div(rdst, rdst, denom); + div(idst, idst, denom); + copy(dst.real, rdst); + copy(dst.imag, idst); +} + + diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/conversions.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/conversions.cuh new file mode 100644 index 0000000..44d2cfc --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/conversions.cuh @@ -0,0 +1,917 @@ + + + + + + + + + + + + + + + + + +__device__ static inline void swap_layout_8(bf16_2 &dst, const bf16_2 &src) { + KITTENS_CHECK_WARP + asm volatile ( + "movmatrix.sync.aligned.m8n8.trans.b16 %0, %1;\n" + : "+r"(*(uint32_t*)(&dst)) + : "r"(*(uint32_t*)(&src)) + ); +} + + + + + + + + + + + +template +__device__ static inline void swap_layout(rt_base::type> &dst, const rt_base &src) { + swap_layout_8(dst.data[0], src.data[0]); + + + typename rt_base::T2 data1_cache = src.data[1]; + swap_layout_8(dst.data[1], src.data[2]); + swap_layout_8(dst.data[2], data1_cache); + swap_layout_8(dst.data[3], src.data[3]); +} + + + + + + + + + + + + + +template +__device__ static inline void swap_layout(rt::type> &dst, const rt &src) { + #pragma unroll + for(int i = 0; i < dst.height; i++) { + #pragma unroll + for(int j = 0; j < dst.width; j++) { + swap_layout(dst.tiles[i][j], src.tiles[i][j]); + } + } +} + + + + + + + + + + + + +template +__device__ static inline rt_base::type>& swap_layout_inplace(const rt_base &src) { + rt_base::type> &dst = *(rt_base::type>*)(&src); + swap_layout(dst, src); + return dst; +} + + + + + + + + + + + + + +template +__device__ static inline rt::type>& swap_layout_inplace(rt &tile) { + #pragma unroll + for(int i = 0; i < tile.height; i++) { + #pragma unroll + for(int j = 0; j < tile.width; j++) { + swap_layout_inplace(tile.tiles[i][j]); + } + } + return *(rt::type>*)(&tile); +} + + + + + + + + + + + +template +__device__ static inline void transpose(rt_base &dst, const rt_base &src) { + swap_layout_8(dst.data[0], src.data[0]); + + + typename rt_base::T2 data1_cache = src.data[1]; + swap_layout_8(dst.data[1], src.data[2]); + swap_layout_8(dst.data[2], data1_cache); + swap_layout_8(dst.data[3], src.data[3]); +} + + + + + + + + + + + + + +template +__device__ static inline void transpose_sep(RT &dst, const rt &src) { + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + transpose(dst.tiles[i][j], src.tiles[j][i]); + } + } +} + + + + + + + + + +template +__device__ static inline rt_base& transpose_inplace(rt_base &src) { + transpose(src, src); + return src; +} + + + + + + + + + + +template +__device__ static inline rt& transpose_inplace(rt &tile) { + static_assert(_cols == _rows, "in-place register tile transpose is only allowed for square tiles."); + #pragma unroll + for(int i = 0; i < tile.height; i++) { + #pragma unroll + for(int j = 0; j < i; j++) { + rt_base tmp; + copy(tmp, tile.tiles[i][j]); + transpose(tile.tiles[i][j], tile.tiles[j][i]); + transpose(tile.tiles[j][i], tmp); + } + transpose_inplace(tile.tiles[i][i]); + } + return tile; +} + + + + + + + + + + + + +template +__device__ static inline void copy(rt_base &dst, const rt_base &src) { + using T2 = typename base_types::packing::packed_type; + using U2 = typename base_types::packing::packed_type; + #pragma unroll + for(int k = 0; k < dst.packed_per_thread; k++) { + dst.data[k] = base_types::convertor::convert(src.data[k]); + } +} +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + + + + +template +__device__ static inline void copy(rt &dst, const rt &src) { + + if constexpr ( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) + ) { + + int laneid = threadIdx.x % 32; + + #pragma unroll + for(int i = 0; i < dst.height; i++) { + #pragma unroll + for(int j = 0; j < dst.width; j++) { + #pragma unroll + for(int k = 0; k < dst.tiles[0][0].packed_per_thread; k++) { + + + using src_t = std::conditional_t, float2, std::conditional_t, bf16_2, half2>>; + src_t val1, val2; + + + if (laneid % 2 == 0) { + + val1 = src.tiles[i][2*j + k/2].data[(k%2)+0]; + val2 = src.tiles[i][2*j + k/2].data[(k%2)+2]; + } else { + + val1 = src.tiles[i][2*j + k/2].data[(k%2)+2]; + val2 = src.tiles[i][2*j + k/2].data[(k%2)+0]; + } + + + int row_mask = 4 * ( laneid / 4 ); + int row_offset = row_mask + ( (laneid-row_mask) / 2 ) + ( laneid % 2 ); + int src_offset = (laneid % 2 == 0 ) ? row_offset + 0 : ( row_offset + 1 ); + src_t val01 = packed_shfl_sync(MASK_ALL, val1, src_offset); + + int src_offset2 = (laneid % 4 < 2 ) ? src_offset + 1 : (src_offset - 1); + src_t val23 = packed_shfl_sync(MASK_ALL, val2, src_offset2); + + + float4 f4; + using fp8_4_t = std::conditional_t, fp8e4m3_4, fp8e5m2_4>; + fp8_4_t f4_fp8; + if ( laneid % 4 < 2 ) { + f4.x = val01.x; + f4.y = val01.y; + f4.z = val23.x; + f4.w = val23.y; + f4_fp8 = base_types::convertor::convert(f4); + dst.tiles[i][j].data[k] = f4_fp8; + } else { + f4.x = val23.x; + f4.y = val23.y; + f4.z = val01.x; + f4.w = val01.y; + f4_fp8 = base_types::convertor::convert(f4); + dst.tiles[i][j].data[k] = f4_fp8; + } + } + } + } + } + else if constexpr ( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) + ) { + + int laneid = threadIdx.x % 32; + + #pragma unroll + for(int i = 0; i < src.height; i++) { + #pragma unroll + for(int j = 0; j < src.width; j++) { + #pragma unroll + for(int k = 0; k < src.tiles[0][0].packed_per_thread; k++) { + int dst_j = 2*j + k/2; + + + using fp8_4_t = std::conditional_t, fp8e4m3_4, fp8e5m2_4>; + fp8_4_t val = src.tiles[i][j].data[k]; + float4 f4 = base_types::convertor::convert(val); + float2 f2_0, f2_1; + if ( laneid % 4 < 2 ) { + f2_0 = make_float2(f4.x, f4.y); + f2_1 = make_float2(f4.z, f4.w); + } + else { + f2_0 = make_float2(f4.z, f4.w); + f2_1 = make_float2(f4.x, f4.y); + } + + int row_offset = 4 * (laneid/4) + (laneid%2) * 2 + (laneid%4) / 2; + float2 f2_0_shfl = packed_shfl_sync(MASK_ALL, f2_0, row_offset); + float2 f2_1_shfl = packed_shfl_sync(MASK_ALL, f2_1, row_offset^2); + + + using dst_t = std::conditional_t, float2, std::conditional_t, bf16_2, half2>>; + if constexpr (!(std::is_same_v)) { + dst_t f2_0_shfl_t = base_types::convertor::convert(f2_0_shfl); + dst_t f2_1_shfl_t = base_types::convertor::convert(f2_1_shfl); + if (laneid % 2 == 0) { + dst.tiles[i][dst_j].data[(k%2)+0] = f2_0_shfl_t; + dst.tiles[i][dst_j].data[(k%2)+2] = f2_1_shfl_t; + } else { + dst.tiles[i][dst_j].data[(k%2)+0] = f2_1_shfl_t; + dst.tiles[i][dst_j].data[(k%2)+2] = f2_0_shfl_t; + } + } else { + if (laneid % 2 == 0) { + dst.tiles[i][dst_j].data[(k%2)+0] = f2_0_shfl; + dst.tiles[i][dst_j].data[(k%2)+2] = f2_1_shfl; + } else { + dst.tiles[i][dst_j].data[(k%2)+0] = f2_1_shfl; + dst.tiles[i][dst_j].data[(k%2)+2] = f2_0_shfl; + } + } + } + } + } + } + + else { + #pragma unroll + for(int i = 0; i < dst.height; i++) { + #pragma unroll + for(int j = 0; j < dst.width; j++) { + copy(dst.tiles[i][j], src.tiles[i][j]); + } + } + } +} +#else + + + + + + + + + + + +template +__device__ static inline void copy(rt &dst, const rt &src) { + #pragma unroll + for(int i = 0; i < dst.height; i++) { + #pragma unroll + for(int j = 0; j < dst.width; j++) { + copy(dst.tiles[i][j], src.tiles[i][j]); + } + } +} +#endif + + + + + + + + + + + + + + + +template +__device__ static inline void make_causal(RT &dst, const RT &src, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + const typename RT::dtype packed_val = base_types::packing::pack(val); + #if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert(!std::is_same_v && !std::is_same_v, "Unsupported type for make_causal"); + #endif + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + if(j < i) { + #pragma unroll + for(int k = 0; k < RT::packed_per_tile; k++) { + dst.tiles[i][j].data[k] = src.tiles[i][j].data[k]; + } + } + else if(j > i) { + #pragma unroll + for(int k = 0; k < RT::packed_per_tile; k++) { + dst.tiles[i][j].data[k] = packed_val; + } + } + else { + constexpr uint32_t MASK_X = 0xFF773311, MASK_Y = 0xF7733110; + dst.tiles[i][j].data[1] = src.tiles[i][j].data[1]; + dst.tiles[i][j].data[2] = packed_val; + if((MASK_X >> laneid()) & 1) { + dst.tiles[i][j].data[0].x = src.tiles[i][j].data[0].x; + dst.tiles[i][j].data[3].x = src.tiles[i][j].data[3].x; + } + else { + dst.tiles[i][j].data[0].x = val; + dst.tiles[i][j].data[3].x = val; + } + if((MASK_Y >> laneid()) & 1) { + dst.tiles[i][j].data[0].y = src.tiles[i][j].data[0].y; + dst.tiles[i][j].data[3].y = src.tiles[i][j].data[3].y; + } + else { + dst.tiles[i][j].data[0].y = val; + dst.tiles[i][j].data[3].y = val; + } + } + __syncwarp(); + } + } +} + + + + + + + + + + + + + +template +__device__ static inline void make_causal_t(RT &dst, const RT &src, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + const typename RT::dtype packed_val = base_types::packing::pack(val); + #if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert(!std::is_same_v && !std::is_same_v, "Unsupported type for make_causal"); + #endif + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + if(j > i) { + #pragma unroll + for(int k = 0; k < RT::packed_per_tile; k++) { + dst.tiles[i][j].data[k] = src.tiles[i][j].data[k]; + } + } + else if(j < i) { + #pragma unroll + for(int k = 0; k < RT::packed_per_tile; k++) { + dst.tiles[i][j].data[k] = packed_val; + } + } + else { + constexpr uint32_t MASK_X = 0x88CCEEF; + constexpr uint32_t MASK_Y = 0x88CCEEFF; + + dst.tiles[i][j].data[1] = packed_val; + dst.tiles[i][j].data[2] = src.tiles[i][j].data[2]; + + + if((MASK_X >> laneid()) & 1) { + dst.tiles[i][j].data[0].x = src.tiles[i][j].data[0].x; + dst.tiles[i][j].data[3].x = src.tiles[i][j].data[3].x; + } + + else { + dst.tiles[i][j].data[0].x = val; + dst.tiles[i][j].data[3].x = val; + } + + + if((MASK_Y >> laneid()) & 1) { + dst.tiles[i][j].data[0].y = src.tiles[i][j].data[0].y; + dst.tiles[i][j].data[3].y = src.tiles[i][j].data[3].y; + } + + else { + dst.tiles[i][j].data[0].y = val; + dst.tiles[i][j].data[3].y = val; + } + + } + __syncwarp(); + } + } +} + + + + + + + + + + + + +template +__device__ static inline void tril(RT &dst, const RT &src, const int diagonal, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + const typename RT::dtype packed_val = base_types::packing::pack(val); + + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int global_row_idx = (i * RT::tile_size_row) + ((k % 2) * 8) + (laneid() / 4); + const int global_col_idx_x = (j * RT::tile_size_col) + ((k / 2) * 8) + ((laneid() % 4) * 2); + const int global_col_idx_y = (j * RT::tile_size_col) + ((k / 2) * 8) + ((laneid() % 4) * 2) + 1; + + if (global_col_idx_x <= global_row_idx + diagonal) { dst.tiles[i][j].data[k].x = src.tiles[i][j].data[k].x; } + else { dst.tiles[i][j].data[k].x = val; } + + if (global_col_idx_y <= global_row_idx + diagonal) { dst.tiles[i][j].data[k].y = src.tiles[i][j].data[k].y; } + else { dst.tiles[i][j].data[k].y = val; } + } + } + __syncwarp(); + } +} +template +__device__ static inline void tril(RT &dst, const RT &src, const int diagonal, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int global_row_idx_x = (i * RT::tile_size_row) + ((k / 2) * 8) + ((laneid() % 4) * 2); + const int global_row_idx_y = (i * RT::tile_size_row) + ((k / 2) * 8) + ((laneid() % 4) * 2) + 1; + const int global_col_idx = (j * RT::tile_size_col) + ((k % 2) * 8) + (laneid() / 4); + + if (global_col_idx <= global_row_idx_x + diagonal) { dst.tiles[i][j].data[k].x = src.tiles[i][j].data[k].x; } + else { dst.tiles[i][j].data[k].x = val; } + if (global_col_idx <= global_row_idx_y + diagonal) { dst.tiles[i][j].data[k].y = src.tiles[i][j].data[k].y; } + else { dst.tiles[i][j].data[k].y = val; } + } + } + __syncwarp(); + } +} + + + + + + + + + + +template +__device__ static inline void triu(RT &dst, const RT &src, const int diagonal, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + const typename RT::dtype packed_val = base_types::packing::pack(val); + + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int global_row_idx = (i * RT::tile_size_row) + ((k % 2) * 8) + (laneid() / 4); + const int global_col_idx_x = (j * RT::tile_size_col) + ((k / 2) * 8) + ((laneid() % 4) * 2); + const int global_col_idx_y = (j * RT::tile_size_col) + ((k / 2) * 8) + ((laneid() % 4) * 2) + 1; + + if (global_col_idx_x >= global_row_idx + diagonal) { dst.tiles[i][j].data[k].x = src.tiles[i][j].data[k].x; } + else { dst.tiles[i][j].data[k].x = val; } + + if (global_col_idx_y >= global_row_idx + diagonal) { dst.tiles[i][j].data[k].y = src.tiles[i][j].data[k].y; } + else { dst.tiles[i][j].data[k].y = val; } + } + } + __syncwarp(); + } +} +template +__device__ static inline void triu(RT &dst, const RT &src, const int diagonal, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int global_row_idx_x = (i * RT::tile_size_row) + ((k / 2) * 8) + ((laneid() % 4) * 2); + const int global_row_idx_y = (i * RT::tile_size_row) + ((k / 2) * 8) + ((laneid() % 4) * 2) + 1; + const int global_col_idx = (j * RT::tile_size_col) + ((k % 2) * 8) + (laneid() / 4); + + if (global_col_idx >= global_row_idx_x + diagonal) { dst.tiles[i][j].data[k].x = src.tiles[i][j].data[k].x; } + else { dst.tiles[i][j].data[k].x = val; } + + if (global_col_idx >= global_row_idx_y + diagonal) { dst.tiles[i][j].data[k].y = src.tiles[i][j].data[k].y; } + else { dst.tiles[i][j].data[k].y = val; } + } + } + __syncwarp(); + } +} + + + + + + + + + + + + +template +__device__ static inline void right_fill(RT &dst, const RT &src, const int col_idx, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + if(col_idx >= RT::cols) return; + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int col_idx_x = (j * RT::tile_size_col) + ((k / 2) * 8) + ((laneid() % 4) * 2); + const int col_idx_y = (j * RT::tile_size_col) + ((k / 2) * 8) + ((laneid() % 4) * 2) + 1; + if (col_idx_x >= col_idx) { dst.tiles[i][j].data[k].x = val; } + else { dst.tiles[i][j].data[k].x = src.tiles[i][j].data[k].x; } + if (col_idx_y >= col_idx) { dst.tiles[i][j].data[k].y = val; } + else { dst.tiles[i][j].data[k].y = src.tiles[i][j].data[k].y; } + } + } + } +} +template +__device__ static inline void right_fill(RT &dst, const RT &src, const int col_idx, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + const typename RT::dtype packed_val = base_types::packing::pack(val); + + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int t_col_idx = (j * RT::tile_size_col) + ((k % 2) * 8) + (laneid() / 4); + if (t_col_idx >= col_idx) { dst.tiles[i][j].data[k] = packed_val; } + else { dst.tiles[i][j].data[k] = src.tiles[i][j].data[k]; } + } + } + __syncwarp(); + } +} + + + + + + + + + + +template +__device__ static inline void left_fill(RT &dst, const RT &src, const int col_idx, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + if(col_idx <= 0) return; + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int col_idx_x = (j * RT::tile_size_col) + ((k / 2) * 8) + ((laneid() % 4) * 2); + const int col_idx_y = (j * RT::tile_size_col) + ((k / 2) * 8) + ((laneid() % 4) * 2) + 1; + if (col_idx_x < col_idx) { dst.tiles[i][j].data[k].x = val; } + else { dst.tiles[i][j].data[k].x = src.tiles[i][j].data[k].x; } + if (col_idx_y < col_idx) { dst.tiles[i][j].data[k].y = val; } + else { dst.tiles[i][j].data[k].y = src.tiles[i][j].data[k].y; } + } + } + } +} +template +__device__ static inline void left_fill(RT &dst, const RT &src, const int col_idx, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + const typename RT::dtype packed_val = base_types::packing::pack(val); + + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int thread_col = (j * RT::tile_size_col) + ((k % 2) * 8) + ((laneid() / 4)); + if (thread_col < col_idx) { dst.tiles[i][j].data[k] = packed_val; } + else { dst.tiles[i][j].data[k] = src.tiles[i][j].data[k]; } + } + } + __syncwarp(); + } +} + + + + + + + + + + +template +__device__ static inline void upper_fill(RT &dst, const RT &src, const int row_idx, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + if(row_idx <= 0) return; + const typename RT::dtype packed_val = base_types::packing::pack(val); + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int thread_row = (i * RT::tile_size_row) + ((k % 2) * 8) + ((laneid() / 4)); + if (thread_row < row_idx) { dst.tiles[i][j].data[k] = packed_val; } + else { dst.tiles[i][j].data[k] = src.tiles[i][j].data[k]; } + } + } + } +} +template +__device__ static inline void upper_fill(RT &dst, const RT &src, const int row_idx, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int row_idx_x = (i * RT::tile_size_row) + ((k / 2) * 8) + ((laneid() % 4) * 2); + const int row_idx_y = (i * RT::tile_size_row) + ((k / 2) * 8) + ((laneid() % 4) * 2) + 1; + if (row_idx_x < row_idx) { dst.tiles[i][j].data[k].x = val; } + else { dst.tiles[i][j].data[k].x = src.tiles[i][j].data[k].x; } + if (row_idx_y < row_idx) { dst.tiles[i][j].data[k].y = val; } + else { dst.tiles[i][j].data[k].y = src.tiles[i][j].data[k].y; } + } + } + } +} + + + + + + + + + + +template +__device__ static inline void lower_fill(RT &dst, const RT &src, const int row_idx, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + if(row_idx >= RT::rows) return; + const typename RT::dtype packed_val = base_types::packing::pack(val); + + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int thread_row = (i * RT::tile_size_row) + ((k % 2) * 8) + ((laneid() / 4)); + if (thread_row >= row_idx) { dst.tiles[i][j].data[k] = packed_val; } + else { dst.tiles[i][j].data[k] = src.tiles[i][j].data[k]; } + } + } + } +} +template +__device__ static inline void lower_fill(RT &dst, const RT &src, const int row_idx, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int row_idx_x = (i * RT::tile_size_row) + ((k / 2) * 8) + ((laneid() % 4) * 2); + const int row_idx_y = (i * RT::tile_size_row) + ((k / 2) * 8) + ((laneid() % 4) * 2) + 1; + if (row_idx_x >= row_idx) { dst.tiles[i][j].data[k].x = val; } + else { dst.tiles[i][j].data[k].x = src.tiles[i][j].data[k].x; } + if (row_idx_y >= row_idx) { dst.tiles[i][j].data[k].y = val; } + else { dst.tiles[i][j].data[k].y = src.tiles[i][j].data[k].y; } + } + } + } +} + + + + + + + + + + + + + +template +__device__ static inline void upper_right_fill(RT &dst, const RT &src, const int row_idx, const int col_idx, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + if(col_idx >= RT::cols || row_idx < 0) return; + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int row_idx_xy = (i * RT::tile_size_row) + ((k % 2) * 8) + ((laneid() / 4)); + const int col_idx_x = (j * RT::tile_size_col) + ((k / 2) * 8) + ((laneid() % 4) * 2); + const int col_idx_y = (j * RT::tile_size_col) + ((k / 2) * 8) + ((laneid() % 4) * 2) + 1; + if (col_idx_x >= col_idx && row_idx_xy < row_idx) { dst.tiles[i][j].data[k].x = val; } + else { dst.tiles[i][j].data[k].x = src.tiles[i][j].data[k].x; } + if (col_idx_y >= col_idx && row_idx_xy < row_idx) { dst.tiles[i][j].data[k].y = val; } + else { dst.tiles[i][j].data[k].y = src.tiles[i][j].data[k].y; } + } + } + __syncwarp(); + } +} +template +__device__ static inline void upper_right_fill(RT &dst, const RT &src, const int row_idx, const int col_idx, const typename base_types::packing::unpacked_type &val=0) { + KITTENS_CHECK_WARP + const typename RT::dtype packed_val = base_types::packing::pack(val); + if(col_idx >= RT::cols || row_idx < 0) return; + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for (int k = 0; k < RT::packed_per_tile; k++) { + const int row_idx_x = (i * RT::tile_size_row) + ((k / 2) * 8) + ((laneid() % 4) * 2); + const int row_idx_y = (i * RT::tile_size_row) + ((k / 2) * 8) + ((laneid() % 4) * 2) + 1; + const int col_idx_xy = (j * RT::tile_size_col) + ((k % 2) * 8) + (laneid() / 4); + if (row_idx_x < row_idx && col_idx_xy < col_idx) { dst.tiles[i][j].data[k].x = val; } + else { dst.tiles[i][j].data[k].x = src.tiles[i][j].data[k].x; } + if (row_idx_y < row_idx && col_idx_xy < col_idx) { dst.tiles[i][j].data[k].y = val; } + else { dst.tiles[i][j].data[k].y = src.tiles[i][j].data[k].y; } + } + } + __syncwarp(); + } +} + + + + + + + + + + + + + +template +__device__ static inline rt &subtile_inplace(RT & src, int idx) { + KITTENS_CHECK_WARP + using T = typename RT::T; + static_assert(RT::height % (subtile_rows / TILE_ROW_DIM) == 0, "subtile height should evenly divide tile height."); + return reinterpret_cast&>( + src.tiles[idx*(subtile_rows / TILE_ROW_DIM)] + ); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/maps.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/maps.cuh new file mode 100644 index 0000000..0142037 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/maps.cuh @@ -0,0 +1,890 @@ + + + + + + + + + + + + + + + +template +__device__ static inline void unary_map(T &dst, const T &src) { + #pragma unroll + for(int i = 0; i < T::height; i++) { + #pragma unroll + for(int j = 0; j < T::width; j++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile; k++) { + dst.tiles[i][j].data[k] = op::template op(src.tiles[i][j].data[k]); + } + } + } +} + + + + + + + + + + +template +__device__ static inline void bin_map(T &dst, const T &src, const typename T::dtype ¶m) { + #pragma unroll + for(int i = 0; i < T::height; i++) { + #pragma unroll + for(int j = 0; j < T::width; j++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile; k++) { + dst.tiles[i][j].data[k] = op::template op(src.tiles[i][j].data[k], param); + } + } + } +} + + + + + + + + + +template +__device__ static inline void bin_map(T &dst, const T &src, const typename base_types::packing::unpacked_type ¶m) { + + bin_map(dst, src, base_types::packing::pack(param)); +} + + + + + + + + + +template +__device__ static inline void bin_map(T &dst, const T &lhs, const T &rhs) { + #pragma unroll + for(int i = 0; i < T::height; i++) { + #pragma unroll + for(int j = 0; j < T::width; j++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile; k++) { + dst.tiles[i][j].data[k] = op::template op(lhs.tiles[i][j].data[k], rhs.tiles[i][j].data[k]); + } + } + } +} + +template +__device__ static inline void apply(RT &dst, const RT &src, Lambda &&lambda) { + int row_offset = 0; + if constexpr(GROUP_WARPS > 1) { + row_offset = warpid()*RT::height; + } + static_assert(sizeof(RT::T) != 1, "Cannot apply lambda to 8-bit types"); + if constexpr (ducks::rt::row_layout) { + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for(int k = 0; k < RT::packed_per_tile; k++) { + int row = row_offset + i*TILE_ROW_DIM + (k%2) * (TILE_ROW_DIM/2) + ::kittens::laneid()/4; + int col = j*TILE_COL_DIM + (k/2) * (TILE_COL_DIM/2) + (::kittens::laneid()%4)*2; + dst.tiles[i][j].data[k].x = lambda(row, col+0, src.tiles[i][j].data[k].x); + dst.tiles[i][j].data[k].y = lambda(row, col+1, src.tiles[i][j].data[k].y); + } + } + } + } + else { + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for(int k = 0; k < RT::packed_per_tile; k++) { + int row = row_offset + i*TILE_ROW_DIM + (k/2) * (TILE_ROW_DIM/2) + (::kittens::laneid()%4)*2; + int col = j*TILE_COL_DIM + (k%2) * (TILE_COL_DIM/2) + ::kittens::laneid()/4; + dst.tiles[i][j].data[k].x = lambda(row+0, col, src.tiles[i][j].data[k].x); + dst.tiles[i][j].data[k].y = lambda(row+1, col, src.tiles[i][j].data[k].y); + } + } + } + } +} +template +__device__ static inline RT apply(const RT &src, Lambda &&lambda) { + RT dst; + apply(dst, src, std::forward(lambda)); + return dst; +} + + + + + + + + + + + + + +template +__device__ static inline void row_map(T &dst, const T &src, const V &row_values) { + + static_assert(std::is_same_v::col_vec_layout>); + static_assert(std::is_same_v); + static_assert(V::outer_dim == T::height); + + using dtype = T::dtype; + + #pragma unroll + for(int i = 0; i < T::height; i++) { + dtype packed_top_row = base_types::packing::pack(row_values[i][0].x); + dtype packed_bottom_row = base_types::packing::pack(row_values[i][0].y); + #pragma unroll + for(int j = 0; j < T::width; j++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile; k+=2) { + dst.tiles[i][j].data[k+0] = op::template op(src.tiles[i][j].data[k+0], packed_top_row); + dst.tiles[i][j].data[k+1] = op::template op(src.tiles[i][j].data[k+1], packed_bottom_row); + } + } + } +} + + + + + + + + + + +template +__device__ static inline void row_map(T &dst, const T &src, const V &row_values) { + + static_assert(std::is_same_v); + static_assert(std::is_same_v::col_vec_layout>); + static_assert(V::outer_dim == T::height); + + using dtype = T::dtype; + + #pragma unroll + for(int i = 0; i < T::height; i++) { + #pragma unroll + for(int j = 0; j < T::width; j++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile/2; k++) { + dst.tiles[i][j].data[k+0] = op::template op(src.tiles[i][j].data[k+0], row_values[i][0]); + dst.tiles[i][j].data[k+2] = op::template op(src.tiles[i][j].data[k+2], row_values[i][1]); + } + } + } +} + + + + + + + + + + + + + + + +template +__device__ static inline void row_map(T &dst, const T &a, const T &b, const V &row_values) { + + static_assert(std::is_same_v::col_vec_layout>); + static_assert(std::is_same_v); + static_assert(V::outer_dim == T::height); + + using dtype = T::dtype; + + #pragma unroll + for(int i = 0; i < T::height; i++) { + dtype packed_top_row = base_types::packing::pack(row_values[i][0].x); + dtype packed_bottom_row = base_types::packing::pack(row_values[i][0].y); + #pragma unroll + for(int j = 0; j < T::width; j++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile; k+=2) { + dst.tiles[i][j].data[k+0] = op::template op(a.tiles[i][j].data[k+0], b.tiles[i][j].data[k+0], packed_top_row); + dst.tiles[i][j].data[k+1] = op::template op(a.tiles[i][j].data[k+1], b.tiles[i][j].data[k+1], packed_bottom_row); + } + } + } +} + + + + + + + + + + + +template +__device__ static inline void row_map(T &dst, const T &a, const T &b, const V &row_values) { + + static_assert(std::is_same_v::col_vec_layout>); + static_assert(std::is_same_v); + static_assert(V::outer_dim == T::height); + + using dtype = T::dtype; + + #pragma unroll + for(int i = 0; i < T::height; i++) { + #pragma unroll + for(int j = 0; j < T::width; j++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile/2; k++) { + dst.tiles[i][j].data[k+0] = op::template op(a.tiles[i][j].data[k+0], b.tiles[i][j].data[k+0], row_values[i][0]); + dst.tiles[i][j].data[k+2] = op::template op(a.tiles[i][j].data[k+2], b.tiles[i][j].data[k+2], row_values[i][1]); + } + } + } +} + + + + + + + + + + + + + +template +__device__ static inline void col_map(T &dst, const T &src, const V &col_values) { + KITTENS_CHECK_WARP + + static_assert(std::is_same_v::row_vec_layout>); + static_assert(std::is_same_v); + static_assert(V::outer_dim == T::width); + + using dtype = T::dtype; + + #pragma unroll + for(int j = 0; j < T::width; j++) { + #pragma unroll + for(int i = 0; i < T::height; i++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile/2; k++) { + dst.tiles[i][j].data[k+0] = op::template op(src.tiles[i][j].data[k+0], col_values[j][0]); + dst.tiles[i][j].data[k+2] = op::template op(src.tiles[i][j].data[k+2], col_values[j][1]); + } + } + } +} + + + + + + + + + + +template +__device__ static inline void col_map(T &dst, const T &src, const V &col_values) { + KITTENS_CHECK_WARP + + static_assert(std::is_same_v::row_vec_layout>); + static_assert(std::is_same_v); + static_assert(V::outer_dim == T::width); + + using dtype = T::dtype; + + #pragma unroll + for(int j = 0; j < T::width; j++) { + dtype packed_left_col = base_types::packing::pack(col_values[j][0].x); + dtype packed_right_col = base_types::packing::pack(col_values[j][0].y); + #pragma unroll + for(int i = 0; i < T::height; i++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile; k+=2) { + dst.tiles[i][j].data[k+0] = op::template op(src.tiles[i][j].data[k+0], packed_left_col); + dst.tiles[i][j].data[k+1] = op::template op(src.tiles[i][j].data[k+1], packed_right_col); + } + } + } +} + + + + + + + + + + + + + +template +__device__ static inline void col_map(T &dst, const T &a, const T &b, const V &col_values) { + KITTENS_CHECK_WARP + + static_assert(std::is_same_v::row_vec_layout>); + static_assert(std::is_same_v); + static_assert(V::outer_dim == T::width); + + using dtype = T::dtype; + + #pragma unroll + for(int j = 0; j < T::width; j++) { + #pragma unroll + for(int i = 0; i < T::height; i++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile/2; k++) { + dst.tiles[i][j].data[k+0] = op::template op(a.tiles[i][j].data[k+0], b.tiles[i][j].data[k+0], col_values[j][0]); + dst.tiles[i][j].data[k+2] = op::template op(a.tiles[i][j].data[k+2], b.tiles[i][j].data[k+2], col_values[j][1]); + } + } + } +} + + + + + + + + + + + +template +__device__ static inline void col_map(T &dst, const T &a, const T &b, const V &col_values) { + KITTENS_CHECK_WARP + + static_assert(std::is_same_v); + static_assert(std::is_same_v::row_vec_layout>); + static_assert(V::outer_dim == T::width); + + using dtype = T::dtype; + #pragma unroll + for(int j = 0; j < T::width; j++) { + dtype packed_left_col = base_types::packing::pack(col_values[j][0].x); + dtype packed_right_col = base_types::packing::pack(col_values[j][0].y); + #pragma unroll + for(int i = 0; i < T::height; i++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile; k+=2) { + dst.tiles[i][j].data[k+0] = op::template op(a.tiles[i][j].data[k+0], b.tiles[i][j].data[k+0], packed_left_col); + dst.tiles[i][j].data[k+1] = op::template op(a.tiles[i][j].data[k+1], b.tiles[i][j].data[k+1], packed_right_col); + } + } + } +} + + + + + + + + + + + + + +template +__device__ static inline void zero(T &dst) { + unary_map(dst, dst); +} + + + + + + +template +__device__ static inline void one(T &dst) { + unary_map(dst, dst); +} + + + + + + +template +__device__ static inline void pos_infty(T &dst) { + unary_map(dst, dst); +} + + + + + + +template +__device__ static inline void neg_infty(T &dst) { + unary_map(dst, dst); +} + + + + + + + + +template +__device__ static inline void exp(T &dst, const T &src) { + unary_map(dst, src); +} +template +__device__ static inline T exp(const T &src) { + T dst; + exp(dst, src); + return dst; +} + + + + + + + + +template +__device__ static inline void exp2(T &dst, const T &src) { + unary_map(dst, src); +} +template +__device__ static inline T exp2(const T &src) { + T dst; + exp2(dst, src); + return dst; +} + + + + + + + + +template +__device__ static inline void log(T &dst, const T &src) { + unary_map(dst, src); +} +template +__device__ static inline T log(const T &src) { + T dst; + log(dst, src); + return dst; +} + + + + + + + + +template +__device__ static inline void log2(T &dst, const T &src) { + unary_map(dst, src); +} +template +__device__ static inline T log2(const T &src) { + T dst; + log2(dst, src); + return dst; +} + + + + + + + + +template +__device__ static inline void abs(T &dst, const T &src) { + unary_map(dst, src); +} +template +__device__ static inline T abs(const T &src) { + T dst; + abs(dst, src); + return dst; +} + + + + + + + + +template +__device__ static inline void relu(T &dst, const T &src) { + unary_map(dst, src); +} +template +__device__ static inline T relu(const T &src) { + T dst; + relu(dst, src); + return dst; +} + + + + + + + + +template +__device__ static inline void neg(T &dst, const T &src) { + unary_map(dst, src); +} +template +__device__ static inline T neg(const T &src) { + T dst; + neg(dst, src); + return dst; +} + + + + + + + + +template +__device__ static inline void sqrt(T &dst, const T &src) { + unary_map(dst, src); +} +template +__device__ static inline T sqrt(const T &src) { + T dst; + sqrt(dst, src); + return dst; +} + + + + + + + + +template +__device__ static inline void rsqrt(T &dst, const T &src) { + unary_map(dst, src); +} +template +__device__ static inline T rsqrt(const T &src) { + T dst; + rsqrt(dst, src); + return dst; +} + + + + + + + + + +template +__device__ static inline void copy(T &dst, const U &src) { + bin_map(dst, src); +} + + + + + + + + + + +template +__device__ static inline void max(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} +template +__device__ static inline T max(const T &lhs, const U &rhs) { + T dst; + max(dst, lhs, rhs); + return dst; +} + + + + + + + + + + +template +__device__ static inline void min(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} +template +__device__ static inline T min(const T &lhs, const U &rhs) { + T dst; + min(dst, lhs, rhs); + return dst; +} + + + + + + + + + + +template +__device__ static inline void add(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + + + + + + + + + + +template +__device__ static inline void sub(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void mul(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + + + + + + + + + + +template +__device__ static inline void div(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + + + + + + + + + + +template +__device__ static inline void add_row(T &dst, const T &src, const V &row_values) { + row_map(dst, src, row_values); +} + + + + + + + + + + +template +__device__ static inline void sub_row(T &dst, const T &src, const V &row_values) { + row_map(dst, src, row_values); +} + + + + + + + + + + +template +__device__ static inline void mul_row(T &dst, const T &src, const V &row_values) { + row_map(dst, src, row_values); +} + + + + + + + + + + +template +__device__ static inline void div_row(T &dst, const T &src, const V &row_values) { + row_map(dst, src, row_values); +} + + + + + + + + + +template +__device__ static inline void broadcast_row(T &dst, const V &row_values) { + row_map(dst, dst, row_values); +} +template +__device__ static inline T broadcast_row(const V &row_values) { + T dst; + broadcast_row(dst, row_values); + return dst; +} + + + + + + + + + + + + +template +__device__ static inline void add_col(T &dst, const T &src, const V &col_values) { + col_map(dst, src, col_values); +} + + + + + + + + + + +template +__device__ static inline void sub_col(T &dst, const T &src, const V &col_values) { + col_map(dst, src, col_values); +} + + + + + + + + + + +template +__device__ static inline void mul_col(T &dst, const T &src, const V &col_values) { + col_map(dst, src, col_values); +} + + + + + + + + + + +template +__device__ static inline void div_col(T &dst, const T &src, const V &col_values) { + col_map(dst, src, col_values); +} + + + + + + + + + +template +__device__ static inline void broadcast_col(T &dst, const V &col_values) { + col_map(dst, dst, col_values); +} +template +__device__ static inline T broadcast_col(const V &col_values) { + T dst; + broadcast_col(dst, col_values); + return dst; +} + + +template +__device__ static inline void tril(RT &dst, const RT &src, int diagonal=0, const typename base_types::packing::unpacked_type &val=0) { + apply(dst, src, [val, diagonal]__device__(int row, int col, auto &src_val) { + return col <= row + diagonal ? src_val : val; + }); +} +template +__device__ static inline void triu(RT &dst, const RT &src, int diagonal=0, const typename base_types::packing::unpacked_type &val=0) { + apply(dst, src, [val, diagonal]__device__(int row, int col, auto &src_val) { + return col >= row + diagonal ? src_val : val; + }); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/reductions.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/reductions.cuh new file mode 100644 index 0000000..16322b1 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/reductions.cuh @@ -0,0 +1,540 @@ + + + + + + + + + + + + + + + + + + + +template +__device__ static inline void row_reduce(V &row_accum, const T &src, const V &src_accum) { + + static_assert(std::is_same_v::col_vec_layout>); + static_assert(std::is_same_v); + static_assert(V::outer_dim == T::height); + + using dtype = V::dtype; + + #pragma unroll + for(int i = 0; i < T::height; i++) { + dtype accum_top_row = op::template op(src.tiles[i][0].data[0], src.tiles[i][0].data[2]); + dtype accum_bottom_row = op::template op(src.tiles[i][0].data[1], src.tiles[i][0].data[3]); + #pragma unroll + for(int j = 1; j < T::width; j++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile; k+=2) { + accum_top_row = op::template op(accum_top_row, src.tiles[i][j].data[k+0]); + accum_bottom_row = op::template op(accum_bottom_row, src.tiles[i][j].data[k+1]); + } + } + dtype accum_packed; + accum_packed.x = op::template op::unpacked_type>(accum_top_row.x, accum_top_row.y); + accum_packed.y = op::template op::unpacked_type>(accum_bottom_row.x, accum_bottom_row.y); + + + + accum_packed = op::template op(accum_packed, packed_shfl_xor_sync(MASK_ALL, accum_packed, 2)); + accum_packed = op::template op(accum_packed, packed_shfl_xor_sync(MASK_ALL, accum_packed, 1)); + + if(reset) { + row_accum[i][0] = accum_packed; + } + else { + row_accum[i][0] = op::template op(src_accum[i][0], accum_packed); + } + } +} + + + + + + + + + + + + + + +template +__device__ static inline void row_reduce(V &row_accum, const T &src, const V &src_accum) { + + static_assert(std::is_same_v::col_vec_layout>); + static_assert(std::is_same_v); + static_assert(V::outer_dim == T::height); + + using dtype = V::dtype; + + #pragma unroll + for(int i = 0; i < T::height; i++) { + dtype accum_top_rows = op::template op(src.tiles[i][0].data[0], src.tiles[i][0].data[1]); + dtype accum_bottom_rows = op::template op(src.tiles[i][0].data[2], src.tiles[i][0].data[3]); + #pragma unroll + for(int j = 1; j < T::width; j++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile/2; k++) { + accum_top_rows = op::template op(accum_top_rows, src.tiles[i][j].data[k+0]); + accum_bottom_rows = op::template op(accum_bottom_rows, src.tiles[i][j].data[k+2]); + } + } + + + + accum_top_rows = op::template op(accum_top_rows, packed_shfl_xor_sync(MASK_ALL, accum_top_rows, 16)); + accum_top_rows = op::template op(accum_top_rows, packed_shfl_xor_sync(MASK_ALL, accum_top_rows, 8)); + accum_top_rows = op::template op(accum_top_rows, packed_shfl_xor_sync(MASK_ALL, accum_top_rows, 4)); + + accum_bottom_rows = op::template op(accum_bottom_rows, packed_shfl_xor_sync(MASK_ALL, accum_bottom_rows, 16)); + accum_bottom_rows = op::template op(accum_bottom_rows, packed_shfl_xor_sync(MASK_ALL, accum_bottom_rows, 8)); + accum_bottom_rows = op::template op(accum_bottom_rows, packed_shfl_xor_sync(MASK_ALL, accum_bottom_rows, 4)); + + if(reset) { + row_accum[i][0] = accum_top_rows; + row_accum[i][1] = accum_bottom_rows; + } + else { + row_accum[i][0] = op::template op(src_accum[i][0], accum_top_rows); + row_accum[i][1] = op::template op(src_accum[i][1], accum_bottom_rows); + } + } +} + + + + + + + + + + + + + + + + +template +__device__ static inline void col_reduce(V &col_accum, const T &src, const V &src_accum) { + + KITTENS_CHECK_WARP + static_assert(std::is_same_v::row_vec_layout>); + static_assert(std::is_same_v); + static_assert(V::outer_dim == T::width); + + using dtype = V::dtype; + + #pragma unroll + for(int j = 0; j < T::width; j++) { + dtype accum_left_cols = op::template op(src.tiles[0][j].data[0], src.tiles[0][j].data[1]); + dtype accum_right_cols = op::template op(src.tiles[0][j].data[2], src.tiles[0][j].data[3]); + #pragma unroll + for(int i = 1; i < T::height; i++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile/2; k++) { + accum_left_cols = op::template op(accum_left_cols, src.tiles[i][j].data[k+0]); + accum_right_cols = op::template op(accum_right_cols, src.tiles[i][j].data[k+2]); + } + } + + + + accum_left_cols = op::template op(accum_left_cols, packed_shfl_xor_sync(MASK_ALL, accum_left_cols, 16)); + accum_left_cols = op::template op(accum_left_cols, packed_shfl_xor_sync(MASK_ALL, accum_left_cols, 8)); + accum_left_cols = op::template op(accum_left_cols, packed_shfl_xor_sync(MASK_ALL, accum_left_cols, 4)); + + accum_right_cols = op::template op(accum_right_cols, packed_shfl_xor_sync(MASK_ALL, accum_right_cols, 16)); + accum_right_cols = op::template op(accum_right_cols, packed_shfl_xor_sync(MASK_ALL, accum_right_cols, 8)); + accum_right_cols = op::template op(accum_right_cols, packed_shfl_xor_sync(MASK_ALL, accum_right_cols, 4)); + + if(reset) { + col_accum[j][0] = accum_left_cols; + col_accum[j][1] = accum_right_cols; + } + else { + col_accum[j][0] = op::template op(src_accum[j][0], accum_left_cols); + col_accum[j][1] = op::template op(src_accum[j][1], accum_right_cols); + } + } +} + + + + + + + + + + + + + + +template +__device__ static inline void col_reduce(V &col_accum, const T &src, const V &src_accum) { + + KITTENS_CHECK_WARP + static_assert(std::is_same_v::row_vec_layout>); + static_assert(std::is_same_v); + static_assert(V::outer_dim == T::width); + + using dtype = V::dtype; + #pragma unroll + for(int j = 0; j < T::width; j++) { + dtype accum_left_col = op::template op(src.tiles[0][j].data[0], src.tiles[0][j].data[2]); + dtype accum_right_col = op::template op(src.tiles[0][j].data[1], src.tiles[0][j].data[3]); + #pragma unroll + for(int i = 1; i < T::height; i++) { + #pragma unroll + for(int k = 0; k < T::packed_per_tile; k+=2) { + accum_left_col = op::template op(accum_left_col, src.tiles[i][j].data[k+0]); + accum_right_col = op::template op(accum_right_col, src.tiles[i][j].data[k+1]); + } + } + dtype accum_packed; + accum_packed.x = op::template op::unpacked_type>(accum_left_col.x, accum_left_col.y); + accum_packed.y = op::template op::unpacked_type>(accum_right_col.x, accum_right_col.y); + + + + accum_packed = op::template op(accum_packed, packed_shfl_xor_sync(MASK_ALL, accum_packed, 2)); + accum_packed = op::template op(accum_packed, packed_shfl_xor_sync(MASK_ALL, accum_packed, 1)); + + if(reset) { + col_accum[j][0] = accum_packed; + } + else { + col_accum[j][0] = op::template op(src_accum[j][0], accum_packed); + } + } +} + + + + + + + + + + + + + +template +__device__ static inline void row_max(V &row_accum, const T &src) { + row_reduce(row_accum, src, row_accum); +} + + + + + + + + +template +__device__ static inline void row_min(V &row_accum, const T &src) { + row_reduce(row_accum, src, row_accum); +} + + + + + + + + +template +__device__ static inline void row_sum(V &row_accum, const T &src) { + row_reduce(row_accum, src, row_accum); +} + + + + + + + + +template +__device__ static inline void row_prod(V &row_accum, const T &src) { + row_reduce(row_accum, src, row_accum); +} + + + + + + + + + + +template +__device__ static inline void row_max(V &row_accum, const T &src, const V &src_accum) { + row_reduce(row_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void row_min(V &row_accum, const T &src, const V &src_accum) { + row_reduce(row_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void row_sum(V &row_accum, const T &src, const V &src_accum) { + row_reduce(row_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void row_prod(V &row_accum, const T &src, const V &src_accum) { + row_reduce(row_accum, src, src_accum); +} + + + + + + + + + + + +template +__device__ static inline void col_max(V &col_accum, const T &src) { + col_reduce(col_accum, src, col_accum); +} + + + + + + + + +template +__device__ static inline void col_min(V &col_accum, const T &src) { + col_reduce(col_accum, src, col_accum); +} + + + + + + + + +template +__device__ static inline void col_sum(V &col_accum, const T &src) { + col_reduce(col_accum, src, col_accum); +} + + + + + + + + +template +__device__ static inline void col_prod(V &col_accum, const T &src) { + col_reduce(col_accum, src, col_accum); +} + + + + + + + + + + +template +__device__ static inline void col_max(V &col_accum, const T &src, const V &src_accum) { + col_reduce(col_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void col_min(V &col_accum, const T &src, const V &src_accum) { + col_reduce(col_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void col_sum(V &col_accum, const T &src, const V &src_accum) { + col_reduce(col_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void col_prod(V &col_accum, const T &src, const V &src_accum) { + col_reduce(col_accum, src, src_accum); +} + + + +template +__device__ static inline void max(RV &dst, const T &src, const RV &src_accum) { + if constexpr (ax == axis::COL) row_max(dst, src, src_accum); + else col_max(dst, src, src_accum); +} +template +__device__ static inline auto max(const T &src, const RV &src_accum) { + RV dst; + if constexpr (ax == axis::COL) row_max(dst, src, src_accum); + else col_max(dst, src, src_accum); + return dst; +} +template +__device__ static inline void max(RV &dst, const T &src) { + if constexpr (ax == axis::COL) row_max(dst, src); + else col_max(dst, src); +} +template +__device__ static inline auto max(const T &src) { + using RV = std::conditional_t; + RV dst; + if constexpr (ax == axis::COL) row_max(dst, src); + else col_max(dst, src); + return dst; +} + +template +__device__ static inline void min(RV &dst, const T &src, const RV &src_accum) { + if constexpr (ax == axis::COL) row_min(dst, src, src_accum); + else col_min(dst, src, src_accum); +} +template +__device__ static inline auto min(const T &src, const RV &src_accum) { + RV dst; + if constexpr (ax == axis::COL) row_min(dst, src, src_accum); + else col_min(dst, src, src_accum); + return dst; +} +template +__device__ static inline void min(RV &dst, const T &src) { + if constexpr (ax == axis::COL) row_min(dst, src); + else col_min(dst, src); +} +template +__device__ static inline auto min(const T &src) { + using RV = std::conditional_t; + RV dst; + if constexpr (ax == axis::COL) row_min(dst, src); + else col_min(dst, src); + return dst; +} + +template +__device__ static inline void sum(RV &dst, const T &src, const RV &src_accum) { + if constexpr (ax == axis::COL) row_sum(dst, src, src_accum); + else col_sum(dst, src, src_accum); +} +template +__device__ static inline auto sum(const T &src, const RV &src_accum) { + RV dst; + if constexpr (ax == axis::COL) row_sum(dst, src, src_accum); + else col_sum(dst, src, src_accum); + return dst; +} +template +__device__ static inline void sum(RV &dst, const T &src) { + if constexpr (ax == axis::COL) row_sum(dst, src); + else col_sum(dst, src); +} +template +__device__ static inline auto sum(const T &src) { + using RV = std::conditional_t; + RV dst; + if constexpr (ax == axis::COL) row_sum(dst, src); + else col_sum(dst, src); + return dst; +} + +template +__device__ static inline void prod(RV &dst, const T &src, const RV &src_accum) { + if constexpr (ax == axis::COL) row_prod(dst, src, src_accum); + else col_prod(dst, src, src_accum); +} +template +__device__ static inline auto prod(const T &src, const RV &src_accum) { + RV dst; + if constexpr (ax == axis::COL) row_prod(dst, src, src_accum); + else col_prod(dst, src, src_accum); + return dst; +} +template +__device__ static inline void prod(RV &dst, const T &src) { + if constexpr (ax == axis::COL) row_prod(dst, src); + else col_prod(dst, src); +} +template +__device__ static inline auto prod(const T &src) { + using RV = std::conditional_t; + RV dst; + if constexpr (ax == axis::COL) row_prod(dst, src); + else col_prod(dst, src); + return dst; +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/tile.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/tile.cuh new file mode 100644 index 0000000..5d736b6 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/tile.cuh @@ -0,0 +1,46 @@ + + + + + +#include "conversions.cuh" +#include "maps.cuh" +#include "reductions.cuh" + +template +__device__ static inline bool hasnan(const RT &src) { + KITTENS_CHECK_WARP + bool nan_detected = false; + #pragma unroll + for(int i = 0; i < RT::height; i++) { + #pragma unroll + for(int j = 0; j < RT::width; j++) { + #pragma unroll + for(int k = 0; k < RT::packed_per_tile; k++) { + if constexpr (std::is_same_v) { + if(isnan(src.tiles[i][j].data[k].x) || isnan(src.tiles[i][j].data[k].y)) { + nan_detected = true; + } + } + else if constexpr (std::is_same_v) { + if(isnan(__bfloat162float(src.tiles[i][j].data[k].x)) || isnan(__bfloat162float(src.tiles[i][j].data[k].y))) { + nan_detected = true; + } + } + else if constexpr (std::is_same_v) { + if(isnan(__half2float(src.tiles[i][j].data[k].x)) || isnan(__half2float(src.tiles[i][j].data[k].y))) { + nan_detected = true; + } + } + else { + static_assert(sizeof(typename RT::T) == 999, "Unsupported dtype"); + } + } + } + } + + return (__ballot_sync(0xffffffff, nan_detected) != 0); +} + +#include "complex_conversions.cuh" +#include "complex_maps.cuh" diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/vec/conversions.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/vec/conversions.cuh new file mode 100644 index 0000000..ab37d25 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/vec/conversions.cuh @@ -0,0 +1,153 @@ + + + + + +struct vec_conversion_detail { + + + +__device__ static inline int row_from_indices_dim2(int laneid, int inner_dim, int x_or_y) { + return 8*inner_dim + (laneid%4)*2 + x_or_y; +} +__device__ static inline int row_from_indices_dim1(int laneid, int x_or_y) { + return 8*x_or_y + (laneid/4); +} +__device__ static inline int canonical_src_lane_dim2(int row) { + return (row/2)%4 + 4*(row%2); +} +__device__ static inline int canonical_src_lane_dim1(int row) { + return (row*4)%32; +} + +}; + + + + + + + + + +template +__device__ static inline void copy(RV1 &dst, const RV2 &src) { + KITTENS_CHECK_WARP + static_assert(RV1::length == RV2::length, "Register vectors must be the same length."); + using D1 = RV1::dtype; + using D2 = RV2::dtype; + if constexpr (std::is_same_v) { + #pragma unroll + for(int i = 0; i < RV1::outer_dim; i++) { + #pragma unroll + for(int j = 0; j < RV1::inner_dim; j++) { + dst[i][j] = base_types::convertor::convert(src[i][j]); + } + } + } + else { + int laneid = ::kittens::laneid(); + if constexpr (std::is_same_v && std::is_same_v) { + #pragma unroll + for(int i = 0; i < RV1::outer_dim; i++) { + dst[i][0].x = packed_shfl_sync( + kittens::MASK_ALL, + laneid < 4 ? src[i][0].x : src[i][0].y, + vec_conversion_detail::canonical_src_lane_dim2(vec_conversion_detail::row_from_indices_dim1(laneid, 0)) + ); + dst[i][0].y = packed_shfl_sync( + kittens::MASK_ALL, + laneid < 4 ? src[i][1].x : src[i][1].y, + vec_conversion_detail::canonical_src_lane_dim2(vec_conversion_detail::row_from_indices_dim1(laneid, 1)) + ); + } + } + else if constexpr (std::is_same_v && std::is_same_v) { + #pragma unroll + for(int i = 0; i < RV1::outer_dim; i++) { + dst[i][0].x = packed_shfl_sync( + kittens::MASK_ALL, + src[i][0].x, + vec_conversion_detail::canonical_src_lane_dim1(vec_conversion_detail::row_from_indices_dim2(laneid, 0, 0)) + ); + dst[i][0].y = packed_shfl_sync( + kittens::MASK_ALL, + src[i][0].x, + vec_conversion_detail::canonical_src_lane_dim1(vec_conversion_detail::row_from_indices_dim2(laneid, 0, 1)) + ); + dst[i][1].x = packed_shfl_sync( + kittens::MASK_ALL, + src[i][0].y, + vec_conversion_detail::canonical_src_lane_dim1(vec_conversion_detail::row_from_indices_dim2(laneid, 1, 0)) + ); + dst[i][1].y = packed_shfl_sync( + kittens::MASK_ALL, + src[i][0].y, + vec_conversion_detail::canonical_src_lane_dim1(vec_conversion_detail::row_from_indices_dim2(laneid, 1, 1)) + ); + } + } + else if constexpr (std::is_same_v && std::is_same_v) { + #pragma unroll + for(int i = 0; i < RV1::outer_dim; i++) { + dst[i][0].x = packed_shfl_sync( + kittens::MASK_ALL, src[i/2][0], + 16*(i%2) + 0 + (laneid/4) + ); + dst[i][0].y = packed_shfl_sync( + kittens::MASK_ALL, src[i/2][0], + 16*(i%2) + 8 + (laneid/4) + ); + } + } + else if constexpr (std::is_same_v && std::is_same_v) { + int lane_replication = laneid%4; + #pragma unroll + for(int i = 0; i < RV1::outer_dim; i++) { + D1 tmp = 0; + if(RV1::length%32==0 || i < RV1::outer_dim-1 || lane_replication<2) { + tmp = lane_replication%2 ? src[2*i + (lane_replication>=2)][0].y : src[2*i + (lane_replication>=2)][0].x; + } + dst[i][0] = packed_shfl_sync( + kittens::MASK_ALL, tmp, + (laneid%8)*4 + (laneid/8) + ); + } + } + else if constexpr (std::is_same_v && std::is_same_v) { + #pragma unroll + for(int i = 0; i < RV1::outer_dim; i++) { + dst[i][0].x = packed_shfl_sync( + kittens::MASK_ALL, src[i/2][0], + 16*(i%2) + 0 + 2*(laneid%4) + 0 + ); + dst[i][0].y = packed_shfl_sync( + kittens::MASK_ALL, src[i/2][0], + 16*(i%2) + 0 + 2*(laneid%4) + 1 + ); + dst[i][1].x = packed_shfl_sync( + kittens::MASK_ALL, src[i/2][0], + 16*(i%2) + 8 + 2*(laneid%4) + 0 + ); + dst[i][1].y = packed_shfl_sync( + kittens::MASK_ALL, src[i/2][0], + 16*(i%2) + 8 + 2*(laneid%4) + 1 + ); + } + } + else if constexpr (std::is_same_v && std::is_same_v) { + int lane_replication = laneid/8; + #pragma unroll + for(int i = 0; i < RV1::outer_dim; i++) { + D1 tmp = 0; + if(RV1::length%32==0 || i < RV1::outer_dim-1 || laneid<16) { + tmp = (laneid%8)<4 ? src[2*i + (lane_replication>=2)][lane_replication%2].x : src[2*i + (lane_replication>=2)][lane_replication%2].y; + } + dst[i][0] = packed_shfl_sync( + kittens::MASK_ALL, tmp, + 4*(laneid%2) + (laneid%8)/2 + (laneid&0b11000) + ); + } + } + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/vec/maps.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/vec/maps.cuh new file mode 100644 index 0000000..91fefb2 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/vec/maps.cuh @@ -0,0 +1,374 @@ + + + + + + + + + + + + + + + +template +__device__ static inline void unary_op(T &dst, const T &src) { + #pragma unroll + for(int i = 0; i < T::outer_dim; i++) { + #pragma unroll + for(int j = 0; j < T::inner_dim; j++) { + dst[i][j] = op::template op(src[i][j]); + } + } +} + + + + + + + + + +template +__device__ static inline void bin_op(T &dst, const T &lhs, const T &rhs) { + #pragma unroll + for(int i = 0; i < T::outer_dim; i++) { + #pragma unroll + for(int j = 0; j < T::inner_dim; j++) { + dst[i][j] = op::template op(lhs[i][j], rhs[i][j]); + } + } +} + + + + + + + + + +template +__device__ static inline void bin_op(T &dst, const T &src, const typename T::dtype ¶m) { + #pragma unroll + for(int i = 0; i < T::outer_dim; i++) { + #pragma unroll + for(int j = 0; j < T::inner_dim; j++) { + dst[i][j] = op::template op(src[i][j], param); + } + } +} + + + + + + + + + +template +__device__ static inline void bin_op(T &dst, const T &src, const typename base_types::packing::unpacked_type ¶m) { + bin_op(dst, src, base_types::packing::pack(param)); +} + + +template +__device__ static inline void apply(RV &dst, const RV &src, Lambda &&lambda) { + int group_offset = 0; + if constexpr(GROUP_WARPS > 1) { + group_offset = warpid()*RV::length; + } + static_assert(sizeof(RV::T) != 1, "Cannot apply lambda to 8-bit types"); + if constexpr (ducks::rv::ortho_layout) { + #pragma unroll + for(int i = 0; i < RV::outer_dim; i++) { + int base_idx = group_offset + i*16 + ::kittens::laneid()/4; + dst[i][0].x = lambda(base_idx+0, src[i][0].x); + dst[i][0].y = lambda(base_idx+8, src[i][0].y); + } + } + else if constexpr (ducks::rv::align_layout) { + #pragma unroll + for(int i = 0; i < RV::outer_dim; i++) { + int base_idx = group_offset + i*16 + 2*(::kittens::laneid()%4); + dst[i][0].x = lambda(base_idx+0, src[i][0].x); + dst[i][0].y = lambda(base_idx+1, src[i][0].y); + dst[i][1].x = lambda(base_idx+8, src[i][1].x); + dst[i][1].y = lambda(base_idx+9, src[i][1].y); + } + } + else { + #pragma unroll + for(int i = 0; i < RV::outer_dim; i++) { + int base_idx = group_offset + i*32 + ::kittens::laneid(); + if (i < RV::outer_dim-1 || RV::length%32 == 0 || ::kittens::laneid()<16) { + dst[i][0] = lambda(base_idx, src[i][0]); + } + } + } +} +template +__device__ static inline RV apply(const RV &src, Lambda &&lambda) { + RV dst; + apply(dst, src, std::forward(lambda)); + return dst; +} + + + + + + + + + + + +template +__device__ static inline void zero(T &dst) { + unary_op(dst, dst); +} + + + + + + +template +__device__ static inline void one(T &dst) { + unary_op(dst, dst); +} + + + + + + +template +__device__ static inline void pos_infty(T &dst) { + unary_op(dst, dst); +} + + + + + + +template +__device__ static inline void neg_infty(T &dst) { + unary_op(dst, dst); +} + + + + + + + + + + + +template +__device__ static inline void copy(T &dst, const U &src) { + bin_op(dst, dst, src); +} + + + + + + + +template +__device__ static inline void exp(T &dst, const T &src) { + unary_op(dst, src); +} +template +__device__ static inline T exp(const T &src) { + T dst; + exp(dst, src); + return dst; +} + + + + + + + +template +__device__ static inline void exp2(T &dst, const T &src) { + unary_op(dst, src); +} +template +__device__ static inline T exp2(const T &src) { + T dst; + exp2(dst, src); + return dst; +} + + + + + + + +template +__device__ static inline void log(T &dst, const T &src) { + unary_op(dst, src); +} +template +__device__ static inline T log(const T &src) { + T dst; + log(dst, src); + return dst; +} + + + + + + + +template +__device__ static inline void log2(T &dst, const T &src) { + unary_op(dst, src); +} +template +__device__ static inline T log2(const T &src) { + T dst; + log2(dst, src); + return dst; +} + + + + + + + +template +__device__ static inline void abs(T &dst, const T &src) { + unary_op(dst, src); +} +template +__device__ static inline T abs(const T &src) { + T dst; + abs(dst, src); + return dst; +} + + + + + + + +template +__device__ static inline void relu(T &dst, const T &src) { + unary_op(dst, src); +} +template +__device__ static inline T relu(const T &src) { + T dst; + relu(dst, src); + return dst; +} + + + + + + + + + + + + +template +__device__ static inline void max(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} +template +__device__ static inline T max(const T &lhs, const U &rhs) { + T dst; + max(dst, lhs, rhs); + return dst; +} + + + + + + + + + +template +__device__ static inline void min(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} +template +__device__ static inline T min(const T &lhs, const U &rhs) { + T dst; + min(dst, lhs, rhs); + return dst; +} + + + + + + + + + +template +__device__ static inline void add(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void sub(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void mul(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void div(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/vec/reductions.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/vec/reductions.cuh new file mode 100644 index 0000000..62808d0 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/vec/reductions.cuh @@ -0,0 +1,233 @@ + + + + + + + + + + + + + + + + + + + + + +template +__device__ static inline void reduce( + typename base_types::packing::unpacked_type &dst_accum, + const RV &src, + const typename base_types::packing::unpacked_type &src_accum) { + KITTENS_CHECK_WARP + using T = base_types::packing::unpacked_type; + int laneid = kittens::laneid(); + if constexpr (std::is_same_v) { + T accum = op::template op(src[0][0].x, src[0][0].y); + #pragma unroll + for(int i = 1; i < RV::outer_dim; i++) { + accum = op::template op(accum, src[i][0].x); + accum = op::template op(accum, src[i][0].y); + } + + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 16)); + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 8)); + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 4)); + + if constexpr (!reset) accum = op::template op(accum, src_accum); + + dst_accum = packed_shfl_sync(kittens::MASK_ALL, accum, 0); + } + else if constexpr (std::is_same_v) { + T accum = op::template op(src[0][0].x, src[0][0].y); + accum = op::template op(accum, src[0][1].x); + accum = op::template op(accum, src[0][1].y); + #pragma unroll + for(int i = 1; i < RV::outer_dim; i++) { + + accum = op::template op(accum, src[i][0].x); + accum = op::template op(accum, src[i][0].y); + accum = op::template op(accum, src[i][1].x); + accum = op::template op(accum, src[i][1].y); + } + + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 2)); + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 1)); + + if constexpr (!reset) accum = op::template op(accum, src_accum); + + dst_accum = packed_shfl_sync(kittens::MASK_ALL, accum, 0); + } + else if constexpr (std::is_same_v) { + T accum = src[0][0]; + #pragma unroll + for(int i = 1; i < RV::outer_dim; i++) { + if (i < RV::outer_dim-1 || i*kittens::TILE_ROW_DIM*2 + laneid < RV::length) { + accum = op::template op(accum, src[i][0]); + } + } + if(RV::length > 16) accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 16)); + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 8)); + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 4)); + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 2)); + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 1)); + if constexpr (!reset) accum = op::template op(accum, src_accum); + dst_accum = packed_shfl_sync(kittens::MASK_ALL, accum, 0); + } +} + + + + + + + + + +template +__device__ static inline void max(typename base_types::packing::unpacked_type &max_val, const RV &src) { + reduce(max_val, src, max_val); +} +template +__device__ static inline typename base_types::packing::unpacked_type max(const RV &src) { + typename base_types::packing::unpacked_type max_val; + reduce(max_val, src, max_val); + return max_val; +} + + + + + + + + +template +__device__ static inline void min(typename base_types::packing::unpacked_type &min_val, const RV &src) { + reduce(min_val, src, min_val); +} +template +__device__ static inline typename base_types::packing::unpacked_type min(const RV &src) { + typename base_types::packing::unpacked_type min_val; + reduce(min_val, src, min_val); + return min_val; +} + + + + + + + + +template +__device__ static inline void sum(typename base_types::packing::unpacked_type &sum_val, const RV &src) { + reduce(sum_val, src, sum_val); +} +template +__device__ static inline typename base_types::packing::unpacked_type sum(const RV &src) { + typename base_types::packing::unpacked_type sum_val; + reduce(sum_val, src, sum_val); + return sum_val; +} + + + + + + + + +template +__device__ static inline void prod(typename base_types::packing::unpacked_type &prod_val, const RV &src) { + reduce(prod_val, src, prod_val); +} +template +__device__ static inline typename base_types::packing::unpacked_type prod(const RV &src) { + typename base_types::packing::unpacked_type prod_val; + reduce(prod_val, src, prod_val); + return prod_val; +} + + + + + + + + + + + +template +__device__ static inline void max(typename base_types::packing::unpacked_type &max_val, const RV &src, const typename base_types::packing::unpacked_type &src_accum) { + reduce(max_val, src, src_accum); +} +template +__device__ static inline typename base_types::packing::unpacked_type max(const RV &src, const typename base_types::packing::unpacked_type &src_accum) { + typename base_types::packing::unpacked_type max_val; + reduce(max_val, src, src_accum); + return max_val; +} + + + + + + + + + +template +__device__ static inline void min(typename base_types::packing::unpacked_type &min_val, const RV &src, const typename base_types::packing::unpacked_type &src_accum) { + reduce(min_val, src, src_accum); +} +template +__device__ static inline typename base_types::packing::unpacked_type min(const RV &src, const typename base_types::packing::unpacked_type &src_accum) { + typename base_types::packing::unpacked_type min_val; + reduce(min_val, src, src_accum); + return min_val; +} + + + + + + + + + +template +__device__ static inline void sum(typename base_types::packing::unpacked_type &sum_val, const RV &src, const typename base_types::packing::unpacked_type &src_accum) { + reduce(sum_val, src, src_accum); +} +template +__device__ static inline typename base_types::packing::unpacked_type sum(const RV &src, const typename base_types::packing::unpacked_type &src_accum) { + typename base_types::packing::unpacked_type sum_val; + reduce(sum_val, src, src_accum); + return sum_val; +} + + + + + + + + + +template +__device__ static inline void prod(typename base_types::packing::unpacked_type &prod_val, const RV &src, const typename base_types::packing::unpacked_type &src_accum) { + reduce(prod_val, src, src_accum); +} +template +__device__ static inline typename base_types::packing::unpacked_type prod(const RV &src, const typename base_types::packing::unpacked_type &src_accum) { + typename base_types::packing::unpacked_type prod_val; + reduce(prod_val, src, src_accum); + return prod_val; +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/vec/vec.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/vec/vec.cuh new file mode 100644 index 0000000..fa67f5d --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/vec/vec.cuh @@ -0,0 +1,59 @@ + + + + + +#include "conversions.cuh" +#include "maps.cuh" +#include "reductions.cuh" + +template +__device__ static inline bool hasnan(const RV &src) { + KITTENS_CHECK_WARP + bool nan_detected = false; + #pragma unroll + for(int i = 0; i < RV::outer_dim; i++) { + #pragma unroll + for(int j = 0; j < RV::inner_dim; j++) { + if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { + if(isnan(src[i][j].x) || isnan(src[i][j].y)) { + nan_detected = true; + } + } + else if constexpr (std::is_same_v) { + if(isnan(__bfloat162float(src[i][j].x)) || isnan(__bfloat162float(src[i][j].y))) { + nan_detected = true; + } + } + else if constexpr (std::is_same_v) { + if(isnan(__half2float(src[i][j].x)) || isnan(__half2float(src[i][j].y))) { + nan_detected = true; + } + } + } + else if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { + if(isnan(src[i][j])) { + nan_detected = true; + } + } + else if constexpr (std::is_same_v) { + if(isnan(__bfloat162float(src[i][j]))) { + nan_detected = true; + } + } + else if constexpr (std::is_same_v) { + if(isnan(__half2float(src[i][j]))) { + nan_detected = true; + } + } + } + else { + static_assert(sizeof(typename RV::dtype) == 999, "Unsupported dtype"); + } + } + } + + return (__ballot_sync(0xffffffff, nan_detected) != 0); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/shared.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/shared.cuh new file mode 100644 index 0000000..404f8af --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/shared.cuh @@ -0,0 +1,7 @@ + + + + + +#include "tile/tile.cuh" +#include "vec/vec.cuh" diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/tile/conversions.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/tile/conversions.cuh new file mode 100644 index 0000000..45810c1 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/tile/conversions.cuh @@ -0,0 +1,16 @@ + + + + + + + +template +__device__ static inline void copy(ST1 &dst, const ST2 &src) { + static_assert(ST1::rows == ST2::rows && ST1::cols == ST2::cols, "Tiles must have the same rows and cols"); + #pragma unroll + for(int i = laneid(); i < ST1::num_elements; i+=GROUP_THREADS) { + int row = i/ST1::cols, col = i%ST1::cols; + dst[{row, col}] = base_types::convertor::convert(src[{row, col}]); + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/tile/maps.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/tile/maps.cuh new file mode 100644 index 0000000..c5c1b36 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/tile/maps.cuh @@ -0,0 +1,487 @@ + + + + + + + + + + + + + + + + +template +__device__ static inline void unary_map(T &dst, const T &src) { + #pragma unroll + for(int i = laneid(); i < T::num_elements; i += GROUP_THREADS) { + dst.data[i] = op::template op(src.data[i]); + } +} + + + + + + + + + + + + + +template +__device__ static inline void bin_map(T &dst, const T &src, const typename T::dtype ¶m) { + #pragma unroll + for(int i = laneid(); i < T::num_elements; i += GROUP_THREADS) { + dst.data[i] = op::template op(src.data[i], param); + } +} + + + + + + + + + + + + + +template +__device__ static inline void bin_map(T &dst, const T &lhs, const T &rhs) { + #pragma unroll + for(int i = laneid(); i < T::num_elements; i += GROUP_THREADS) { + dst.data[i] = op::template op(lhs.data[i], rhs.data[i]); + } +} + + + + + + + + + + + + + + + +template +__device__ static inline void row_map(T &dst, const T &src, const V &vec) { + static_assert(std::is_same::value, "Tile and vector must have the same data type"); + static_assert(V::length == T::rows, "Vector length must match the number of rows in the tile"); + #pragma unroll + for(int i = laneid(); i < T::num_elements; i += GROUP_THREADS) { + int row = i/T::cols, col = i%T::cols; + dst[{row, col}] = op::template op(src[{row, col}], vec[row]); + } +} + + + + + + + + + + + + + + + +template +__device__ static inline void col_map(T &dst, const T &src, const V &vec) { + static_assert(std::is_same::value, "Tile and vector must have the same data type"); + static_assert(V::length == T::cols, "Vector length must match the number of columns in the tile"); + #pragma unroll + for(int i = laneid(); i < T::num_elements; i += GROUP_THREADS) { + int row = i/T::cols, col = i%T::cols; + dst[{row, col}] = op::template op(src[{row, col}], vec[col]); + } +} + + + + + + + + + + + + + + +template +__device__ static inline void zero(T &dst) { + unary_map(dst, dst); +} + + + + + + +template +__device__ static inline void one(T &dst) { + unary_map(dst, dst); +} + + + + + + +template +__device__ static inline void pos_infty(T &dst) { + unary_map(dst, dst); +} + + + + + + +template +__device__ static inline void neg_infty(T &dst) { + unary_map(dst, dst); +} + + + + + + + + + +template +__device__ static inline void exp(T &dst, const T &src) { + unary_map(dst, src); +} + + + + + + + +template +__device__ static inline void exp2(T &dst, const T &src) { + unary_map(dst, src); +} + + + + + + + +template +__device__ static inline void log(T &dst, const T &src) { + unary_map(dst, src); +} + + + + + + + +template +__device__ static inline void log2(T &dst, const T &src) { + unary_map(dst, src); +} + + + + + + + +template +__device__ static inline void abs(T &dst, const T &src) { + unary_map(dst, src); +} + + + + + + + +template +__device__ static inline void relu(T &dst, const T &src) { + unary_map(dst, src); +} + + + + + + + + +template +__device__ static inline void copy(T &dst, const U &src) { + bin_map(dst, src); +} + + + + + + + + + + + +template +__device__ static inline void max(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void min(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void add(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void sub(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void mul(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void div(T &dst, const T &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + + + + + + + + + + + + +template +__device__ static inline void add_row(T &dst, const T &src, const V &row_values) { + row_map(dst, src, row_values); +} + + + + + + + + + +template +__device__ static inline void sub_row(T &dst, const T &src, const V &row_values) { + row_map(dst, src, row_values); +} + + + + + + + + + +template +__device__ static inline void mul_row(T &dst, const T &src, const V &row_values) { + row_map(dst, src, row_values); +} + + + + + + + + + +template +__device__ static inline void div_row(T &dst, const T &src, const V &row_values) { + row_map(dst, src, row_values); +} + + + + + + + + +template +__device__ static inline void broadcast_row(T &dst, const V &row_values) { + row_map(dst, dst, row_values); +} + + + + + + + + + + + + +template +__device__ static inline void add_col(T &dst, const T &src, const V &col_values) { + col_map(dst, src, col_values); +} + + + + + + + + + +template +__device__ static inline void sub_col(T &dst, const T &src, const V &col_values) { + col_map(dst, src, col_values); +} + + + + + + + + + +template +__device__ static inline void mul_col(T &dst, const T &src, const V &col_values) { + col_map(dst, src, col_values); +} + + + + + + + + + +template +__device__ static inline void div_col(T &dst, const T &src, const V &col_values) { + col_map(dst, src, col_values); +} + + + + + + + + +template +__device__ static inline void broadcast_col(T &dst, const V &col_values) { + col_map(dst, dst, col_values); +} + + + +template +__device__ static inline void add(T &dst, const T &src, const V &col_values) { + if constexpr (axis == axis::COL) add_col(dst, src, col_values); + else add_row(dst, src, col_values); +} + +template +__device__ static inline void sub(T &dst, const T &src, const V &col_values) { + if constexpr (axis == axis::COL) sub_col(dst, src, col_values); + else sub_row(dst, src, col_values); +} + +template +__device__ static inline void mul(T &dst, const T &src, const V &col_values) { + if constexpr (axis == axis::COL) mul_col(dst, src, col_values); + else mul_row(dst, src, col_values); +} + +template +__device__ static inline void div(T &dst, const T &src, const V &col_values) { + if constexpr (axis == axis::COL) div_col(dst, src, col_values); + else div_row(dst, src, col_values); +} + +template +__device__ static inline void broadcast(T &dst, const V &col_values) { + if constexpr (axis == axis::COL) broadcast_col(dst, col_values); + else broadcast_row(dst, col_values); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/tile/reductions.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/tile/reductions.cuh new file mode 100644 index 0000000..b4926e1 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/tile/reductions.cuh @@ -0,0 +1,372 @@ + + + + + + + + + + + + + + + + +template +__device__ static inline void row_reduce(V &row_accum, const T &src, const V &src_accum) { + using dtype = typename V::dtype; + for (int row = laneid(); row < T::rows; row += GROUP_THREADS) { + dtype accum = src[{row, 0}]; + #pragma unroll + for (int col = 1; col < T::cols; col++) { + accum = op::template op(accum, src[{row, col}]); + } + if (reset) { + row_accum[row] = accum; + } else { + row_accum[row] = op::template op(src_accum[row], accum); + } + } +} + + + + + + + + + + + + +template +__device__ static inline void col_reduce(V &col_accum, const T &src, const V &src_accum) { + using dtype = typename V::dtype; + for (int col = laneid(); col < T::cols; col += GROUP_THREADS) { + dtype accum = src[{0, col}]; + #pragma unroll + for (int row = 1; row < T::rows; row++) { + accum = op::template op(accum, src[{row, col}]); + } + if (reset) { + col_accum[col] = accum; + } else { + col_accum[col] = op::template op(src_accum[col], accum); + } + } +} + + + + + + + + + + + +template +__device__ static inline void row_max(V &row_accum, const T &src) { + row_reduce(row_accum, src, row_accum); +} + + + + + + + + +template +__device__ static inline void row_min(V &row_accum, const T &src) { + row_reduce(row_accum, src, row_accum); +} + + + + + + + + +template +__device__ static inline void row_sum(V &row_accum, const T &src) { + row_reduce(row_accum, src, row_accum); +} + + + + + + + + +template +__device__ static inline void row_prod(V &row_accum, const T &src) { + row_reduce(row_accum, src, row_accum); +} + + + + + + + + + + +template +__device__ static inline void row_max(V &row_accum, const T &src, const V &src_accum) { + row_reduce(row_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void row_min(V &row_accum, const T &src, const V &src_accum) { + row_reduce(row_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void row_sum(V &row_accum, const T &src, const V &src_accum) { + row_reduce(row_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void row_prod(V &row_accum, const T &src, const V &src_accum) { + row_reduce(row_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void col_max(V &col_accum, const T &src) { + col_reduce(col_accum, src, col_accum); +} + + + + + + + + +template +__device__ static inline void col_min(V &col_accum, const T &src) { + col_reduce(col_accum, src, col_accum); +} + + + + + + + + +template +__device__ static inline void col_sum(V &col_accum, const T &src) { + col_reduce(col_accum, src, col_accum); +} + + + + + + + + +template +__device__ static inline void col_prod(V &col_accum, const T &src) { + col_reduce(col_accum, src, col_accum); +} + + + + + + + + + + +template +__device__ static inline void col_max(V &col_accum, const T &src, const V &src_accum) { + col_reduce(col_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void col_min(V &col_accum, const T &src, const V &src_accum) { + col_reduce(col_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void col_sum(V &col_accum, const T &src, const V &src_accum) { + col_reduce(col_accum, src, src_accum); +} + + + + + + + + + +template +__device__ static inline void col_prod(V &col_accum, const T &src, const V &src_accum) { + col_reduce(col_accum, src, src_accum); +} + + + +template +__device__ static inline void max(V &dst, const T &src, const V &src_accum) { + if constexpr (ax == axis::COL) row_max(dst, src, src_accum); + else col_max(dst, src, src_accum); +} +template +__device__ static inline auto max(const T &src, const V &src_accum) { + V dst; + if constexpr (ax == axis::COL) row_max(dst, src, src_accum); + else col_max(dst, src, src_accum); + return dst; +} +template +__device__ static inline void max(V &dst, const T &src) { + if constexpr (ax == axis::COL) row_max(dst, src); + else col_max(dst, src); +} +template +__device__ static inline auto max(const T &src) { + using V = std::conditional_t; + V dst; + if constexpr (ax == axis::COL) row_max(dst, src); + else col_max(dst, src); + return dst; +} + +template +__device__ static inline void min(V &dst, const T &src, const V &src_accum) { + if constexpr (ax == axis::COL) row_min(dst, src, src_accum); + else col_min(dst, src, src_accum); +} +template +__device__ static inline auto min(const T &src, const V &src_accum) { + V dst; + if constexpr (ax == axis::COL) row_min(dst, src, src_accum); + else col_min(dst, src, src_accum); + return dst; +} +template +__device__ static inline void min(V &dst, const T &src) { + if constexpr (ax == axis::COL) row_min(dst, src); + else col_min(dst, src); +} +template +__device__ static inline auto min(const T &src) { + using V = std::conditional_t; + V dst; + if constexpr (ax == axis::COL) row_min(dst, src); + else col_min(dst, src); + return dst; +} + +template +__device__ static inline void sum(V &dst, const T &src, const V &src_accum) { + if constexpr (ax == axis::COL) row_sum(dst, src, src_accum); + else col_sum(dst, src, src_accum); +} +template +__device__ static inline auto sum(const T &src, const V &src_accum) { + V dst; + if constexpr (ax == axis::COL) row_sum(dst, src, src_accum); + else col_sum(dst, src, src_accum); + return dst; +} +template +__device__ static inline void sum(V &dst, const T &src) { + if constexpr (ax == axis::COL) row_sum(dst, src); + else col_sum(dst, src); +} +template +__device__ static inline auto sum(const T &src) { + using V = std::conditional_t; + V dst; + if constexpr (ax == axis::COL) row_sum(dst, src); + else col_sum(dst, src); + return dst; +} + +template +__device__ static inline void prod(V &dst, const T &src, const V &src_accum) { + if constexpr (ax == axis::COL) row_prod(dst, src, src_accum); + else col_prod(dst, src, src_accum); +} +template +__device__ static inline auto prod(const T &src, const V &src_accum) { + V dst; + if constexpr (ax == axis::COL) row_prod(dst, src, src_accum); + else col_prod(dst, src, src_accum); + return dst; +} +template +__device__ static inline void prod(V &dst, const T &src) { + if constexpr (ax == axis::COL) row_prod(dst, src); + else col_prod(dst, src); +} +template +__device__ static inline auto prod(const T &src) { + using V = std::conditional_t; + V dst; + if constexpr (ax == axis::COL) row_prod(dst, src); + else col_prod(dst, src); + return dst; +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/tile/tile.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/tile/tile.cuh new file mode 100644 index 0000000..cf76bda --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/tile/tile.cuh @@ -0,0 +1,37 @@ + + + + + +#include "conversions.cuh" +#include "maps.cuh" +#include "reductions.cuh" + +template +__device__ static inline bool hasnan(const ST &src) { + KITTENS_CHECK_WARP + bool nan_detected = false; + #pragma unroll + for(int i = laneid(); i < ST::num_elements; i+=GROUP_THREADS) { + if constexpr (std::is_same_v) { + if(isnan(src[i])) { + nan_detected = true; + } + } + else if constexpr (std::is_same_v) { + if(isnan(__bfloat162float(src[i]))) { + nan_detected = true; + } + } + else if constexpr (std::is_same_v) { + if(isnan(__half2float(src[i]))) { + nan_detected = true; + } + } + else { + static_assert(sizeof(typename ST::T) == 999, "Unsupported dtype"); + } + } + + return (__ballot_sync(0xffffffff, nan_detected) != 0); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/vec/conversions.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/vec/conversions.cuh new file mode 100644 index 0000000..5a0fa4a --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/vec/conversions.cuh @@ -0,0 +1,27 @@ + + + + + + + + + + + + + + + + + + + +template +__device__ static inline void copy(SV1 &dst, const SV2 &src) { + static_assert(SV1::length == SV2::length, "Source and destination vectors must have the same length."); + #pragma unroll + for(int i = laneid(); i < SV1::length; i+=GROUP_THREADS) { + dst[i] = base_types::convertor::convert(src[i]); + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/vec/maps.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/vec/maps.cuh new file mode 100644 index 0000000..e491c3e --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/vec/maps.cuh @@ -0,0 +1,259 @@ + + + + + + + + + + + + + +template +__device__ static inline void unary_op(T &dst, const T &src) { + #pragma unroll + for(auto cur = laneid(); cur < T::length; cur+=GROUP_THREADS) { + dst[cur] = op::template op(src[cur]); + } +} + + + + + + + + + +template +__device__ static inline void bin_op(T &dst, const T &lhs, const T &rhs) { + #pragma unroll + for(auto cur = laneid(); cur < T::length; cur+=GROUP_THREADS) { + dst[cur] = op::template op(lhs[cur], rhs[cur]); + } +} + + + + + + + + + +template +__device__ static inline void bin_op(T &dst, const T &src, const typename T::dtype ¶m) { + #pragma unroll + for(auto cur = laneid(); cur < T::length; cur+=GROUP_THREADS) { + dst[cur] = op::template op(src[cur], param); + } +} + + + + + + + + + + + +template +__device__ static inline void zero(T &dst) { + unary_op(dst, dst); +} + + + + + + +template +__device__ static inline void one(T &dst) { + unary_op(dst, dst); +} + + + + + + +template +__device__ static inline void pos_infty(T &dst) { + unary_op(dst, dst); +} + + + + + + +template +__device__ static inline void neg_infty(T &dst) { + unary_op(dst, dst); +} + + + + + + + + + + + +template +__device__ static inline void copy(T &dst, const U &src) { + bin_op(dst, dst, src); +} + + + + + + + +template +__device__ static inline void exp(T &dst, const T &src) { + unary_op(dst, src); +} + + + + + + + +template +__device__ static inline void exp2(T &dst, const T &src) { + unary_op(dst, src); +} + + + + + + + +template +__device__ static inline void log(T &dst, const T &src) { + unary_op(dst, src); +} + + + + + + + +template +__device__ static inline void log2(T &dst, const T &src) { + unary_op(dst, src); +} + + + + + + + +template +__device__ static inline void abs(T &dst, const T &src) { + unary_op(dst, src); +} + + + + + + + +template +__device__ static inline void relu(T &dst, const T &src) { + unary_op(dst, src); +} + + + + + + + + + + + + +template +__device__ static inline void max(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void min(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void add(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void sub(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void mul(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} + + + + + + + + + +template +__device__ static inline void div(T &dst, const T &lhs, const U &rhs) { + bin_op(dst, lhs, rhs); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/vec/reductions.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/vec/reductions.cuh new file mode 100644 index 0000000..89212da --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/vec/reductions.cuh @@ -0,0 +1,207 @@ + + + + + + + + + + + + + + + + + + + + + + +template +__device__ static inline void reduce(typename SV::dtype &dst_accum, const SV &src, const typename SV::dtype &src_accum) { + if constexpr (GROUP_WARPS == 1) { + using T = SV::dtype; + int lane = laneid(); + T accum; + if(lane < SV::length) accum = src[lane]; + __syncwarp(); + for(int i = lane+kittens::WARP_THREADS; i < SV::length; i+=kittens::WARP_THREADS) { + accum = op::template op(accum, src[i]); + } + __syncwarp(); + + if constexpr (SV::length > 16) { + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 16)); + __syncwarp(); + } + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 8)); + __syncwarp(); + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 4)); + __syncwarp(); + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 2)); + __syncwarp(); + accum = op::template op(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 1)); + __syncwarp(); + if constexpr (!reset) accum = op::template op(accum, src_accum); + + dst_accum = packed_shfl_sync(kittens::MASK_ALL, accum, 0); + } + else { + ::kittens::group<1>::reduce(dst_accum, src, src_accum); + } +} + + + + + + + + + + +template +__device__ static inline void max(typename SV::dtype &max_val, const SV &src) { + reduce(max_val, src, max_val); +} +template +__device__ static inline typename SV::dtype max(const SV &src) { + typename SV::dtype max_val; + reduce(max_val, src, max_val); + return max_val; +} + + + + + + + + +template +__device__ static inline void min(typename SV::dtype &min_val, const SV &src) { + reduce(min_val, src, min_val); +} +template +__device__ static inline typename SV::dtype min(const SV &src) { + typename SV::dtype min_val; + reduce(min_val, src, min_val); + return min_val; +} + + + + + + + + +template +__device__ static inline void sum(typename SV::dtype &sum_val, const SV &src) { + reduce(sum_val, src, sum_val); +} +template +__device__ static inline typename SV::dtype sum(const SV &src) { + typename SV::dtype sum_val; + reduce(sum_val, src, sum_val); + return sum_val; +} + + + + + + + + +template +__device__ static inline void prod(typename SV::dtype &prod_val, const SV &src) { + reduce(prod_val, src, prod_val); +} +template +__device__ static inline typename SV::dtype prod(const SV &src) { + typename SV::dtype prod_val; + reduce(prod_val, src, prod_val); + return prod_val; +} + + + + + + + + + + + +template +__device__ static inline void max(typename SV::dtype &max_val, const SV &src, const typename SV::dtype &src_accum) { + reduce(max_val, src, src_accum); +} +template +__device__ static inline typename SV::dtype max(const SV &src, const typename SV::dtype &src_accum) { + typename SV::dtype max_val; + reduce(max_val, src, src_accum); + return max_val; +} + + + + + + + + + +template +__device__ static inline void min(typename SV::dtype &min_val, const SV &src, const typename SV::dtype &src_accum) { + reduce(min_val, src, src_accum); +} +template +__device__ static inline typename SV::dtype min(const SV &src, const typename SV::dtype &src_accum) { + typename SV::dtype min_val; + reduce(min_val, src, src_accum); + return min_val; +} + + + + + + + + + +template +__device__ static inline void sum(typename SV::dtype &sum_val, const SV &src, const typename SV::dtype &src_accum) { + reduce(sum_val, src, src_accum); +} +template +__device__ static inline typename SV::dtype sum(const SV &src, const typename SV::dtype &src_accum) { + typename SV::dtype sum_val; + reduce(sum_val, src, src_accum); + return sum_val; +} + + + + + + + + + +template +__device__ static inline void prod(typename SV::dtype &prod_val, const SV &src, const typename SV::dtype &src_accum) { + reduce(prod_val, src, src_accum); +} +template +__device__ static inline typename SV::dtype prod(const SV &src, const typename SV::dtype &src_accum) { + typename SV::dtype prod_val; + reduce(prod_val, src, src_accum); + return prod_val; +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/vec/vec.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/vec/vec.cuh new file mode 100644 index 0000000..5354d34 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/vec/vec.cuh @@ -0,0 +1,37 @@ + + + + + +#include "conversions.cuh" +#include "maps.cuh" +#include "reductions.cuh" + +template +__device__ static inline bool hasnan(const SV &src) { + KITTENS_CHECK_WARP + bool nan_detected = false; + #pragma unroll + for(int i = laneid(); i < SV::length; i+=GROUP_THREADS) { + if constexpr (std::is_same_v) { + if(isnan(src[i])) { + nan_detected = true; + } + } + else if constexpr (std::is_same_v) { + if(isnan(__bfloat162float(src[i]))) { + nan_detected = true; + } + } + else if constexpr (std::is_same_v) { + if(isnan(__half2float(src[i]))) { + nan_detected = true; + } + } + else { + static_assert(sizeof(typename SV::T) == 999, "Unsupported dtype"); + } + } + + return (__ballot_sync(0xffffffff, nan_detected) != 0); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/util/sync.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/util/sync.cuh new file mode 100644 index 0000000..e500aba --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/util/sync.cuh @@ -0,0 +1,224 @@ + + + + + +template __device__ static inline void load_async_wait(int bar_id) { + asm volatile("cp.async.wait_group %0;\n" : : "n"(N) : "memory"); + sync(bar_id); +} +template __device__ static inline void load_async_wait() { + KITTENS_CHECK_WARP + asm volatile("cp.async.wait_group %0;\n" : : "n"(N) : "memory"); + __syncwarp(); +} + +__device__ static inline void arrive(barrier bar) { + asm volatile("bar.arrive %0, %1;\n" :: "r"(bar.barrier_id), "n"(GROUP_WARPS*WARP_THREADS) : "memory"); +} +__device__ static inline void arrive_and_wait(barrier bar) { + asm volatile("bar.sync %0, %1;\n" :: "r"(bar.barrier_id), "n"(GROUP_WARPS*WARP_THREADS) : "memory"); +} + + + + + + + + + + + + + +__device__ static inline void init_semaphore(semaphore& bar, int thread_count, int transaction_count=0) { + if (laneid() == 0) { + void const* const ptr = &bar; + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + + asm volatile ( + "mbarrier.init.shared::cta.b64 [%0], %1;\n" + :: "r"(bar_ptr), "r"(thread_count+transaction_count) + ); + } +} + + + + + +__device__ static inline void invalidate_semaphore(semaphore& bar) { + if (laneid() == 0) { + void const* const ptr = &bar; + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + asm volatile ( + "mbarrier.inval.shared::cta.b64 [%0];\n" + :: "r"(bar_ptr) + ); + } +} + + + + + + + + +__device__ static inline void arrive(semaphore& sem) { + if(laneid() == 0) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(&sem)); + asm volatile ( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];\n" + : + : "r"(mbar_ptr) + : "memory" + ); + } +} +template __device__ static inline void arrive(barrier bar) { + asm volatile("bar.arrive %0, %1;\n" :: "r"(bar.barrier_id), "n"(num_warps*WARP_THREADS) : "memory"); +} + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + +__device__ static inline void arrive(semaphore& sem, uint32_t count) { + if(laneid() == 0) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(&sem)); + asm volatile ( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0], %1;\n" + : + : "r"(mbar_ptr), "r"(count) + : "memory" + ); + } +} +#endif + + + + + + + +__device__ static inline void wait(semaphore& sem, int kPhaseBit) { + void const* const ptr = &sem; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +#else + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.test_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "nanosleep.u32 5;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +#endif +} + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +__device__ static inline bool try_wait(semaphore &sem, int kPhaseBit) { + void const* const ptr = &sem; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + uint32_t success; + + asm volatile( + "{\n" + ".reg .pred P1; \n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2; \n" + "selp.b32 %0, 1, 0, P1; \n" + "}\n" + : "=r"(success) + : "r"(mbar_ptr), "r"(kPhaseBit) + : "memory" + ); + + return static_cast(success); +} +#endif + + + + + + + +__device__ static inline int test_wait(semaphore& sem, int kPhaseBit) { + void const* const ptr = &sem; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + int result; + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "mbarrier.test_wait.parity.shared::cta.b64 P1, [%1], %2;\n" + "selp.u32 %0,1,0,P1;" + "}\n" + : "=r"(result) + : "r"(mbar_ptr), "r"(kPhaseBit) + ); + return result; +} + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + +__device__ static inline bool elect_leader() { + if constexpr (GROUP_WARPS == 1) { + uint32_t elected = 0; + asm volatile( + "{.reg .pred P;\n" + " elect.sync _|P, %1;\n" + " selp.u32 %0, 1, 0, P;}\n" + : "+r"(elected) + : "r"(0xFFFFFFFF) + ); + return static_cast(elected); + } else { + return (warpid() == 0 && ::kittens::group<1>::elect_leader()); + } +} + +__device__ static inline bool elect_leader_from_active() { + if constexpr (GROUP_WARPS == 1) { + uint32_t elected = 0; + asm volatile( + "{.reg .pred P;\n" + " .reg .u32 mask;\n" + " activemask.b32 mask;\n" + " elect.sync _|P, mask;\n" + " selp.u32 %0, 1, 0, P;}\n" + : "+r"(elected) + ); + return static_cast(elected); + } else { + return (warpid() == 0 && ::kittens::group<1>::elect_leader_from_active()); + } +} + +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/util/tma.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/util/tma.cuh new file mode 100644 index 0000000..015b31a --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/util/tma.cuh @@ -0,0 +1,90 @@ + + + + + + + + + + + + + + + + + +__device__ static inline void expect_bytes(semaphore& bar, uint32_t bytes) { + if(laneid() == 0) { + ::kittens::tma::expect_bytes(bar, bytes); + } +} + + + + + +template +__device__ static inline void expect(semaphore& bar, const T& _1, const args&... _2) { + expect_bytes(bar, size_bytes); +} + + + + + + +__device__ static inline void store_commit_group() { + asm volatile("cp.async.bulk.commit_group;"); +} + + + + + +template +__device__ static inline void store_async_wait() { + asm volatile ( + "cp.async.bulk.wait_group %0;" + : + : "n"(N) + : "memory" + ); +} + + + + + +template +__device__ static inline void store_async_read_wait() { + asm volatile ( + "cp.async.bulk.wait_group.read %0;" + : + : "n"(N) + : "memory" + ); +} + + + +__device__ static inline void load_async(void *dst, void *src, uint32_t size_bytes, semaphore& bar) { + if(laneid() == 0) { + ::kittens::tma::load_async(dst, src, size_bytes, bar); + } +} +template +__device__ static inline void load_async(T &dst, T &src, uint32_t size_bytes, semaphore& bar) { + load_async(reinterpret_cast(&dst), reinterpret_cast(&src), size_bytes, bar); +} + +__device__ static inline void store_async(void *dst, void *src, uint32_t size_bytes) { + if(laneid() == 0) { + ::kittens::tma::store_async(dst, src, size_bytes); + } +} +template +__device__ static inline void store_async(T &dst, T &src, uint32_t size_bytes) { + store_async(reinterpret_cast(&dst), reinterpret_cast(&src), size_bytes); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/util/tma_cluster.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/util/tma_cluster.cuh new file mode 100644 index 0000000..7404629 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/util/tma_cluster.cuh @@ -0,0 +1,118 @@ + + + + + + + + + + + +__device__ static inline void wait(semaphore& bar, int kPhaseBit) { + void const* const ptr = &bar; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +} + +__device__ static inline bool try_wait(semaphore &bar, int kPhaseBit) { + void const* const ptr = &bar; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + uint32_t success; + + asm volatile( + "{\n" + ".reg .pred P1; \n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [%1], %2; \n" + "selp.b32 %0, 1, 0, P1; \n" + "}\n" + : "=r"(success) + : "r"(mbar_ptr), "r"(kPhaseBit) + : "memory" + ); + + return static_cast(success); +} + + + + + + + + + + + + + + + +__device__ static inline void expect_bytes(semaphore& bar, uint32_t bytes) { + if(laneid() == 0) { + ::kittens::tma::cluster::expect_bytes(bar, bytes); + } +} +__device__ static inline void expect_bytes(semaphore& bar, uint32_t bytes, int dst_cta) { + if(laneid() == 0) { + ::kittens::tma::cluster::expect_bytes(bar, bytes, dst_cta); + } +} + + + + + + + + + + + + + + + +template +__device__ static inline void expect(semaphore& bar, const T& _1, const args&... _2) { + expect_bytes(bar, size_bytes); +} + + + + + + + + + + +__device__ static inline void arrive(semaphore& bar, int dst_cta, uint32_t count=1) { + if(laneid() == 0) { + ::kittens::tma::cluster::arrive(bar, dst_cta, count); + } +} + + + +__device__ static inline void load_async(void *dst, void *src, uint32_t size_bytes, semaphore& bar, uint16_t cta_mask) { + if(laneid() == 0) { + ::kittens::tma::cluster::load_async(dst, src, size_bytes, bar, cta_mask); + } +} +template +__device__ static inline void load_async(T &dst, T &src, uint32_t size_bytes, semaphore& bar, uint16_t cta_mask) { + load_async(reinterpret_cast(&dst), reinterpret_cast(&src), size_bytes, bar, cta_mask); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/util/util.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/util/util.cuh new file mode 100644 index 0000000..295407b --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/util/util.cuh @@ -0,0 +1,94 @@ + + + + + +#include "sync.cuh" + + + + +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + +struct clc { + + + + + + + + +__device__ static inline void schedule(kittens::clc::handle &h, kittens::semaphore &sem) { + if (laneid() == 0) { + asm volatile("{clusterlaunchcontrol.try_cancel.async.shared::cta.mbarrier::complete_tx::bytes.multicast::cluster::all.b128 [%0], [%1];}" + :: "r"(static_cast(__cvta_generic_to_shared(&h.internal_value))), "r"(static_cast(__cvta_generic_to_shared(&sem))) + : "memory" + ); + } +} + + + + + +__device__ static inline kittens::clc::result query(kittens::clc::handle &h) { + kittens::clc::result r; + asm volatile( + "{\n" + ".reg .pred SUCCESS;\n" + ".reg .b128 CLC_HANDLE;\n" + "ld.shared.b128 CLC_HANDLE, [%4];\n" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 SUCCESS, CLC_HANDLE;\n" + "selp.u32 %0, 1, 0, SUCCESS;\n" + "@!SUCCESS bra.uni DONE;\n" + "clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%1, %2, %3, _}, CLC_HANDLE;\n" + "fence.proxy.async.shared::cta;\n" + "DONE:\n" + "}" + : "=r"(r.success), "=r"(r.x), "=r"(r.y), "=r"(r.z) + : "r"(static_cast(__cvta_generic_to_shared(&h.internal_value))) + : "memory" + ); + return r; +} + +}; + +#endif + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + + +struct pdl { + + + + + + + + +__device__ static inline void arrive() { + if (laneid() == 0) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +} + + + + +__device__ static inline void wait() { + asm volatile("griddepcontrol.wait;" ::: "memory"); +} + +}; + +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/ops.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/ops.cuh new file mode 100644 index 0000000..0c58623 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/ops.cuh @@ -0,0 +1,261 @@ + + + + + +#pragma once + +#include "thread/thread.cuh" +#include "group/group.cuh" + +namespace kittens { + + + + + +template +__device__ static inline T operator+(const T &lhs, const U &rhs) { + T dst; + warp::add(dst, lhs, rhs); + return dst; +} +template +__device__ static inline void operator+=(T &lhs, const U &rhs) { + warp::add(lhs, lhs, rhs); +} +template +__device__ static inline T operator-(const T &lhs, const U &rhs) { + T dst; + warp::sub(dst, lhs, rhs); + return dst; +} +template +__device__ static inline void operator-=(T &lhs, const U &rhs) { + warp::sub(lhs, lhs, rhs); +} +template +__device__ static inline T operator*(const T &lhs, const U &rhs) { + T dst; + warp::mul(dst, lhs, rhs); + return dst; +} +template +__device__ static inline void operator*=(T &lhs, const U &rhs) { + warp::mul(lhs, lhs, rhs); +} +template +__device__ static inline T operator/(const T &lhs, const U &rhs) { + T dst; + warp::div(dst, lhs, rhs); + return dst; +} +template +__device__ static inline void operator/=(T &lhs, const U &rhs) { + warp::div(lhs, lhs, rhs); +} +template +__device__ static inline T operator+(const T &src, const V &row_values) { + T dst; + warp::add_row(dst, src, row_values); + return dst; +} +template +__device__ static inline T operator+(const T &src, const V &row_values) { + T dst; + warp::add_row(dst, src, row_values); + return dst; +} +template +__device__ static inline void operator+=(T &lhs, const V &row_values) { + warp::add_row(lhs, lhs, row_values); +} +template +__device__ static inline void operator+=(T &lhs, const V &row_values) { + warp::add_row(lhs, lhs, row_values); +} +template +__device__ static inline T operator-(const T &src, const V &row_values) { + T dst; + warp::sub_row(dst, src, row_values); + return dst; +} +template +__device__ static inline T operator-(const T &src, const V &row_values) { + T dst; + warp::sub_row(dst, src, row_values); + return dst; +} +template +__device__ static inline void operator-=(T &lhs, const V &row_values) { + warp::sub_row(lhs, lhs, row_values); +} +template +__device__ static inline void operator-=(T &lhs, const V &row_values) { + warp::sub_row(lhs, lhs, row_values); +} +template +__device__ static inline T operator*(const T &src, const V &row_values) { + T dst; + warp::mul_row(dst, src, row_values); + return dst; +} +template +__device__ static inline T operator*(const T &src, const V &row_values) { + T dst; + warp::mul_row(dst, src, row_values); + return dst; +} +template +__device__ static inline void operator*=(T &lhs, const V &row_values) { + warp::mul_row(lhs, lhs, row_values); +} +template +__device__ static inline void operator*=(T &lhs, const V &row_values) { + warp::mul_row(lhs, lhs, row_values); +} +template +__device__ static inline T operator/(const T &src, const V &row_values) { + T dst; + warp::div_row(dst, src, row_values); + return dst; +} +template +__device__ static inline T operator/(const T &src, const V &row_values) { + T dst; + warp::div_row(dst, src, row_values); + return dst; +} +template +__device__ static inline void operator/=(T &lhs, const V &row_values) { + warp::div_row(lhs, lhs, row_values); +} +template +__device__ static inline void operator/=(T &lhs, const V &row_values) { + warp::div_row(lhs, lhs, row_values); +} +template +__device__ static inline T operator+(const T &src, const V &col_values) { + T dst; + warp::add_col(dst, src, col_values); + return dst; +} +template +__device__ static inline T operator+(const T &src, const V &col_values) { + T dst; + warp::add_col(dst, src, col_values); + return dst; +} +template +__device__ static inline void operator+=(T &lhs, const V &col_values) { + warp::add_col(lhs, lhs, col_values); +} +template +__device__ static inline void operator+=(T &lhs, const V &col_values) { + warp::add_col(lhs, lhs, col_values); +} +template +__device__ static inline T operator-(const T &src, const V &col_values) { + T dst; + warp::sub_col(dst, src, col_values); + return dst; +} +template +__device__ static inline T operator-(const T &src, const V &col_values) { + T dst; + warp::sub_col(dst, src, col_values); + return dst; +} +template +__device__ static inline void operator-=(T &lhs, const V &col_values) { + warp::sub_col(lhs, lhs, col_values); +} +template +__device__ static inline void operator-=(T &lhs, const V &col_values) { + warp::sub_col(lhs, lhs, col_values); +} +template +__device__ static inline T operator*(const T &src, const V &col_values) { + T dst; + warp::mul_col(dst, src, col_values); + return dst; +} +template +__device__ static inline T operator*(const T &src, const V &col_values) { + T dst; + warp::mul_col(dst, src, col_values); + return dst; +} +template +__device__ static inline void operator*=(T &lhs, const V &col_values) { + warp::mul_col(lhs, lhs, col_values); +} +template +__device__ static inline void operator*=(T &lhs, const V &col_values) { + warp::mul_col(lhs, lhs, col_values); +} +template +__device__ static inline T operator/(const T &src, const V &col_values) { + T dst; + warp::div_col(dst, src, col_values); + return dst; +} +template +__device__ static inline T operator/(const T &src, const V &col_values) { + T dst; + warp::div_col(dst, src, col_values); + return dst; +} +template +__device__ static inline void operator/=(T &lhs, const V &col_values) { + warp::div_col(lhs, lhs, col_values); +} +template +__device__ static inline void operator/=(T &lhs, const V &col_values) { + warp::div_col(lhs, lhs, col_values); +} + + + +template +__device__ static inline T operator+(const T &lhs, const U &rhs) { + T dst; + warp::add(dst, lhs, rhs); + return dst; +} +template +__device__ static inline void operator+=(T &lhs, const U &rhs) { + warp::add(lhs, lhs, rhs); +} +template +__device__ static inline T operator-(const T &lhs, const U &rhs) { + T dst; + warp::sub(dst, lhs, rhs); + return dst; +} +template +__device__ static inline void operator-=(T &lhs, const U &rhs) { + warp::sub(lhs, lhs, rhs); +} +template +__device__ static inline T operator*(const T &lhs, const U &rhs) { + T dst; + warp::mul(dst, lhs, rhs); + return dst; +} +template +__device__ static inline void operator*=(T &lhs, const U &rhs) { + warp::mul(lhs, lhs, rhs); +} +template +__device__ static inline T operator/(const T &lhs, const U &rhs) { + T dst; + warp::div(dst, lhs, rhs); + return dst; +} +template +__device__ static inline void operator/=(T &lhs, const U &rhs) { + warp::div(lhs, lhs, rhs); +} + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/memory.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/memory.cuh new file mode 100644 index 0000000..3b0cb5f --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/memory.cuh @@ -0,0 +1,9 @@ + + + + + +#pragma once + +#include "tile/tile.cuh" +#include "vec/vec.cuh" diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/tile/shared_to_tensor.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/tile/shared_to_tensor.cuh new file mode 100644 index 0000000..72eed49 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/tile/shared_to_tensor.cuh @@ -0,0 +1,39 @@ + + + + + +#pragma once + +#include "../../mma/tcgen05.cuh" + +template +__device__ inline static void load_mxnv_scale_async(TT &dst, const ST &src) { + static_assert(ncta == 1 || ncta == 2, "ncta must be 1 or 2"); + static_assert(std::is_same_v || std::is_same_v, "Scale TT must be fp8e8m0 or fp8e4m3"); + static_assert(std::is_same_v, "Scale TT and ST must have the same type"); + static_assert(TT::rows == 128 && TT::cols == 16, "Tensor memory scale tile must always be 128x16"); + static_assert(ST::rows == 32 && ST::cols == 16, "Shared memory scale tile must always be 32x16"); + static_assert(!ST::swizzle, "Shared memory scale tile must not be TMA-swizzled"); + + uint64_t st_desc = kittens::detail::matrix_descriptor_raw(&src.data[0], 128, 128, 0); + + if constexpr (ncta == 1) { + asm volatile("{tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;}" :: "r"(dst.addr), "l"(st_desc)); + } else { + asm volatile("{tcgen05.cp.cta_group::2.32x128b.warpx4 [%0], %1;}" :: "r"(dst.addr), "l"(st_desc)); + } +} +template +__device__ inline static void load_mxnv_scale_async(TT &dst, const ST &src, kittens::semaphore &sem) { + load_mxnv_scale_async(dst, src); + kittens::detail::tcgen05::commit(sem); +} +template +__device__ inline static void load_mxnv_scale_async2(TT &dst, const ST &src) { + load_mxnv_scale_async<2>(dst, src); +} +template +__device__ inline static void load_mxnv_scale_async2(TT &dst, const ST &src, kittens::semaphore &sem) { + load_mxnv_scale_async<2>(dst, src, sem); +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/tile/tile.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/tile/tile.cuh new file mode 100644 index 0000000..9b41b67 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/tile/tile.cuh @@ -0,0 +1,13 @@ + + + + + +#pragma once + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +#include "tma.cuh" +#endif +#if defined(KITTENS_SM10X) +#include "shared_to_tensor.cuh" +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/tile/tma.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/tile/tma.cuh new file mode 100644 index 0000000..0a063ad --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/tile/tma.cuh @@ -0,0 +1,880 @@ + + + + + +#pragma once + +#include "../../../../common/common.cuh" +#include "../../../../types/types.cuh" +#include "../../util/util.cuh" + +namespace kittens { +namespace tma { + +namespace detail { +template __device__ inline int4 tma_coords(const coord &unit_coord) { + static_assert(ST::swizzle, "tma_coords only should be called for swizzled tiles"); + constexpr int swizzle_elements = ST::swizzle_bytes / sizeof(typename ST::dtype); + if constexpr (axis == 2) return {unit_coord.r, unit_coord.c / swizzle_elements, unit_coord.d, unit_coord.b}; + else if constexpr (axis == 1) return {unit_coord.d, unit_coord.c / swizzle_elements, unit_coord.r, unit_coord.b}; + else if constexpr (axis == 0) return {unit_coord.b, unit_coord.c / swizzle_elements, unit_coord.r, unit_coord.d}; +} +} + + + + + + + + + + + +template> +__device__ static inline void prefetch(ST &dst, const GL &src, const COORD &idx) { + uint64_t tma_ptr = reinterpret_cast(src.template get_tma()); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.prefetch.tensor.5d.L2.global.tile" + " [%0, {%1, %2, %3, %4, %5}];" + : + : "l"(tma_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.prefetch.tensor.5d.L2.global.tile.L2::cache_hint" + " [%0, {%1, %2, %3, %4, %5}], %6;" + : + : "l"(tma_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.prefetch.tensor.4d.L2.global.tile" + " [%0, {%1, %2, %3, %4}];" + : + : "l"(tma_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.prefetch.tensor.4d.L2.global.tile.L2::cache_hint" + " [%0, {%1, %2, %3, %4}], %5;" + : + : "l"(tma_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } + } +} +template> +__device__ static inline void prefetch(ST &dst, const GL &src, const COORD &idx) { + prefetch(dst, src, idx); +} + + + + + + + + + + + + + +template> +__device__ static inline void store_async(const GL &dst, const ST &src, const COORD &idx) { + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.5d.global.shared::cta.tile.bulk_group" + " [%0, {%2, %3, %4, %5, %6}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.tensor.5d.global.shared::cta.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5, %6}], [%1], %7;" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } + } + store_commit_group(); +} +template> +__device__ static inline void store_async(const GL &dst, const ST &src, const COORD &idx) { + store_async(dst, src, idx); +} +template> +__device__ static inline void store_async(const PGL &dst, const ST &src, const COORD &idx) { + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.5d.global.shared::cta.tile.bulk_group" + " [%0, {%2, %3, %4, %5, %6}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.tensor.5d.global.shared::cta.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5, %6}], [%1], %7;" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } + } + store_commit_group(); +} +template> +__device__ static inline void store_async(const PGL &dst, const ST &src, const COORD &idx) { + store_async(dst, src, idx); +} + + + + + + + + + + + + + +template> +__device__ static inline void store_add_async(const GL &dst, const ST &src, const COORD &idx) { + + static_assert(!(std::is_same_v || + std::is_same_v), + "TMA does not support async add reductions for fp8 types."); + + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.add.tile.bulk_group" + " [%0, {%2, %3, %4, %5, %6}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.add.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5, %6}], [%1], %7;" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.add.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.add.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } + } + store_commit_group(); +} +template> +__device__ static inline void store_add_async(const GL &dst, const ST &src, const COORD &idx) { + store_add_async(dst, src, idx); +} +template> +__device__ static inline void store_add_async(const PGL &dst, const ST &src, const COORD &idx) { + + static_assert(!(std::is_same_v || + std::is_same_v), + "TMA does not support async add reductions for fp8 types."); + + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.add.tile.bulk_group" + " [%0, {%2, %3, %4, %5, %6}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.add.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5, %6}], [%1], %7;" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.add.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.add.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } + } + store_commit_group(); +} +template> +__device__ static inline void store_add_async(const PGL &dst, const ST &src, const COORD &idx) { + store_add_async(dst, src, idx); +} + + + + + + + + + + + +template> +__device__ static inline void store_min_async(const GL &dst, const ST &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + + static_assert(!(std::is_same_v || + std::is_same_v), + "TMA does not support async min reductions for fp8 types."); + + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.min.tile.bulk_group" + " [%0, {%2, %3, %4, %5, %6}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.min.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5, %6}], [%1], %7;" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.min.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.min.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } + } + store_commit_group(); +} +template> +__device__ static inline void store_min_async(const GL &dst, const ST &src, const COORD &idx) { + store_min_async(dst, src, idx); +} +template> +__device__ static inline void store_min_async(const PGL &dst, const ST &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + + static_assert(!(std::is_same_v || + std::is_same_v), + "TMA does not support async min reductions for fp8 types."); + + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.min.tile.bulk_group" + " [%0, {%2, %3, %4, %5, %6}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.min.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5, %6}], [%1], %7;" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.min.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.min.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } + } + store_commit_group(); +} +template> +__device__ static inline void store_min_async(const PGL &dst, const ST &src, const COORD &idx) { + store_min_async(dst, src, idx); +} + + + + + + + + + + + +template> +__device__ static inline void store_max_async(const GL &dst, const ST &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + + static_assert(!(std::is_same_v || + std::is_same_v), + "TMA does not support async max reductions for fp8 types."); + + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.max.tile.bulk_group" + " [%0, {%2, %3, %4, %5, %6}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.max.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5, %6}], [%1], %7;" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.max.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.max.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } + } + store_commit_group(); +} +template> +__device__ static inline void store_max_async(const GL &dst, const ST &src, const COORD &idx) { + store_max_async(dst, src, idx); +} +template> +__device__ static inline void store_max_async(const PGL &dst, const ST &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + + static_assert(!(std::is_same_v || + std::is_same_v), + "TMA does not support async max reductions for fp8 types."); + + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.max.tile.bulk_group" + " [%0, {%2, %3, %4, %5, %6}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.5d.global.shared::cta.max.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5, %6}], [%1], %7;" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.max.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.max.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } + } + store_commit_group(); +} +template> +__device__ static inline void store_max_async(const PGL &dst, const ST &src, const COORD &idx) { + store_max_async(dst, src, idx); +} + + + + + + + + + + + + +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar) { + uint64_t tma_ptr = reinterpret_cast(src.template get_tma()); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(&bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst)); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + + if constexpr (policy == cache_policy::NORMAL) { + asm volatile( + "cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5, %6, %7}], [%2];" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w) + : "memory" + ); + } + else { + asm volatile( + "cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + + if constexpr (policy == cache_policy::NORMAL) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5, %6}], [%2];" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b) + : "memory" + ); + } + else { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%3, %4, %5, %6}], [%2], %7;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } + } +} +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar) { + load_async(dst, src, idx, bar); +} + +namespace cluster { + + + + + + + + + + + + + +#ifdef KITTENS_SM10X +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) +#else +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask) +#endif +{ + uint64_t tma_ptr = reinterpret_cast(src.template get_tma()); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(&bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst)); + coord unit_coord = idx.template unit_coord(); + + if constexpr (ST::swizzle) { + int4 tma_coords = detail::tma_coords(unit_coord); + +#ifdef KITTENS_SM10X + if(dst_mbar_cta != -1) { + uint32_t neighbor_mbar_ptr; + asm volatile ( + "mapa.shared::cluster.u32 %0, %1, %2;\n" + : "=r"(neighbor_mbar_ptr) + : "r"(mbar_ptr), "r"(dst_mbar_cta) + ); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.cta_group::2.multicast::cluster" + " [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(neighbor_mbar_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "h"(cluster_mask) + : "memory" + ); + } else { + asm volatile ( + "cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.cta_group::2.multicast::cluster.L2::cache_hint" + " [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8, %9;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(neighbor_mbar_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "h"(cluster_mask), "l"(make_cache_policy()) + : "memory" + ); + } + } else +#endif + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "h"(cluster_mask) + : "memory" + ); + } else { + asm volatile ( + "cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster.L2::cache_hint" + " [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8, %9;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "h"(cluster_mask), "l"(make_cache_policy()) + : "memory" + ); + } + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + +#ifdef KITTENS_SM10X + if(dst_mbar_cta != -1) { + uint32_t neighbor_mbar_ptr; + asm volatile ( + "mapa.shared::cluster.u32 %0, %1, %2;\n" + : "=r"(neighbor_mbar_ptr) + : "r"(mbar_ptr), "r"(dst_mbar_cta) + ); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.cta_group::2.multicast::cluster" + " [%0], [%1, {%3, %4, %5, %6}], [%2], %7;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(neighbor_mbar_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "h"(cluster_mask) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.cta_group::2.multicast::cluster.L2::cache_hint" + " [%0], [%1, {%3, %4, %5, %6}], [%2], %7, %8;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(neighbor_mbar_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "h"(cluster_mask), "l"(make_cache_policy()) + : "memory" + ); + } + } else +#endif + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%3, %4, %5, %6}], [%2], %7;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "h"(cluster_mask) + : "memory" + ); + } else { + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster.L2::cache_hint" + " [%0], [%1, {%3, %4, %5, %6}], [%2], %7, %8;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "r"(unit_coord.c), "r"(unit_coord.r), "r"(unit_coord.d), "r"(unit_coord.b), "h"(cluster_mask), "l"(make_cache_policy()) + : "memory" + ); + } + } +} +#ifdef KITTENS_SM10X +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) { + load_async(dst, src, idx, bar, cluster_mask, dst_mbar_cta); +} +#else +template> +__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask) { + load_async(dst, src, idx, bar, cluster_mask); +} +#endif + +} +} + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/vec/tma.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/vec/tma.cuh new file mode 100644 index 0000000..f288ef4 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/vec/tma.cuh @@ -0,0 +1,418 @@ + + + + + +#pragma once + +#include "../../../../common/common.cuh" +#include "../../../../types/types.cuh" +#include "../../util/util.cuh" + + +#define __KITTENS_TMA_DEFINE_DEFAULT_LOAD_CACHE_VEC__(function_name) \ +template> \ +__device__ static inline void function_name(SV &dst, const GL &src, const COORD &idx) { \ + function_name(dst, src, idx); \ +} +#define __KITTENS_TMA_DEFINE_PGL_DEFAULT_LOAD_CACHE_VEC__(function_name) \ +template> \ +__device__ static inline void function_name(SV &dst, const PGL &src, const COORD &idx) { \ + function_name(dst, src, idx); \ +} +#define __KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(function_name) \ +template> \ +__device__ static inline void function_name(const GL &dst, const SV &src, const COORD &idx) { \ + function_name(dst, src, idx); \ +} +#define __KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(function_name) \ +template> \ +__device__ static inline void function_name(const PGL &dst, const SV &src, const COORD &idx) { \ + function_name(dst, src, idx); \ +} +#define __KITTENS_TMA_DEFINE_SEMAPHORE_CACHE_VEC__(function_name) \ +template> \ +__device__ static inline void function_name(SV &dst, const GL &src, const COORD &idx, semaphore& bar) { \ + function_name(dst, src, idx, bar); \ +} +#define __KITTENS_TMA_DEFINE_PGL_SEMAPHORE_CACHE_VEC__(function_name) \ +template> \ +__device__ static inline void function_name(SV &dst, const PGL &src, const COORD &idx, semaphore& bar) { \ + function_name(dst, src, idx, bar); \ +} +#define __KITTENS_TMA_DEFINE_CLUSTER_SEMAPHORE_CACHE_VEC__(function_name) \ +template> \ +__device__ static inline void function_name(SV &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) { \ + function_name(dst, src, idx, bar, cluster_mask, dst_mbar_cta); \ +} +#define __KITTENS_TMA_DEFINE_PGL_CLUSTER_SEMAPHORE_CACHE_VEC__(function_name) \ +template> \ +__device__ static inline void function_name(SV &dst, const PGL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) { \ + function_name(dst, src, idx, bar, cluster_mask, dst_mbar_cta); \ +} + + +namespace kittens { + +namespace detail { +namespace tma { + +template __device__ static inline void vec_prefetch_tma_internal(uint64_t tma_ptr, coord<> tma_coord) { + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.prefetch.tensor.4d.L2.global.tile" + " [%0, {%1, %2, %3, %4}];" + : + : "l"(tma_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.prefetch.tensor.4d.L2.global.tile.L2::cache_hint" + " [%0, {%1, %2, %3, %4}], %5;" + : + : "l"(tma_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } +} + +template __device__ static inline void vec_store_async_tma_internal(uint64_t tma_ptr, uint32_t src_i_ptr, coord<> tma_coord) { + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } +} + +template __device__ static inline void vec_store_add_async_tma_internal(uint64_t tma_ptr, uint32_t src_i_ptr, coord<> tma_coord) { + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.add.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.add.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } +} + +template __device__ static inline void vec_store_min_async_tma_internal(uint64_t tma_ptr, uint32_t src_i_ptr, coord<> tma_coord) { + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.min.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.min.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } +} + +template __device__ static inline void vec_store_max_async_tma_internal(uint64_t tma_ptr, uint32_t src_i_ptr, coord<> tma_coord) { + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.max.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.reduce.async.bulk.tensor.4d.global.shared::cta.max.tile.bulk_group.L2::cache_hint" + " [%0, {%2, %3, %4, %5}], [%1], %6;" + : + : "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } +} + +template __device__ static inline void vec_load_async_tma_internal(uint64_t tma_ptr, uint32_t dst_i_ptr, uint32_t mbar_ptr, coord<> tma_coord) { + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5, %6}], [%2];" + : + : "r"(dst_i_ptr), "l"(tma_ptr), "r"(mbar_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%3, %4, %5, %6}], [%2], %7;" + : + : "r"(dst_i_ptr), "l"(tma_ptr), "r"(mbar_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "l"(make_cache_policy()) + : "memory" + ); + } +} + +namespace cluster { +template __device__ static inline void vec_load_async_tma_internal(uint64_t tma_ptr, uint32_t dst_i_ptr, uint32_t mbar_ptr, coord<> tma_coord, uint16_t cluster_mask, int dst_mbar_cta=-1) { +#ifdef KITTENS_SM10X + if(dst_mbar_cta != -1) { + uint32_t neighbor_mbar_ptr; + asm volatile ( + "mapa.shared::cluster.u32 %0, %1, %2;\n" + : "=r"(neighbor_mbar_ptr) + : "r"(mbar_ptr), "r"(dst_mbar_cta) + ); + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.cta_group::2.multicast::cluster" + " [%0], [%1, {%3, %4, %5, %6}], [%2], %7;" + : + : "r"(dst_i_ptr), "l"(tma_ptr), "r"(neighbor_mbar_ptr), + "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "h"(cluster_mask) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.cta_group::2.multicast::cluster.L2::cache_hint" + " [%0], [%1, {%3, %4, %5, %6}], [%2], %7, %8;" + : + : "r"(dst_i_ptr), "l"(tma_ptr), "r"(neighbor_mbar_ptr), + "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "h"(cluster_mask), "l"(make_cache_policy()) + : "memory" + ); + } + } else +#endif + if constexpr (policy == cache_policy::NORMAL) { + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%3, %4, %5, %6}], [%2], %7;" + : + : "r"(dst_i_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "h"(cluster_mask) + : "memory" + ); + } + else { + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster.L2::cache_hint" + " [%0], [%1, {%3, %4, %5, %6}], [%2], %7, %8;" + : + : "r"(dst_i_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "h"(cluster_mask), "l"(make_cache_policy()) + : "memory" + ); + } +} +} + +} +} + +namespace tma { + +template> +__device__ static inline void prefetch(SV &dst, const GL &src, const COORD &idx) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(src.template get_tma()); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + ::kittens::detail::tma::vec_prefetch_tma_internal(tma_ptr, tma_coord); + } +} +__KITTENS_TMA_DEFINE_DEFAULT_LOAD_CACHE_VEC__(prefetch) + +template> +__device__ static inline void store_async(const GL &dst, const SV &src, const COORD &idx) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + ::kittens::tma::store_commit_group(); +} +__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_async) + +template> +__device__ static inline void store_async(const PGL &dst, const SV &src, const COORD &idx) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + ::kittens::tma::store_commit_group(); +} +__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_async) + +template> +__device__ static inline void store_add_async(const GL &dst, const SV &src, const COORD &idx) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_add_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + ::kittens::tma::store_commit_group(); +} +__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_add_async) + +template> +__device__ static inline void store_add_async(const PGL &dst, const SV &src, const COORD &idx) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_add_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + ::kittens::tma::store_commit_group(); +} +__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_add_async) + +template> +__device__ static inline void store_min_async(const GL &dst, const SV &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_min_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + ::kittens::tma::store_commit_group(); +} +__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_min_async) + +template> +__device__ static inline void store_min_async(const PGL &dst, const SV &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_min_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + ::kittens::tma::store_commit_group(); +} +__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_min_async) + +template> +__device__ static inline void store_max_async(const GL &dst, const SV &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_max_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + ::kittens::tma::store_commit_group(); +} +__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_max_async) + +template> +__device__ static inline void store_max_async(const PGL &dst, const SV &src, const COORD &idx) { + static_assert(!std::is_same_v, "TMA does not support async min/max reductions for fp32 types."); + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(dst.template get_tma()); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(&src)); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_store_max_async_tma_internal(tma_ptr, src_i_ptr, tma_coord); + } + ::kittens::tma::store_commit_group(); +} +__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_max_async) + +template> +__device__ static inline void load_async(SV &dst, const GL &src, const COORD &idx, semaphore& bar) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(src.template get_tma()); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(&bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst)); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t dst_i_ptr = dst_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::vec_load_async_tma_internal(tma_ptr, dst_i_ptr, mbar_ptr, tma_coord); + } +} +__KITTENS_TMA_DEFINE_SEMAPHORE_CACHE_VEC__(load_async) + +namespace cluster { +template> +__device__ static inline void load_async(SV &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) { + coord<> unit_coord = idx.template unit_coord<-1, 3>(); + uint64_t tma_ptr = reinterpret_cast(src.template get_tma()); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(&bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(&dst)); + for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2; i++) { + coord<> tma_coord = unit_coord; + tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1; + uint32_t dst_i_ptr = dst_ptr + i*::kittens::detail::tma::sv_tma_dim1*sizeof(typename SV::dtype); + ::kittens::detail::tma::cluster::vec_load_async_tma_internal(tma_ptr, dst_i_ptr, mbar_ptr, tma_coord, cluster_mask, dst_mbar_cta); + } +} +__KITTENS_TMA_DEFINE_CLUSTER_SEMAPHORE_CACHE_VEC__(load_async) +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/vec/vec.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/vec/vec.cuh new file mode 100644 index 0000000..8730f36 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/vec/vec.cuh @@ -0,0 +1,10 @@ + + + + + +#pragma once + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +#include "tma.cuh" +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/mma/mma.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/mma/mma.cuh new file mode 100644 index 0000000..9bf9e9c --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/mma/mma.cuh @@ -0,0 +1,10 @@ + + + + + +#pragma once + +#if defined(KITTENS_SM10X) +#include "tcgen05.cuh" +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/mma/tcgen05.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/mma/tcgen05.cuh new file mode 100644 index 0000000..6fc4f05 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/mma/tcgen05.cuh @@ -0,0 +1,985 @@ + + + + + +#pragma once + +#include "../../../common/common.cuh" +#include "../../../types/types.cuh" +#include "../util/sync.cuh" + +namespace kittens { +namespace detail { +namespace tcgen05 { + + +template +__device__ static inline constexpr uint32_t instruction_descriptor() { + uint32_t desc = 0; + if constexpr (std::is_same_v || std::is_same_v) { + + static_assert(std::is_same_v || std::is_same_v); + desc |= 0b00 << 0; + desc |= 0b0 << 2; + desc |= 0b0 << 3; + if constexpr (std::is_same_v) { + desc |= 0b01 << 4; + } + else { + desc |= 0b00 << 4; + } + desc |= 0b0 << 6; + if constexpr (std::is_same_v) { + desc |= 0b000 << 7; + desc |= 0b000 << 10; + } else if constexpr (std::is_same_v) { + desc |= 0b001 << 7; + desc |= 0b001 << 10; + } + if constexpr (neg) { + desc |= 0b1 << 13; + } + else { + desc |= 0b0 << 13; + } + desc |= 0b0 << 14; + if constexpr (trans_a) { + desc |= 0b1 << 15; + } + else { + desc |= 0b0 << 15; + } + if constexpr (trans_b) { + desc |= 0b1 << 16; + } + else { + desc |= 0b0 << 16; + } + desc |= (N >> 3) << 17; + desc |= 0b0 << 23; + desc |= (M >> 4) << 24; + desc |= 0b0 << 29; + desc |= 0b00 << 30; + } else if constexpr (std::is_same_v || std::is_same_v) { + static_assert(std::is_same_v, "Integer MMA must accumulate to int32."); + desc |= 0b00 << 0; + desc |= 0b0 << 2; + desc |= 0b1 << 3; + desc |= 0b10 << 4; + desc |= 0b0 << 6; + if constexpr (std::is_same_v) { + desc |= 0b000 << 7; + desc |= 0b000 << 10; + } else if constexpr (std::is_same_v) { + desc |= 0b001 << 7; + desc |= 0b001 << 10; + } else { + static_assert(sizeof(AB) == 999, "Invalid AB type."); + } + static_assert(!neg, "Negation is not supported for integer matrix multiplies."); + desc |= 0b0 << 13; + desc |= 0b0 << 14; + if constexpr (trans_a) { + desc |= 0b1 << 15; + } + else { + desc |= 0b0 << 15; + } + if constexpr (trans_b) { + desc |= 0b1 << 16; + } + else { + desc |= 0b0 << 16; + } + desc |= (N >> 3) << 17; + desc |= 0b0 << 23; + desc |= (M >> 4) << 24; + desc |= 0b0 << 29; + desc |= 0b00 << 30; + } else if constexpr (std::is_same_v || std::is_same_v || std::is_same_v) { + static_assert(std::is_same_v || std::is_same_v); + desc |= 0b00 << 0; + desc |= 0b0 << 2; + desc |= 0b0 << 3; + if constexpr (std::is_same_v) { + desc |= 0b01 << 4; + } + else { + desc |= 0b00 << 4; + } + desc |= 0b0 << 6; + if constexpr (std::is_same_v) { + desc |= 0b000 << 7; + desc |= 0b000 << 10; + } else if constexpr (std::is_same_v) { + desc |= 0b001 << 7; + desc |= 0b001 << 10; + } else if constexpr (std::is_same_v) { + desc |= 0b101 << 7; + desc |= 0b101 << 10; + } + if constexpr (neg) { + desc |= 0b1 << 13; + } + else { + desc |= 0b0 << 13; + } + desc |= 0b0 << 14; + if constexpr (trans_a) { + desc |= 0b1 << 15; + } + else { + desc |= 0b0 << 15; + } + if constexpr (trans_b) { + desc |= 0b1 << 16; + } + else { + desc |= 0b0 << 16; + } + desc |= (N >> 3) << 17; + desc |= 0b0 << 23; + desc |= (M >> 4) << 24; + desc |= 0b0 << 29; + desc |= 0b00 << 30; + } else { + static_assert(sizeof(AB) == 999, "Invalid AB type size; not implemented yet."); + } + return desc; +}; + +template +__device__ static inline constexpr uint32_t instruction_descriptor() { + + static_assert(std::is_same_v || std::is_same_v, "AB must be fp8e4m3 for f4e2m1"); + static_assert(std::is_same_v || std::is_same_v, "SAB must be either fp8e4m3 or fp8e8m0"); +#if !defined(KITTENS_SM103) + static_assert(mma_k_bytes == 32, "32B block-scaled MMA is only supported on non-SM103."); +#endif + static_assert(mma_k_bytes == 32 || (std::is_same_v && mma_k_bytes == 48), "48B (K96) block-scaled MMA is only supported for packed NVFP4 (E4M3 or E8M0 scales)."); + constexpr int scale_type = std::is_same_v ? 0 : std::is_same_v ? 1 : -1; + + uint32_t desc = 0; + desc |= 0b00 << 0; + desc |= 0b0 << 2; + desc |= 0b0 << 3; + desc |= scale_factor_id << 4; + desc |= 0b0 << 6; + if constexpr (std::is_same_v) { + desc |= (0b000 << 7); + desc |= (0b000 << 10); + } else if constexpr (std::is_same_v) { + desc |= 0b001 << 7; + desc |= 0b01 << 10; + desc |= 0b0 << 12; + } else { + static_assert(sizeof(AB) == 999, "Invalid AB type."); + } + if constexpr (neg) { + desc |= 0b1 << 13; + } + else { + desc |= 0b0 << 13; + } + desc |= 0b0 << 14; + desc |= 0b0 << 15; + desc |= 0b0 << 16; + desc |= (N >> 3) << 17; + desc |= scale_type << 23; + desc |= 0b000 << 24; + desc |= (M >> 7) << 27; + desc |= scale_factor_id << 29; + desc |= (mma_k_bytes == 48 ? 1u : 0u) << 31; + + return desc; +} + +template +__device__ static inline void tt_st(uint32_t d_tt_addr, uint32_t a_tt_addr, uint64_t b_desc, uint32_t idesc) { + if constexpr (std::is_same_v || std::is_same_v) { + if constexpr (ncta == 1) { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::1.kind::i8 [%0], [%1], %2, %3, p;}\n" + :: "r"(d_tt_addr), "r"(a_tt_addr), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + else { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::2.kind::i8 [%0], [%1], %2, %3, p;}\n" + :: "r"(d_tt_addr), "r"(a_tt_addr), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + } else if constexpr (std::is_same_v || std::is_same_v || std::is_same_v) { + if constexpr (ncta == 1) { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [%1], %2, %3, p;}\n" + :: "r"(d_tt_addr), "r"(a_tt_addr), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + else { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], [%1], %2, %3, p;}\n" + :: "r"(d_tt_addr), "r"(a_tt_addr), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + } else { + if constexpr (ncta == 1) { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::1.kind::f16 [%0], [%1], %2, %3, p;}\n" + :: "r"(d_tt_addr), "r"(a_tt_addr), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + else { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::2.kind::f16 [%0], [%1], %2, %3, p;}\n" + :: "r"(d_tt_addr), "r"(a_tt_addr), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + } +} + +template +__device__ static inline void st_st(uint32_t d_tt_addr, uint64_t a_desc, uint64_t b_desc, uint32_t idesc) { + if constexpr (std::is_same_v || std::is_same_v) { + if constexpr (ncta == 1) { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::1.kind::i8 [%0], %1, %2, %3, p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + else { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::2.kind::i8 [%0], %1, %2, %3, p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + } else if constexpr (std::is_same_v || std::is_same_v || std::is_same_v) { + if constexpr (ncta == 1) { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + else { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], %1, %2, %3, p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + } else { + if constexpr (ncta == 1) { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + else { + asm volatile( + "{.reg .pred p;\n" \ + "setp.eq.u32 p, 1, %4;\n" \ + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "n"(acc) + ); + } + } +} + +template +__device__ static inline void st_st(uint32_t d_tt_addr, uint64_t a_desc, uint64_t b_desc, uint32_t sa_tt_addr, uint32_t sb_tt_addr, uint32_t idesc) { + static_assert(std::is_same_v || std::is_same_v, "AB must be fp8e4m3 for f4e2m1"); + if constexpr (ncta == 1) { + if constexpr (std::is_same_v) { + asm volatile( + "{.reg .pred p;\n\t" \ + "setp.eq.u32 p, 1, %6;\n\t" \ + "tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [%0], %1, %2, %3, [%4], [%5], p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "r"(sa_tt_addr), "r"(sb_tt_addr), "n"(acc) + ); + } else if constexpr (std::is_same_v) { + if constexpr (block_size == 32) { + asm volatile( + "{.reg .pred p;\n\t" \ + "setp.eq.u32 p, 1, %6;\n\t" \ + "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::2X [%0], %1, %2, %3, [%4], [%5], p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "r"(sa_tt_addr), "r"(sb_tt_addr), "n"(acc) + ); + } + else { + asm volatile( + "{.reg .pred p;\n\t" \ + "setp.eq.u32 p, 1, %6;\n\t" \ + "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::4X [%0], %1, %2, %3, [%4], [%5], p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "r"(sa_tt_addr), "r"(sb_tt_addr), "n"(acc) + ); + } + } else { + static_assert(sizeof(AB) == 999, "Invalid AB type."); + } + } + else { + if constexpr (std::is_same_v) { + asm volatile( + "{.reg .pred p;\n\t" \ + "setp.eq.u32 p, 1, %6;\n\t" \ + "tcgen05.mma.cta_group::2.kind::mxf8f6f4.block_scale.scale_vec::1X [%0], %1, %2, %3, [%4], [%5], p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "r"(sa_tt_addr), "r"(sb_tt_addr), "n"(acc) + ); + } else if constexpr (std::is_same_v) { + if constexpr (block_size == 32) { + asm volatile( + "{.reg .pred p;\n\t" \ + "setp.eq.u32 p, 1, %6;\n\t" \ + "tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::2X [%0], %1, %2, %3, [%4], [%5], p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "r"(sa_tt_addr), "r"(sb_tt_addr), "n"(acc) + ); + } + else { + asm volatile( + "{.reg .pred p;\n\t" \ + "setp.eq.u32 p, 1, %6;\n\t" \ + "tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X [%0], %1, %2, %3, [%4], [%5], p;}\n" + :: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "r"(sa_tt_addr), "r"(sb_tt_addr), "n"(acc) + ); + } + } else { + static_assert(sizeof(AB) == 999, "Invalid AB type."); + } + } +} + +template +__device__ static inline void commit(kittens::semaphore &sem, uint16_t dst_cta_mask = 0b11) { + if constexpr (ncta == 1) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one.b64 [%0];\n" + :: "l"(__cvta_generic_to_shared(&sem))); + } + else { + asm volatile( + "tcgen05.commit.cta_group::2.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [%0], %1;\n" + :: "l"(__cvta_generic_to_shared(&sem)), "h"(dst_cta_mask)); + } +} + +} +} + + +template constexpr int reduction_dimension = sizeof(T_AB) == 2 ? 16 : sizeof(T_AB) == 4 ? 8 : 32; + + +template +__device__ static inline void mma(D &d, const A &a, const B &b) { + constexpr int trans_b = 1 - n_trans_b; + + + constexpr int M = (trans_a ? A::cols : A::rows) * ncta; + static_assert(M == D::rows*ncta && ((ncta == 1 && (M == 64 || M == 128)) || (ncta == 2 && (M == 128 || M == 256)))); + + constexpr int N = (trans_b ? B::cols : B::rows) * ncta; + static_assert(N == D::cols); + + constexpr int K = trans_a ? A::rows : A::cols; + static_assert((trans_b ? B::rows : B::cols) == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; static_assert(std::is_same_v); + using T_D = D::T; + + constexpr int red_dim = reduction_dimension; + static_assert(K%red_dim == 0, "K dimension must be divisible by red_dim."); + + static_assert( + + (std::is_same_v && ( + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v + )) || + + (std::is_same_v && ( + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v + )) || + + (std::is_same_v && ( + std::is_same_v || + std::is_same_v + )), + "Currently unsupported type combination for matrix multiply." + ); + + static_assert(trans_a == transpose::N, "A from TMEM can't be transposed in TS (tensor-shared) MMA."); + + uint32_t idesc = detail::tcgen05::instruction_descriptor(); + kittens::st_descriptor, trans_b> b_desc(b); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + + detail::tcgen05::template tt_st( + d.addr, + a.template chunk_addr(0), + b_desc.chunk_descriptor(0), + idesc + ); + + #pragma unroll + for(int i = 1; i < K/red_dim; i++) { + detail::tcgen05::template tt_st( + d.addr, + a.template chunk_addr(i), + b_desc.chunk_descriptor(i), + idesc + ); + } +} +template +__device__ static inline void mma(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b); + detail::tcgen05::commit(sem); +} + + +template +__device__ static inline void mma(D &d, const A &a, const B &b) { + constexpr int trans_b = 1 - n_trans_b; + + + constexpr int M = (trans_a ? A::cols : A::rows) * ncta; + static_assert(M == D::rows*ncta && ((ncta == 1 && (M == 64 || M == 128)) || (ncta == 2 && (M == 128 || M == 256)))); + + constexpr int N = (trans_b ? B::cols : B::rows) * ncta; + static_assert(N == D::cols); + + + constexpr int K = (trans_a ? A::rows : A::cols); + static_assert((trans_b ? B::rows : B::cols) == K); + static_assert(std::is_same_v); + + + using T_AB = A::T; static_assert(std::is_same_v); + using T_D = D::T; + + constexpr int red_dim = reduction_dimension; + static_assert(K%red_dim == 0, "K dimension must be divisible by red_dim."); + + static_assert( + + (std::is_same_v && ( + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v + )) || + + (std::is_same_v && ( + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v + )) || + + (std::is_same_v && ( + std::is_same_v || + std::is_same_v + )), + "Currently unsupported type combination for matrix multiply." + ); + + uint32_t idesc = detail::tcgen05::instruction_descriptor(); + kittens::st_descriptor, trans_a> a_desc(a); + kittens::st_descriptor, trans_b> b_desc(b); + + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + + detail::tcgen05::template st_st( + d.addr, + a_desc.chunk_descriptor(0), + b_desc.chunk_descriptor(0), + idesc + ); + #pragma unroll + for(int i = 1; i < K/red_dim; i++) { + detail::tcgen05::template st_st( + d.addr, + a_desc.chunk_descriptor(i), + b_desc.chunk_descriptor(i), + idesc + ); + } +} +template +__device__ static inline void mma(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b); + detail::tcgen05::commit(sem); +} + + +template +__device__ static inline void mma(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + + + static_assert(std::is_same_v); + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "A and B must be fp8e4m3 or fp4e2m1_2" + ); + static_assert( + (std::is_same_v && ( + std::is_same_v && std::is_same_v + )) || + (std::is_same_v && ( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v) + )), + "SAB must be fp8e8m0 for fp8e4m3 element type, or fp8e4m3 / fp8e8m0 for fp4e2m1_2 element type"); + + static_assert(std::is_same_v, "Only float32 accumulator is supported for microscaling formats"); + using T_AB = A::T; + using T_SAB = SA::T; + using T_D = D::T; + + constexpr int block_size = std::is_same_v ? 16 : 32; + constexpr int trans_b = 1 - n_trans_b; + static_assert(trans_a == transpose::N && trans_b == transpose::N, "Only ABt supported for microscaling formats currently"); +#if !defined(KITTENS_SM103) + static_assert(mma_k_bytes == 32, "32B NVFP4 tcgen05 MMA is only supported on non-SM103."); +#endif + static_assert(mma_k_bytes == 32 || mma_k_bytes == 48, "Microscaling MMA K dimension must be 32 or 48 bytes."); + if constexpr (mma_k_bytes == 48) { + static_assert(std::is_same_v, "K96 microscaling MMA is only supported for packed FP4."); + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "K96 expects matching E8M0 (MXFP4, block32) or E4M3 (NVFP4, block16) scale factors."); + } + + + constexpr int M = (trans_a ? A::cols : A::rows) * ncta; + constexpr int N = (trans_b ? B::cols : B::rows) * ncta; + constexpr int K = std::is_same_v ? (trans_a ? A::rows : A::cols) * 2 : (trans_a ? A::rows : A::cols); + constexpr bool is_k96 = std::is_same_v && mma_k_bytes == 48; + constexpr int red_dim = std::is_same_v ? mma_k_bytes * 2 : mma_k_bytes; + static_assert(K % red_dim == 0, "K dimension must be divisible by red_dim."); + + + static_assert(M == D::rows*ncta && ((ncta == 1 && M == 128) || (ncta == 2 && (M == 128 || M == 256)))); + + + static_assert(N == D::cols && ((ncta == 1 && N%8 == 0) || (ncta == 2 && N%16 == 0))); + + + kittens::st_descriptor, trans_a> a_desc(a); + kittens::st_descriptor, trans_b> b_desc(b); + + + + + + + constexpr uint32_t idescs[4] = { + detail::tcgen05::instruction_descriptor(), + detail::tcgen05::instruction_descriptor(), + detail::tcgen05::instruction_descriptor(), + detail::tcgen05::instruction_descriptor() + }; + + + constexpr int N_offset = N / 32; + constexpr int M_offset = M / 32 / ncta; + + uint64_t a_desc0, b_desc0; + if constexpr (is_k96) { + a_desc0 = a_desc.chunk_descriptor_k96(0); + b_desc0 = b_desc.chunk_descriptor_k96(0); + } else { + a_desc0 = a_desc.chunk_descriptor(0); + b_desc0 = b_desc.chunk_descriptor(0); + } + detail::tcgen05::template st_st( + d.addr, + a_desc0, + b_desc0, + sa.addr, + sb.addr, + idescs[0] + ); + + if constexpr (is_k96) { + constexpr int sf_atoms = ((96 / block_size) + 3) / 4; + static_assert((K / red_dim) * M_offset * sf_atoms <= SA::cols / 4, "A scale tensor tile is too narrow for the K96 MMA count."); + static_assert((K / red_dim) * N_offset * sf_atoms <= SB::cols / 4, "B scale tensor tile is too narrow for the K96 MMA count."); + #pragma unroll + for (int i = 1; i < K / red_dim; i++) { + detail::tcgen05::template st_st( + d.addr, + a_desc.chunk_descriptor_k96(i), + b_desc.chunk_descriptor_k96(i), + sa.addr + i * M_offset * sf_atoms, + sb.addr + i * N_offset * sf_atoms, + idescs[0] + ); + } + } else if constexpr (std::is_same_v) { + #pragma unroll + for (int i = 1; i < K / red_dim; i++) { + detail::tcgen05::template st_st( + d.addr, + a_desc.chunk_descriptor(i), + b_desc.chunk_descriptor(i), + sa.addr + (i >> 2) * M_offset, + sb.addr + (i >> 2) * N_offset, + idescs[i % 4] + ); + } + } else if constexpr (std::is_same_v && block_size == 16) { + #pragma unroll + for (int i = 1; i < K / red_dim; i++) { + detail::tcgen05::template st_st( + d.addr, + a_desc.chunk_descriptor(i), + b_desc.chunk_descriptor(i), + sa.addr + i * M_offset, + sb.addr + i * N_offset, + idescs[0] + ); + } + } else if constexpr (std::is_same_v && block_size == 32) { + #pragma unroll + for (int i = 1; i < K / red_dim; i++) { + detail::tcgen05::template st_st( + d.addr, + a_desc.chunk_descriptor(i), + b_desc.chunk_descriptor(i), + sa.addr + (i >> 1) * M_offset, + sb.addr + (i >> 1) * N_offset, + (i & 1) ? idescs[2] : idescs[0] + ); + } + } else { + static_assert(sizeof(T_AB) == 999, "Should not reach here."); + } +} +template +__device__ static inline void mma(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma(d, a, b, sa, sb); + detail::tcgen05::commit(sem); +} + + +template +__device__ static inline void mma2(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mma2(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mma2(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + static_assert(!trans_a && trans_b, "Only ABt supported for microscaling formats currently"); + mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mma2(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + static_assert(!trans_a && trans_b, "Only ABt supported for microscaling formats currently"); + mma(d, a, b, sa, sb); +} +template +__device__ static inline void mm(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + static_assert(!trans_a && trans_b, "Only ABt supported for microscaling formats currently"); + mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mm(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + static_assert(!trans_a && trans_b, "Only ABt supported for microscaling formats currently"); + mma(d, a, b, sa, sb); +} +template +__device__ static inline void mm2(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mm2(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mm2(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + static_assert(!trans_a && trans_b, "Only ABt supported for microscaling formats currently"); + mma2(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mm2(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + static_assert(!trans_a && trans_b, "Only ABt supported for microscaling formats currently"); + mma2(d, a, b, sa, sb); +} + + +template +__device__ static inline void mma_AB(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mma_AB(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mma2_AB(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mma2_AB(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mma_ABt(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mma_ABt(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mma2_ABt(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mma2_ABt(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mma_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mma_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma(d, a, b, sa, sb); +} +template +__device__ static inline void mma2_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mma2_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma(d, a, b, sa, sb); +} +template +__device__ static inline void mma_AtB(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mma_AtB(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mma2_AtB(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mma2_AtB(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mma_AtBt(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mma_AtBt(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mma2_AtBt(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mma2_AtBt(D &d, const A &a, const B &b) { + mma2(d, a, b); +} + +template +__device__ static inline void mm_AB(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm_AB(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm2_AB(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mm2_AB(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mm_ABt(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm_ABt(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm2_ABt(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mm2_ABt(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mm_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mm_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma(d, a, b, sa, sb); +} +template +__device__ static inline void mm2_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb, semaphore &sem) { + mma(d, a, b, sa, sb, sem); +} +template +__device__ static inline void mm2_ABt(D &d, const A &a, const B &b, const SA &sa, const SB &sb) { + mma(d, a, b, sa, sb); +} +template +__device__ static inline void mm_AtB(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm_AtB(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm2_AtB(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mm2_AtB(D &d, const A &a, const B &b) { + mma2(d, a, b); +} +template +__device__ static inline void mm_AtBt(D &d, const A &a, const B &b, semaphore &sem) { + mma(d, a, b, sem); +} +template +__device__ static inline void mm_AtBt(D &d, const A &a, const B &b) { + mma(d, a, b); +} +template +__device__ static inline void mm2_AtBt(D &d, const A &a, const B &b, semaphore &sem) { + mma2(d, a, b, sem); +} +template +__device__ static inline void mm2_AtBt(D &d, const A &a, const B &b) { + mma2(d, a, b); +} + +template +__device__ static inline void mma_chunk(D &d, const A &a, const B &b, const SA &sa, const SB &sb, int ab_chunk_idx, int sf_chunk_idx, bool init) { + static_assert(mma_k_bytes == 48, "Chunked block-scaled MMA currently only supports the packed-FP4 48B (K96) form."); + static_assert(std::is_same_v && std::is_same_v, "A and B must be fp4e2m1_2"); + static_assert( + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "SAB must be fp8e4m3 or fp8e8m0"); + static_assert(std::is_same_v, "Only float32 accumulator is supported for microscaling formats"); + using T_AB = A::T; + using T_SAB = SA::T; + using T_D = D::T; + + constexpr int block_size = std::is_same_v ? 16 : 32; + constexpr int trans_b = 1 - n_trans_b; + static_assert(trans_a == transpose::N && trans_b == transpose::N, "Only ABt supported for microscaling formats currently"); + + + constexpr int M = (trans_a ? A::cols : A::rows) * ncta; + constexpr int N = (trans_b ? B::cols : B::rows) * ncta; + + + kittens::st_descriptor, trans_a> a_desc(a); + kittens::st_descriptor, trans_b> b_desc(b); + + + constexpr uint32_t idescs[4] = { + detail::tcgen05::instruction_descriptor(), + detail::tcgen05::instruction_descriptor(), + detail::tcgen05::instruction_descriptor(), + detail::tcgen05::instruction_descriptor() + }; + + + constexpr int N_offset = N / 32; + constexpr int M_offset = M / 32 / ncta; + + + + uint32_t sa_off, sb_off, idesc; + if constexpr (tight_sf) { + + + const int start_block = sf_chunk_idx * (96 / block_size); + sa_off = (start_block / 4) * M_offset; + sb_off = (start_block / 4) * N_offset; + idesc = idescs[start_block % 4]; + } else { + + constexpr int sf_atoms = ((96 / block_size) + 3) / 4; + sa_off = sf_chunk_idx * M_offset * sf_atoms; + sb_off = sf_chunk_idx * N_offset * sf_atoms; + idesc = idescs[0]; + } + + const uint64_t ad = a_desc.chunk_descriptor_k96(ab_chunk_idx); + const uint64_t bd = b_desc.chunk_descriptor_k96(ab_chunk_idx); + if (init) + detail::tcgen05::template st_st(d.addr, ad, bd, sa.addr + sa_off, sb.addr + sb_off, idesc); + else + detail::tcgen05::template st_st(d.addr, ad, bd, sa.addr + sa_off, sb.addr + sb_off, idesc); +} +template +__device__ static inline void mma2_ABt_chunk(D &d, const A &a, const B &b, const SA &sa, const SB &sb, int ab_chunk_idx, int sf_chunk_idx, bool init) { + mma_chunk(d, a, b, sa, sb, ab_chunk_idx, sf_chunk_idx, init); +} + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/thread.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/thread.cuh new file mode 100644 index 0000000..aa24901 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/thread.cuh @@ -0,0 +1,11 @@ + + + + + +#pragma once + + +#include "memory/memory.cuh" +#include "mma/mma.cuh" +#include "util/util.cuh" diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/util/sync.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/util/sync.cuh new file mode 100644 index 0000000..875d82a --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/util/sync.cuh @@ -0,0 +1,331 @@ + + + + + +#pragma once + +#include "../../../common/common.cuh" +#include "../../../types/types.cuh" + +namespace kittens { + +struct semaphore { +private: + uint64_t value; +}; +template struct barrier { + int barrier_id; + __device__ __forceinline__ barrier(int _id) : barrier_id(_id) {} + __device__ __forceinline__ barrier operator[](int i) { + return barrier(barrier_id + i); + } +}; + + + + + + + + + + + + + +__device__ static inline void init_semaphore(semaphore& bar, int thread_count, int transaction_count=0) { + void const* const ptr = &bar; + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + + asm volatile ( + "mbarrier.init.shared::cta.b64 [%0], %1;\n" + :: "r"(bar_ptr), "r"(thread_count+transaction_count) + ); +} + + + + + +__device__ static inline void invalidate_semaphore(semaphore& bar) { + void const* const ptr = &bar; + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + asm volatile ( + "mbarrier.inval.shared::cta.b64 [%0];\n" + :: "r"(bar_ptr) + ); +} + + + + + + + + +__device__ static inline void arrive(semaphore& sem) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(&sem)); + asm volatile ( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];\n" + : + : "r"(mbar_ptr) + : "memory" + ); +} +template __device__ static inline void arrive(barrier bar) { + asm volatile("bar.arrive %0, %1;\n" :: "r"(bar.barrier_id), "n"(num_warps*WARP_THREADS) : "memory"); +} + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + +__device__ static inline void arrive(semaphore& sem, uint32_t count) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(&sem)); + asm volatile ( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0], %1;\n" + : + : "r"(mbar_ptr), "r"(count) + : "memory" + ); +} +#endif + + + + + + + +__device__ static inline void wait(semaphore& sem, int kPhaseBit) { + void const* const ptr = &sem; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +#else + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.test_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "nanosleep.u32 5;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +#endif +} + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +__device__ static inline bool try_wait(semaphore &sem, int kPhaseBit) { + void const* const ptr = &sem; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + uint32_t success; + + asm volatile( + "{\n" + ".reg .pred P1; \n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2; \n" + "selp.b32 %0, 1, 0, P1; \n" + "}\n" + : "=r"(success) + : "r"(mbar_ptr), "r"(kPhaseBit) + : "memory" + ); + + return static_cast(success); +} +#endif + +__device__ static inline void careful_wait(semaphore& sem, int kPhaseBit) { + void const* const ptr = &sem; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + asm volatile ( + "{\n" + ".reg .b64 start_clock, current_clock;\n" + "mov.b64 start_clock, %clock64;\n" + ".reg .pred P_CLOCK;\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "mov.b64 current_clock, %clock64;\n" + "sub.u64 current_clock, current_clock, start_clock;\n" + "setp.ge.u64 P_CLOCK, current_clock, 1000000;\n" + "@P_CLOCK trap;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +#else + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.test_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "nanosleep.u32 5;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +#endif +} + + + + + + + +__device__ static inline int test_wait(semaphore& sem, int kPhaseBit) { + void const* const ptr = &sem; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + int result; + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "mbarrier.test_wait.parity.shared::cta.b64 P1, [%1], %2;\n" + "selp.u32 %0,1,0,P1;" + "}\n" + : "=r"(result) + : "r"(mbar_ptr), "r"(kPhaseBit) + ); + return result; +} + +__device__ static inline void arrive_and_wait(semaphore& sem, int kPhaseBit) { + arrive(sem); + wait(sem, kPhaseBit); +} +template __device__ static inline void arrive_and_wait(barrier bar) { + asm volatile("bar.sync %0, %1;\n" :: "r"(bar.barrier_id), "n"(num_warps*WARP_THREADS) : "memory"); +} + +template __device__ static inline void load_async_wait() { + if constexpr (N == 0) { + asm volatile("cp.async.wait_all;\n" ::); + } else { + asm volatile("cp.async.wait_group %0;\n" :: "n"(N)); + } + __syncwarp(); +} + + +namespace detail { +template struct size_info { + static constexpr uint32_t bytes = sizeof(std::remove_reference_t); +}; +template struct size_info { + static constexpr uint32_t elements = ST::num_elements; + static constexpr uint32_t bytes = ST::num_elements * sizeof(typename ST::dtype); +}; +template struct size_info { + static constexpr uint32_t elements = SV::length; + static constexpr uint32_t bytes = SV::length * sizeof(typename SV::dtype); +}; +} +template inline constexpr uint32_t size_bytes = 0; +template inline constexpr uint32_t size_bytes = detail::size_info::bytes + size_bytes; + + + +#if defined(KITTENS_SM10X) + +__device__ static inline void tensor_before_thread_sync() { + asm volatile("tcgen05.fence::before_thread_sync;\n"); +} +__device__ static inline void tensor_after_thread_sync() { + asm volatile("tcgen05.fence::after_thread_sync;\n"); +} + +__device__ inline static void tensor_load_wait() { + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} +__device__ inline static void tensor_store_wait() { + asm volatile("tcgen05.wait::st.sync.aligned;"); +} + +template +__device__ static inline void tensor_commit(kittens::semaphore &sem, uint16_t dst_cta_mask = 0b11) { + if constexpr (ncta == 1) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one.b64 [%0];\n" + :: "l"(__cvta_generic_to_shared(&sem))); + } + else { + asm volatile( + "tcgen05.commit.cta_group::2.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [%0], %1;\n" + :: "l"(__cvta_generic_to_shared(&sem)), "h"(dst_cta_mask)); + } +} + +#endif + + + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + +template +__device__ static inline void signal( + const barrier_t &barrier, const coord &idx, const int dst_dev_idx, const int val +) { + asm volatile("{red.release.sys.global.add.s32 [%0], %1;}" :: "l"(&barrier[dst_dev_idx][idx]), "r"(val) : "memory"); +} + +template +__device__ static inline void signal_all( + const barrier_t &barrier, const coord &idx, const int val +) { + asm volatile("{multimem.red.release.sys.global.add.s32 [%0], %1;}" :: "l"(barrier.mc_ptr_at(idx)), "r"(val) : "memory"); +} + +template +__device__ static inline void wait( + const barrier_t &barrier, const coord &idx, const int dev_idx, const int expected +) { + int val; + do { + asm volatile("{ld.relaxed.sys.global.s32 %0, [%1];}" : "=r"(val) : "l"(&barrier[dev_idx][idx]) : "memory"); + } while (val != expected); +} + +template +__device__ static inline void barrier_all( + const barrier_t &barrier, const coord &idx, const int dev_idx +) { + signal_all(barrier, idx, 1); + wait(barrier, idx, dev_idx, NUM_DEVICES); + asm volatile("{red.release.sys.global.add.s32 [%0], %1;}" :: "l"(&barrier[dev_idx][idx]), "r"(-NUM_DEVICES) : "memory"); +} + +#endif + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/util/tma.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/util/tma.cuh new file mode 100644 index 0000000..886aa21 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/util/tma.cuh @@ -0,0 +1,280 @@ + + + + + +#pragma once + +#include "../../../common/common.cuh" +#include "../../../types/types.cuh" +#include "sync.cuh" + +namespace kittens { + + + +namespace tma { + + + + + + + + + + + + + +__device__ static inline void expect_bytes(semaphore& bar, uint32_t bytes) { + void const* const ptr = &bar; + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + + asm volatile ("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" + :: "r"(bar_ptr), "r"(bytes)); +} + + + + + +template +__device__ static inline void expect(semaphore& bar, const T& _1, const args&... _2) { + expect_bytes(bar, size_bytes); +} + + + + + + +__device__ static inline void store_commit_group() { + asm volatile("cp.async.bulk.commit_group;"); +} + + + + + +template +__device__ static inline void store_async_wait() { + asm volatile ( + "cp.async.bulk.wait_group %0;" + : + : "n"(N) + : "memory" + ); +} + + + + + +template +__device__ static inline void store_async_read_wait() { + asm volatile ( + "cp.async.bulk.wait_group.read %0;" + : + : "n"(N) + : "memory" + ); +} + + + +__device__ static inline void load_async(void *dst, void *src, uint32_t size_bytes, semaphore& bar) { + asm volatile ( + "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n" + :: "r"(static_cast(__cvta_generic_to_shared(dst))), + "l"(src), "r"(size_bytes), + "r"(static_cast(__cvta_generic_to_shared(&bar))) + : "memory" + ); +} +template +__device__ static inline void load_async(T &dst, T &src, uint32_t size_bytes, semaphore& bar) { + load_async(reinterpret_cast(&dst), reinterpret_cast(&src), size_bytes, bar); +} + +__device__ static inline void store_async(void *dst, void *src, uint32_t size_bytes) { + asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory"); + asm volatile ( + "cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;\n" + :: "l"(dst), "r"(static_cast(__cvta_generic_to_shared(src))), "r"(size_bytes) + : "memory" + ); + store_commit_group(); +} +template +__device__ static inline void store_async(T &dst, T &src, uint32_t size_bytes) { + store_async(reinterpret_cast(&dst), reinterpret_cast(&src), size_bytes); +} + + + +namespace cluster { + + + + + + + +__device__ static inline void wait(semaphore& bar, int kPhaseBit) { + void const* const ptr = &bar; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +} + +__device__ static inline bool try_wait(semaphore &bar, int kPhaseBit) { + void const* const ptr = &bar; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + uint32_t success; + + asm volatile( + "{\n" + ".reg .pred P1; \n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [%1], %2; \n" + "selp.b32 %0, 1, 0, P1; \n" + "}\n" + : "=r"(success) + : "r"(mbar_ptr), "r"(kPhaseBit) + : "memory" + ); + + return static_cast(success); +} + +__device__ static inline void careful_wait(semaphore& bar, int kPhaseBit) { + void const* const ptr = &bar; + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(ptr)); + + asm volatile ( + "{\n" + ".reg .b64 start_clock, current_clock;\n" + "mov.b64 start_clock, %clock64;\n" + ".reg .pred P_CLOCK;\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "mov.b64 current_clock, %clock64;\n" + "sub.u64 current_clock, current_clock, start_clock;\n" + "setp.ge.u64 P_CLOCK, current_clock, 1000000;\n" + "@P_CLOCK trap;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +} + + + + + + + + + + + + + + + +__device__ static inline void expect_bytes(semaphore& bar, uint32_t bytes) { + uint32_t mbar_addr = static_cast(__cvta_generic_to_shared(&bar)); + asm volatile ("mbarrier.arrive.expect_tx.shared::cluster.b64 _, [%0], %1;\n" + :: "r"(mbar_addr), "r"(bytes)); +} +__device__ static inline void expect_bytes(semaphore& bar, uint32_t bytes, int dst_cta) { + uint32_t mbar_addr = static_cast(__cvta_generic_to_shared(&bar)); + uint32_t neighbor_mbar_addr; + asm volatile ( + "mapa.shared::cluster.u32 %0, %1, %2;\n" + : "=r"(neighbor_mbar_addr) + : "r"(mbar_addr), "r"(dst_cta) + ); + asm volatile ("mbarrier.arrive.expect_tx.shared::cluster.b64 _, [%0], %1;\n" + :: "r"(neighbor_mbar_addr), "r"(bytes)); +} + + + + + + + + + + + + + + + +template +__device__ static inline void expect(semaphore& bar, const T& _1, const args&... _2) { + expect_bytes(bar, size_bytes); +} + + + + + + + + + + +__device__ static inline void arrive(semaphore& bar, int dst_cta, uint32_t count=1) { + uint32_t mbar_addr = static_cast(__cvta_generic_to_shared(&bar)); + uint32_t neighbor_mbar_addr; + asm volatile ( + "mapa.shared::cluster.u32 %0, %1, %2;\n" + : "=r"(neighbor_mbar_addr) + : "r"(mbar_addr), "r"(dst_cta) + ); + asm volatile ( + "mbarrier.arrive.shared::cluster.b64 _, [%0], %1;\n" + : + : "r"(neighbor_mbar_addr), "r" (count) + : "memory" + ); +} + + + +__device__ static inline void load_async(void *dst, void *src, uint32_t size_bytes, semaphore& bar, uint16_t cta_mask) { + asm volatile ( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster [%0], [%1], %2, [%3], %4;\n" + :: "r"(static_cast(__cvta_generic_to_shared(dst))), + "l"(src), "r"(size_bytes), + "r"(static_cast(__cvta_generic_to_shared(&bar))), + "h"(cta_mask) + : "memory" + ); +} +template +__device__ static inline void load_async(T &dst, T &src, uint32_t size_bytes, semaphore& bar, uint16_t cta_mask) { + load_async(reinterpret_cast(&dst), reinterpret_cast(&src), size_bytes, bar, cta_mask); +} + +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/util/util.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/util/util.cuh new file mode 100644 index 0000000..ad52da6 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/util/util.cuh @@ -0,0 +1,327 @@ + + + + + +#pragma once + +#include "sync.cuh" +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +#include "tma.cuh" +#endif + +namespace kittens { + + + +template struct move { + __device__ static inline void lds(T& dst, uint32_t src); + __device__ static inline void sts(uint32_t dst, const T& src); + __device__ static inline void ldg(T& dst, T* src); + __device__ static inline void stg(T* dst, const T& src); +}; + +template<> struct move { + __device__ static inline void lds(bf16& dst, uint32_t src) { + asm volatile("ld.shared.b16 %0, [%1];\n" : "=h"(*(uint16_t*)&dst) : "r"(src)); + } + __device__ static inline void sts(uint32_t dst, const bf16& src) { + asm volatile("st.shared.b16 [%1], %0;\n" : : "h"(*(uint16_t*)&src), "r"(dst)); + } + __device__ static inline void ldg(bf16& dst, bf16* src) { + asm volatile("ld.global.b16 %0, [%1];\n" : "=h"(*(uint16_t*)&dst) : "l"(src)); + } + __device__ static inline void stg(bf16* dst, const bf16& src) { + asm volatile("st.global.b16 [%1], %0;\n" : : "h"(*(uint16_t*)&src), "l"(dst)); + } +}; +template<> struct move { + __device__ static inline void lds(half& dst, uint32_t src) { + asm volatile("ld.shared.b16 %0, [%1];\n" : "=h"(*(uint16_t*)&dst) : "r"(src)); + } + __device__ static inline void sts(uint32_t dst, const half& src) { + asm volatile("st.shared.b16 [%1], %0;\n" : : "h"(*(uint16_t*)&src), "r"(dst)); + } + __device__ static inline void ldg(half& dst, half* src) { + asm volatile("ld.global.b16 %0, [%1];\n" : "=h"(*(uint16_t*)&dst) : "l"(src)); + } + __device__ static inline void stg(half* dst, const half& src) { + asm volatile("st.global.b16 [%1], %0;\n" : : "h"(*(uint16_t*)&src), "l"(dst)); + } +}; +template<> struct move { + __device__ static inline void lds(float& dst, uint32_t src) { + asm volatile("ld.shared.f32 %0, [%1];\n" : "=f"(dst) : "r"(src)); + } + __device__ static inline void sts(uint32_t dst, const float& src) { + asm volatile("st.shared.f32 [%1], %0;\n" : : "f"(src), "r"(dst)); + } + __device__ static inline void ldg(float& dst, float* src) { + asm volatile("ld.global.f32 %0, [%1];\n" : "=f"(dst) : "l"(src)); + } + __device__ static inline void stg(float* dst, const float& src) { + asm volatile("st.global.f32 [%1], %0;\n" : : "f"(src), "l"(dst)); + } +}; +template<> struct move { + __device__ static inline void lds(int& dst, uint32_t src) { + asm volatile("ld.shared.u32 %0, [%1];\n" : "=r"(dst) : "r"(src)); + } + __device__ static inline void sts(uint32_t dst, const int& src) { + asm volatile("st.shared.u32 [%1], %0;\n" : : "r"(src), "r"(dst)); + } + __device__ static inline void ldg(int& dst, int* src) { + asm volatile("ld.global.u32 %0, [%1];\n" : "=r"(dst) : "l"(src)); + } + __device__ static inline void stg(int* dst, const int& src) { + asm volatile("st.global.u32 [%1], %0;\n" : : "r"(src), "l"(dst)); + } +}; + +template<> struct move { + __device__ static inline void lds(bf16_2& dst, uint32_t src) { + asm volatile("ld.shared.b32 %0, [%1];\n" : "=r"(*(uint32_t*)&dst) : "r"(src)); + } + __device__ static inline void sts(uint32_t dst, const bf16_2& src) { + asm volatile("st.shared.b32 [%1], %0;\n" : : "r"(*(uint32_t*)&src), "r"(dst)); + } + __device__ static inline void ldg(bf16_2& dst, bf16_2* src) { + asm volatile("ld.global.b32 %0, [%1];\n" : "=r"(*(uint32_t*)&dst) : "l"(src)); + } + __device__ static inline void stg(bf16_2* dst, const bf16_2& src) { + asm volatile("st.global.b32 [%1], %0;\n" : : "r"(*(uint32_t*)&src), "l"(dst)); + } + __device__ static inline void ldsm4(bf16_2& dst1, bf16_2& dst2, bf16_2& dst3, bf16_2& dst4, uint32_t src) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" : + "=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src)); + } + __device__ static inline void ldsm4t(bf16_2& dst1, bf16_2& dst2, bf16_2& dst3, bf16_2& dst4, uint32_t src) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" : + "=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src)); + } + __device__ static inline void stsm4(uint32_t dst, bf16_2& src1, bf16_2& src2, bf16_2& src3, bf16_2& src4) { + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" :: + "r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst)); + } + __device__ static inline void stsm4t(uint32_t dst, bf16_2& src1, bf16_2& src2, bf16_2& src3, bf16_2& src4) { + asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" :: + "r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst)); + } +}; +template<> struct move { + __device__ static inline void lds(half_2& dst, uint32_t src) { + asm volatile("ld.shared.b32 %0, [%1];\n" : "=r"(*(uint32_t*)&dst) : "r"(src)); + } + __device__ static inline void sts(uint32_t dst, const half_2& src) { + asm volatile("st.shared.b32 [%1], %0;\n" : : "r"(*(uint32_t*)&src), "r"(dst)); + } + __device__ static inline void ldg(half_2& dst, half_2* src) { + asm volatile("ld.global.b32 %0, [%1];\n" : "=r"(*(uint32_t*)&dst) : "l"(src)); + } + __device__ static inline void stg(half_2* dst, const half_2& src) { + asm volatile("st.global.b32 [%1], %0;\n" : : "r"(*(uint32_t*)&src), "l"(dst)); + } + __device__ static inline void ldsm4(half_2& dst1, half_2& dst2, half_2& dst3, half_2& dst4, uint32_t src) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" : + "=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src)); + } + __device__ static inline void ldsm4t(half_2& dst1, half_2& dst2, half_2& dst3, half_2& dst4, uint32_t src) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" : + "=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src)); + } + __device__ static inline void stsm4(uint32_t dst, half_2& src1, half_2& src2, half_2& src3, half_2& src4) { + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" :: + "r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst)); + } + __device__ static inline void stsm4t(uint32_t dst, half_2& src1, half_2& src2, half_2& src3, half_2& src4) { + asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" :: + "r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst)); + } +}; +template<> struct move { + __device__ static inline void lds(float2& dst, uint32_t src) { + asm volatile("ld.shared.v2.f32 {%0, %1}, [%2];\n" : "=f"(dst.x), "=f"(dst.y) : "r"(src)); + } + __device__ static inline void sts(uint32_t dst, const float2& src) { + asm volatile("st.shared.v2.f32 [%2], {%0, %1};\n" : : "f"(src.x), "f"(src.y), "r"(dst)); + } + __device__ static inline void ldg(float2& dst, float2* src) { + asm volatile("ld.global.v2.f32 {%0, %1}, [%2];\n" : "=f"(dst.x), "=f"(dst.y) : "l"(src)); + } + __device__ static inline void stg(float2* dst, const float2& src) { + asm volatile("st.global.v2.f32 [%2], {%0, %1};\n" : : "f"(src.x), "f"(src.y), "l"(dst)); + } +}; +template<> struct move { + __device__ static inline void lds(float4& dst, uint32_t src) { + asm volatile("ld.shared.v4.f32 {%0, %1, %2, %3}, [%4];\n" : "=f"(dst.x), "=f"(dst.y), "=f"(dst.z), "=f"(dst.w) : "r"(src)); + } + __device__ static inline void sts(uint32_t dst, const float4& src) { + asm volatile("st.shared.v4.f32 [%4], {%0, %1, %2, %3};\n" : : "f"(src.x), "f"(src.y), "f"(src.z), "f"(src.w), "r"(dst)); + } + __device__ static inline void ldg(float4& dst, float4* src) { + asm volatile("ld.global.v4.f32 {%0, %1, %2, %3}, [%4];\n" : "=f"(dst.x), "=f"(dst.y), "=f"(dst.z), "=f"(dst.w) : "l"(src)); + } + __device__ static inline void stg(float4* dst, const float4& src) { + asm volatile("st.global.v4.f32 [%4], {%0, %1, %2, %3};\n" : : "f"(src.x), "f"(src.y), "f"(src.z), "f"(src.w), "l"(dst)); + } +}; +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template<> struct move { + __device__ static inline void ldsm4(fp8e4m3_4& dst1, fp8e4m3_4& dst2, fp8e4m3_4& dst3, fp8e4m3_4& dst4, uint32_t src) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" : + "=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src)); + } + __device__ static inline void stsm4(uint32_t dst, fp8e4m3_4& src1, fp8e4m3_4& src2, fp8e4m3_4& src3, fp8e4m3_4& src4) { + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" :: + "r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst)); + } + +}; +template<> struct move { + __device__ static inline void ldsm4(fp8e5m2_4& dst1, fp8e5m2_4& dst2, fp8e5m2_4& dst3, fp8e5m2_4& dst4, uint32_t src) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" : + "=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src)); + } + __device__ static inline void stsm4(uint32_t dst, fp8e5m2_4& src1, fp8e5m2_4& src2, fp8e5m2_4& src3, fp8e5m2_4& src4) { + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" :: + "r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst)); + } +}; +template<> struct move { + __device__ static inline void ldsm4(int8_4& dst1, int8_4& dst2, int8_4& dst3, int8_4& dst4, uint32_t src) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" : + "=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src)); + } + __device__ static inline void stsm4(uint32_t dst, int8_4& src1, int8_4& src2, int8_4& src3, int8_4& src4) { + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" :: + "r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst)); + } +}; +template<> struct move { + __device__ static inline void ldsm4(uint8_4& dst1, uint8_4& dst2, uint8_4& dst3, uint8_4& dst4, uint32_t src) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" : + "=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src)); + } + __device__ static inline void stsm4(uint32_t dst, uint8_4& src1, uint8_4& src2, uint8_4& src3, uint8_4& src4) { + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" :: + "r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst)); + } +}; +#endif + + + +enum cache_policy { + NORMAL = 0, + EVICT_FIRST = 1, + EVICT_LAST = 2 +}; +template __device__ inline uint64_t make_cache_policy() { + uint64_t cache_policy_val; + constexpr float fraction = 1.0f; + static_assert(policy == cache_policy::EVICT_FIRST || policy == cache_policy::EVICT_LAST, "Unexpected cache policy"); + if constexpr (policy == cache_policy::EVICT_FIRST) { + asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, %1;\n" : "=l"(cache_policy_val) : "f"(fraction)); + } + else { + asm volatile("createpolicy.fractional.L2::evict_last.b64 %0, %1;\n" : "=l"(cache_policy_val) : "f"(fraction)); + } + return cache_policy_val; +} + + + +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + +namespace clc { + +struct handle { + uint4 internal_value; +}; + +struct result { + uint32_t success; + uint32_t x; + uint32_t y; + uint32_t z; +}; + + + + + + + + +__device__ static inline void schedule(handle &h, semaphore &sem) { + asm volatile("{clusterlaunchcontrol.try_cancel.async.shared::cta.mbarrier::complete_tx::bytes.multicast::cluster::all.b128 [%0], [%1];}" + :: "r"(static_cast(__cvta_generic_to_shared(&h.internal_value))), "r"(static_cast(__cvta_generic_to_shared(&sem))) + : "memory" + ); +} + + + + + +__device__ static inline result query(handle &h) { + result r; + asm volatile( + "{\n" + ".reg .pred SUCCESS;\n" + ".reg .b128 CLC_HANDLE;\n" + "ld.shared.b128 CLC_HANDLE, [%4];\n" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 SUCCESS, CLC_HANDLE;\n" + "selp.u32 %0, 1, 0, SUCCESS;\n" + "@!SUCCESS bra.uni DONE;\n" + "clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%1, %2, %3, _}, CLC_HANDLE;\n" + "fence.proxy.async.shared::cta;\n" + "DONE:\n" + "}" + : "=r"(r.success), "=r"(r.x), "=r"(r.y), "=r"(r.z) + : "r"(static_cast(__cvta_generic_to_shared(&h.internal_value))) + : "memory" + ); + return r; +} + +} + +#endif + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + + +namespace pdl { + + + + + + + + +__device__ static inline void arrive() { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +} + + + + +__device__ static inline void wait() { + asm volatile("griddepcontrol.wait;" ::: "memory"); +} + +} + +#endif + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/broker.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/broker.cuh new file mode 100644 index 0000000..427debc --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/broker.cuh @@ -0,0 +1,549 @@ + + + + + + + + + + + + + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64) + #error "KittensBroker is not supported on Windows" +#endif + +namespace kittens { + +namespace detail { +namespace broker { + +static constexpr int MAX_LOCAL_WORLD_SIZE = 72; +static constexpr int VAULT_SIZE_PER_RANK = 64; + +struct KittensVault { + static constexpr int INIT_CODE = 0x43617473; + int init; + int barrier; + int sense; + uint8_t data[MAX_LOCAL_WORLD_SIZE * VAULT_SIZE_PER_RANK]; +}; + +static constexpr int SHM_SIZE = (sizeof(KittensVault) + 4095) / 4096 * 4096; + +__host__ inline static void init_sync( + int local_rank, + volatile KittensVault *vault +) { + if (local_rank == 0) { + + vault->barrier = 0; + vault->sense = 0; + __sync_synchronize(); + vault->init = KittensVault::INIT_CODE; + } else { + while (vault->init != KittensVault::INIT_CODE) usleep(1); + __sync_synchronize(); + } +} + +__host__ inline static void sync( + int local_world_size, + volatile KittensVault *vault +) { + if (vault->init != KittensVault::INIT_CODE) + throw std::runtime_error("KittensBroker: KittensVault not initialized"); + + + int arrived = __sync_add_and_fetch(&vault->barrier, 1); + if (arrived == local_world_size) vault->sense = 1; + while (!vault->sense) usleep(1); + + + __sync_synchronize(); + + + arrived = __sync_add_and_fetch(&vault->barrier, -1); + if (arrived == 0) vault->sense = 0; + while (vault->sense) usleep(1); +} + +__host__ inline void *create_shm(const char *key, size_t size) { + int shm_fd; + shm_fd = shm_open(key, O_RDWR | O_CREAT | O_EXCL | O_CLOEXEC, 0600); + + if (shm_fd < 0) { + if (errno == EEXIST) + throw std::runtime_error("KittensBroker: Named shared memory already exists"); + throw std::runtime_error("KittensBroker: Failed to create shared memory"); + } + + if (ftruncate(shm_fd, size) != 0) { + shm_unlink(key); + close(shm_fd); + throw std::runtime_error("KittensBroker: Failed to truncate shared memory"); + } + + void *addr = mmap(0, size, PROT_READ | PROT_WRITE, MAP_SHARED, shm_fd, 0); + close(shm_fd); + if (addr == MAP_FAILED) { + shm_unlink(key); + throw std::runtime_error("KittensBroker: Failed to map to shared memory"); + } + + return addr; +} + +__host__ inline void *open_shm(const char *key, size_t size) { + int shm_fd; + while (true) { + shm_fd = shm_open(key, O_RDWR | O_CLOEXEC, 0); + if (shm_fd >= 0) + break; + if (errno != ENOENT) + throw std::runtime_error("KittensBroker: Failed to open shared memory"); + usleep(1); + } + + struct stat shm_st; + do { + if (fstat(shm_fd, &shm_st) != 0) { + shm_unlink(key); + close(shm_fd); + throw std::runtime_error("KittensBroker: Failed to open shared memory stats"); + } + usleep(1); + } while ((size_t)shm_st.st_size < size); + + void *addr = mmap(0, size, PROT_READ | PROT_WRITE, MAP_SHARED, shm_fd, 0); + close(shm_fd); + if (addr == MAP_FAILED) { + shm_unlink(key); + throw std::runtime_error("KittensBroker: Failed to map to shared memory"); + } + + return addr; +} + +__host__ inline void unlink_shm(const char *key) { + shm_unlink(key); +} + +__host__ inline void unmap_shm(void *addr, size_t size) { + munmap(addr, size); +} + +__host__ inline int create_socket(const char *key, int local_rank) { + int sock_fd; + if ((sock_fd = socket(AF_UNIX, SOCK_DGRAM | SOCK_CLOEXEC, 0)) < 0) + throw std::runtime_error("KittensBroker: Socket creation error"); + + struct sockaddr_un addr; + memset(&addr, 0, sizeof(addr)); + addr.sun_family = AF_UNIX; + + char unique_key[64]; + int n = snprintf(unique_key, sizeof(unique_key), "%s%d", key, local_rank); + if (n < 0 || n >= (int)sizeof(unique_key)) { + close(sock_fd); + throw std::runtime_error("KittensBroker: Socket name too long"); + } + + size_t len = strnlen(unique_key, sizeof(addr.sun_path)); + if (len > (sizeof(addr.sun_path) - 1)) { + close(sock_fd); + throw std::runtime_error("KittensBroker: Socket name too long"); + } + strcpy(addr.sun_path, unique_key); + unlink(unique_key); + + if (bind(sock_fd, (struct sockaddr *)&addr, SUN_LEN(&addr)) < 0) { + close(sock_fd); + throw std::runtime_error("KittensBroker: Failed to bind socket"); + } + + return sock_fd; +} + +__host__ inline void send_fd( + int sock_fd, + int data_fd, + const char *dst_key, + int dst_local_rank, + int src_local_rank +) { + union { + struct cmsghdr cm; + char* control; + } control_un; + + size_t sizeof_control = CMSG_SPACE(sizeof(int)); + control_un.control = reinterpret_cast(malloc(sizeof_control)); + if (!control_un.control) { + close(sock_fd); + close(data_fd); + throw std::runtime_error("KittensBroker: Failed to allocate a control buffer"); + } + + struct msghdr msg {}; + msg.msg_control = control_un.control; + msg.msg_controllen = sizeof_control; + + struct cmsghdr *cmptr = CMSG_FIRSTHDR(&msg); + cmptr->cmsg_len = CMSG_LEN(sizeof(int)); + cmptr->cmsg_level = SOL_SOCKET; + cmptr->cmsg_type = SCM_RIGHTS; + memmove(CMSG_DATA(cmptr), &data_fd, sizeof(data_fd)); + + struct sockaddr_un addr {}; + addr.sun_family = AF_UNIX; + char dst_unique_key[64]; + int n = snprintf(dst_unique_key, sizeof(dst_unique_key), "%s%d", dst_key, dst_local_rank); + if (n < 0 || n >= (int)sizeof(dst_unique_key)) { + free(control_un.control); + close(sock_fd); + close(data_fd); + throw std::runtime_error("KittensBroker: dst path too long"); + } + strcpy(addr.sun_path, dst_unique_key); + msg.msg_name = (void *)&addr; + msg.msg_namelen = sizeof(struct sockaddr_un); + + int payload = src_local_rank; + struct iovec iov[1]; + iov[0].iov_base = &payload; + iov[0].iov_len = sizeof(payload); + msg.msg_iov = iov; + msg.msg_iovlen = 1; + + while (true) { + ssize_t sent = sendmsg(sock_fd, &msg, 0); + if (sent <= 0) { + if (errno == EINTR) continue; + close(sock_fd); + close(data_fd); + free(control_un.control); + throw std::runtime_error("KittensBroker: Failed to send FD over socket"); + } + break; + } + + free(control_un.control); +} + +__host__ inline void recv_fd(int sock_fd, int *data_fd, int *src_local_rank) { + union { + struct cmsghdr cm; + char* control; + } control_un; + + size_t sizeof_control = CMSG_SPACE(sizeof(int)); + control_un.control = reinterpret_cast(malloc(sizeof_control)); + if (!control_un.control) { + close(sock_fd); + throw std::runtime_error("KittensBroker: Failed to allocate a control buffer"); + } + + struct msghdr msg {}; + msg.msg_control = control_un.control; + msg.msg_controllen = sizeof_control; + + int payload = -1; + struct iovec iov[1]; + iov[0].iov_base = &payload; + iov[0].iov_len = sizeof(payload); + msg.msg_iov = iov; + msg.msg_iovlen = 1; + + while (true) { + ssize_t received = recvmsg(sock_fd, &msg, 0); + if (received < 0 && errno == EINTR) { + msg.msg_controllen = sizeof_control; + msg.msg_iovlen = 1; + continue; + } + if (received < static_cast(sizeof(*data_fd))) { + free(control_un.control); + close(sock_fd); + throw std::runtime_error("KittensBroker: Failed to receive data over socket"); + } + break; + } + + if (msg.msg_flags & MSG_CTRUNC) { + free(control_un.control); + close(sock_fd); + throw std::runtime_error("KittensBroker: Control data truncated"); + } + + struct cmsghdr *cmptr = CMSG_FIRSTHDR(&msg); + if (!cmptr || + cmptr->cmsg_len != CMSG_LEN(sizeof(int)) || + cmptr->cmsg_level != SOL_SOCKET || + cmptr->cmsg_type != SCM_RIGHTS) { + free(control_un.control); + close(sock_fd); + throw std::runtime_error("KittensBroker: Failed to receive data over socket"); + } + + memmove(data_fd, CMSG_DATA(cmptr), sizeof(*data_fd)); + free(control_un.control); + *src_local_rank = payload; +} + +__host__ inline void unlink_socket(const char *key, int local_rank) { + char unique_key[64]; + int n = snprintf(unique_key, sizeof(unique_key), "%s%d", key, local_rank); + if (n < 0 || n >= (int)sizeof(unique_key)) + throw std::runtime_error("KittensBroker: Socket name too long"); + unlink(unique_key); +} + +__host__ inline void close_socket(int sock_fd) { + close(sock_fd); +} + +} +} + + + + + + + + + + + + + + + +struct KittensBroker { + + static inline constexpr const char *SHM_KEY_ = "/kittens_broker_shm"; + static inline constexpr const char *SOCK_KEY_ = "/tmp/kittens_broker.sock"; + + int local_rank_; + int local_world_size_; + + void *shm_raw_; + volatile detail::broker::KittensVault *shm_; + int sock_; + + __host__ inline KittensBroker(int local_rank, int local_world_size) + : local_rank_(local_rank), + local_world_size_(local_world_size), + shm_raw_(nullptr), + shm_(nullptr), + sock_(-1) { + if (local_rank_ < 0) + throw std::runtime_error("KittensBroker: Local rank must be non-negative"); + if (local_rank_ >= local_world_size_) + throw std::runtime_error("KittensBroker: Local rank is greater than local world size"); + if (local_world_size_ > detail::broker::MAX_LOCAL_WORLD_SIZE) + throw std::runtime_error("KittensBroker: Local world size is greater than MAX_LOCAL_WORLD_SIZE"); + + if (local_rank_ == 0) { + shm_raw_ = detail::broker::create_shm(SHM_KEY_, sizeof(detail::broker::KittensVault)); + shm_ = reinterpret_cast(shm_raw_); + memset(shm_raw_, 0, sizeof(detail::broker::KittensVault)); + } else { + shm_raw_ = detail::broker::open_shm(SHM_KEY_, sizeof(detail::broker::KittensVault)); + shm_ = reinterpret_cast(shm_raw_); + } + detail::broker::init_sync(local_rank_, shm_); + detail::broker::sync(local_world_size_, shm_); + + if (local_rank_ ==0) + detail::broker::unlink_shm(SHM_KEY_); + detail::broker::sync(local_world_size_, shm_); + + sock_ = detail::broker::create_socket(SOCK_KEY_, local_rank_); + detail::broker::sync(local_world_size_, shm_); + } + + __host__ KittensBroker(const KittensBroker&) = delete; + __host__ KittensBroker& operator=(const KittensBroker&) = delete; + + __host__ inline KittensBroker(KittensBroker&& other) noexcept + : local_rank_(other.local_rank_), + local_world_size_(other.local_world_size_), + shm_raw_(other.shm_raw_), + shm_(other.shm_), + sock_(other.sock_) { + other.local_rank_ = -1; + other.local_world_size_ = -1; + other.shm_raw_ = nullptr; + other.shm_ = nullptr; + other.sock_ = -1; + } + + __host__ inline void destroy() { + if (shm_raw_) { + detail::broker::unmap_shm(shm_raw_, sizeof(detail::broker::KittensVault)); + shm_raw_ = nullptr; + shm_ = nullptr; + } + if (sock_ >= 0) { + detail::broker::unlink_socket(SOCK_KEY_, local_rank_); + detail::broker::close_socket(sock_); + sock_ = -1; + } + local_rank_ = -1; + local_world_size_ = -1; + } + + __host__ inline KittensBroker& operator=(KittensBroker&& other) noexcept { + if (this != &other) { + destroy(); + local_rank_ = other.local_rank_; + local_world_size_ = other.local_world_size_; + shm_raw_ = other.shm_raw_; + shm_ = other.shm_; + sock_ = other.sock_; + other.local_rank_ = -1; + other.local_world_size_ = -1; + other.shm_raw_ = nullptr; + other.shm_ = nullptr; + other.sock_ = -1; + } + return *this; + } + + __host__ inline ~KittensBroker() { + destroy(); + } + + __host__ inline void sync(int num_ranks = -1) { + if (num_ranks == -1) + num_ranks = local_world_size_; + else if (num_ranks < 0 || num_ranks > local_world_size_) + throw std::runtime_error("KittensBroker: Invalid number of ranks"); + + detail::broker::sync(num_ranks, shm_); + } + + __host__ inline void exchange_data(void *dst_, const void *src_, size_t size) { + if (size > detail::broker::VAULT_SIZE_PER_RANK) + throw std::runtime_error("KittensBroker: Size is greater than VAULT_SIZE_PER_RANK"); + + uint8_t *dst = reinterpret_cast(dst_); + const uint8_t *src = reinterpret_cast(src_); + + + sync(); + memcpy(const_cast(shm_->data) + local_rank_ * detail::broker::VAULT_SIZE_PER_RANK, src, size); + sync(); + + + for (int i = 0; i < local_world_size_; i++) + memcpy(dst + i * size, const_cast(shm_->data) + i * detail::broker::VAULT_SIZE_PER_RANK, size); + } + + __host__ inline void exchange_fds(int *dst, const int data_fd) { + if (dst == nullptr) + throw std::runtime_error("KittensBroker: dst is null"); + if (data_fd < 0) + throw std::runtime_error("KittensBroker: source fd is negative"); + + + for (int i = 0; i < local_world_size_; ++i) + dst[i] = -1; + + + sync(); + + if (local_rank_ == 0) { + + dst[0] = data_fd; + for (int i = 0; i < local_world_size_ - 1; i++) { + int received_fd; + int src_local_rank; + detail::broker::recv_fd(sock_, &received_fd, &src_local_rank); + if (received_fd < 0) + throw std::runtime_error("KittensBroker: Failed to receive FD over socket"); + if (src_local_rank == local_rank_) + throw std::runtime_error("KittensBroker: Invalid source rank"); + dst[src_local_rank] = received_fd; + } + for (int dst_local_rank = 1; dst_local_rank < local_world_size_; dst_local_rank++) { + for (int src_local_rank = 0; src_local_rank < local_world_size_; src_local_rank++) { + if (dst_local_rank == src_local_rank) + continue; + detail::broker::send_fd(sock_, dst[src_local_rank], SOCK_KEY_, dst_local_rank, src_local_rank); + } + } + close(dst[0]); + dst[0] = -1; + } else { + + detail::broker::send_fd(sock_, data_fd, SOCK_KEY_, 0, local_rank_); + close(data_fd); + for (int i = 0; i < local_world_size_ - 1; i++) { + int received_fd; + int src_local_rank; + detail::broker::recv_fd(sock_, &received_fd, &src_local_rank); + if (received_fd < 0) + throw std::runtime_error("KittensBroker: Failed to receive FD over socket"); + if (src_local_rank == local_rank_) + throw std::runtime_error("KittensBroker: Invalid source rank"); + dst[src_local_rank] = received_fd; + } + } + + + sync(); + } + + __host__ inline void broadcast_fd(int *dst, const int data_fd, const int src_local_rank) { + if (src_local_rank < 0 || src_local_rank >= local_world_size_) + throw std::runtime_error("KittensBroker: Invalid source rank"); + + + sync(); + + if (local_rank_ == src_local_rank) { + if (data_fd < 0) + throw std::runtime_error("KittensBroker: Source rank has invalid FD"); + for (int dst_local_rank = 0; dst_local_rank < local_world_size_; dst_local_rank++) { + if (dst_local_rank == src_local_rank) + continue; + detail::broker::send_fd(sock_, data_fd, SOCK_KEY_, dst_local_rank, src_local_rank); + } + close(data_fd); + } else { + if (!dst) + throw std::runtime_error("KittensBroker: Destination rank has invalid buffer"); + int _src_local_rank; + detail::broker::recv_fd(sock_, dst, &_src_local_rank); + if (*dst < 0) + throw std::runtime_error("KittensBroker: Failed to receive valid FD over socket"); + if (_src_local_rank != src_local_rank) + throw std::runtime_error("KittensBroker: Invalid source rank"); + } + + + sync(); + } +}; + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/club.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/club.cuh new file mode 100644 index 0000000..5493951 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/club.cuh @@ -0,0 +1,124 @@ +#pragma once + +#include +#include +#include +#include +#include + + + + + + + + + + + + + + + + + + + +class KittensClub { +public: + __host__ inline KittensClub(const int *device_ids, const int num_devices); + __host__ inline KittensClub(const int *device_ids, const cudaStream_t *streams, const int num_devices); + __host__ inline ~KittensClub(); + + + __host__ inline void execute(std::function task); + +private: + + bool stop; + std::vector task_available; + int n_task_done; + + + std::vector workers; + + + std::vector streams; + + + __host__ inline void worker(int worker_id, int device_id); + + + std::function current_task; + + + std::mutex mutex; + std::condition_variable cond_task_available; + std::condition_variable cond_task_done; +}; + +__host__ inline KittensClub::KittensClub(const int *device_ids, const int num_devices) : stop(false), n_task_done(0) { + for (size_t dev_idx = 0; dev_idx < num_devices; ++dev_idx) { + task_available.push_back(false); + streams.push_back(0); + workers.emplace_back([this, dev_idx, device_ids] { worker(dev_idx, device_ids[dev_idx]); }); + } +} + +__host__ inline KittensClub::KittensClub(const int *device_ids, const cudaStream_t *streams_in, const int num_devices) : stop(false), n_task_done(0) { + for (size_t dev_idx = 0; dev_idx < num_devices; ++dev_idx) { + task_available.push_back(false); + streams.push_back(streams_in[dev_idx]); + workers.emplace_back([this, dev_idx, device_ids] { worker(dev_idx, device_ids[dev_idx]); }); + } +} + +__host__ inline KittensClub::~KittensClub() { + { + std::lock_guard lock(mutex); + stop = true; + } + cond_task_available.notify_all(); + for (std::thread &worker : workers) { + worker.join(); + } +} + +__host__ inline void KittensClub::execute(std::function task) { + { + std::lock_guard lock(mutex); + current_task = task; + for (size_t i = 0; i < task_available.size(); ++i) + task_available[i] = true; + } + cond_task_available.notify_all(); + { + std::unique_lock lock(mutex); + cond_task_done.wait(lock, [this] { return n_task_done == workers.size(); }); + n_task_done = 0; + } +} + +__host__ inline void KittensClub::worker(int worker_id, int device_id) { + cudaSetDevice(device_id); + while (true) { + std::function task; + { + std::unique_lock lock(mutex); + cond_task_available.wait(lock, [this, worker_id] { return stop || task_available[worker_id]; }); + + if (stop) + return; + + task = current_task; + task_available[worker_id] = false; + } + task(worker_id, streams[worker_id]); + { + std::lock_guard lock(mutex); + ++n_task_done; + if (n_task_done == workers.size()) + cond_task_done.notify_one(); + } + } +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/parallel_tensor.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/parallel_tensor.cuh new file mode 100644 index 0000000..839a14c --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/parallel_tensor.cuh @@ -0,0 +1,336 @@ +#pragma once + +#include +#include +#include + +#include +#include +#include + +#include "../types/system/vmm.cuh" +#include "../types/system/ipc.cuh" +#include "broker.cuh" + +namespace kittens { +namespace py { + + + + + + +struct TKParallelTensor { + inline static std::map, KittensBroker> brokers_; + + at::Tensor data_; + std::vector shape_; + at::ScalarType dtype_; + + std::vector raw_ptrs_; + size_t allocated_size_; + + int local_rank_; + int local_world_size_; + + bool multicast_; + void *multicast_ptr_; + size_t multicast_allocated_size_; + + detail::ipc::flavor ipc_flavor_; + + __host__ inline TKParallelTensor( + const at::Tensor &tensor, + int local_rank, + int local_world_size, + bool multicast + ) : data_(tensor), + shape_(tensor.sizes().vec()), + dtype_(tensor.scalar_type()), + raw_ptrs_(local_world_size, nullptr), + allocated_size_(tensor.nbytes()), + local_rank_(local_rank), + local_world_size_(local_world_size), + multicast_(multicast), + multicast_ptr_(nullptr), + multicast_allocated_size_(0), + ipc_flavor_(detail::ipc::flavor::LEGACY) { + + TORCH_CHECK(tensor.is_cuda(), "Tensor must be on CUDA device"); + TORCH_CHECK(tensor.is_contiguous(), "Tensor must be contiguous"); + TORCH_CHECK(tensor.dim() <= 4, "Only tensors with dim <= 4 are supported for TKParallelTensor"); + TORCH_CHECK(tensor.device().index() == local_rank_, "Tensor device index must match local_rank"); + TORCH_CHECK(local_rank_ >= 0, "local_rank must be non-negative"); + TORCH_CHECK(local_rank_ < local_world_size_, "local_rank must be less than local_world_size"); + TORCH_CHECK(!multicast, "Multicast is not supported for pre-allocated tensors"); + + brokers_.try_emplace( + {local_rank_, local_world_size_}, + local_rank_, local_world_size_ + ); + + if (brokers_.size() > 1) + std::cerr << "WARNING: 2 KittensBroker instances created in the same process. This is not safe." << std::endl; + + c10::cuda::CUDAGuard device_guard(local_rank_); + exchange_ipc_handles(); + } + + __host__ inline TKParallelTensor( + const std::vector &shape, + const at::ScalarType dtype, + int local_rank, + int local_world_size, + bool multicast + ) : shape_(shape), + dtype_(dtype), + raw_ptrs_(local_world_size, nullptr), + allocated_size_(0), + local_rank_(local_rank), + local_world_size_(local_world_size), + multicast_(multicast), + multicast_ptr_(nullptr), + multicast_allocated_size_(0), + ipc_flavor_(detail::ipc::flavor::VMM) { + + TORCH_CHECK(local_rank_ >= 0, "local_rank must be non-negative"); + TORCH_CHECK(local_rank_ < local_world_size_, "local_rank must be less than local_world_size"); + + brokers_.try_emplace( + {local_rank_, local_world_size_}, + local_rank_, local_world_size_ + ); + + if (brokers_.size() > 1) + std::cerr << "WARNING: 2 KittensBroker instances created in the same process. This is not safe." << std::endl; + + c10::cuda::CUDAGuard device_guard(local_rank_); + create_shareable_cuda_tensor(); + exchange_ipc_handles(); + + if (multicast_) + initialize_multicast(); + } + + __host__ TKParallelTensor(const TKParallelTensor&) = delete; + __host__ TKParallelTensor& operator=(const TKParallelTensor&) = delete; + __host__ TKParallelTensor& operator=(TKParallelTensor&& other) = delete; + + __host__ inline TKParallelTensor(TKParallelTensor&& other) : + data_(std::move(other.data_)), + shape_(std::move(other.shape_)), + dtype_(std::move(other.dtype_)), + raw_ptrs_(std::move(other.raw_ptrs_)), + allocated_size_(other.allocated_size_), + local_rank_(other.local_rank_), + local_world_size_(other.local_world_size_), + multicast_(other.multicast_), + multicast_ptr_(other.multicast_ptr_), + multicast_allocated_size_(other.multicast_allocated_size_), + ipc_flavor_(other.ipc_flavor_) { + other.data_ = at::Tensor(); + other.shape_.clear(); + other.dtype_ = at::ScalarType::Undefined; + other.raw_ptrs_.clear(); + other.allocated_size_ = 0; + other.local_rank_ = -1; + other.local_world_size_ = -1; + other.multicast_ = false; + other.multicast_ptr_ = nullptr; + other.multicast_allocated_size_ = 0; + } + + __host__ inline ~TKParallelTensor() { + destroy(); + } + + __host__ inline at::Tensor data() const { + return data_; + } + + __host__ inline void create_shareable_cuda_tensor() { + c10::cuda::CUDAGuard device_guard(local_rank_); + + TORCH_CHECK(!shape_.empty(), "Shape must be non-empty"); + TORCH_CHECK(shape_.size() <= 4, "Shape must have at most 4 dimensions for TKParallelTensor"); + size_t size = c10::elementSize(dtype_); + for (auto dim : shape_) { + TORCH_CHECK(dim > 0, "Size dimensions must be positive"); + size *= static_cast(dim); + } + + void *raw_ptr; + detail::vmm::vm_alloc_map_set_access( + &raw_ptr, &allocated_size_, size, local_rank_, local_world_size_); + + + int local_rank = local_rank_; + size_t allocated_size = allocated_size_; + + auto deleter = [local_rank, raw_ptr, allocated_size](void* p) mutable { + if (!p) return; + c10::cuda::CUDAGuard device_guard(local_rank); + auto stream = c10::cuda::getCurrentCUDAStream().stream(); + CUDACHECK(cudaStreamSynchronize(stream)); + detail::vmm::vm_unmap(raw_ptr, allocated_size); + }; + + at::TensorOptions options = at::TensorOptions() + .dtype(dtype_) + .device(at::kCUDA, local_rank_); + + data_ = at::from_blob(raw_ptr, shape_, std::move(deleter), options); + } + + template + __host__ inline void exchange_ipc_handles() { + using handle_t = detail::ipc::handle; + + + detail::ipc::check_support(local_rank_); + void *raw_ptr = reinterpret_cast(data_.data_ptr()); + handle_t ipc_handle; + detail::ipc::export_handle(&ipc_handle, raw_ptr); + + + std::vector all_ipc_handles(local_world_size_); + if constexpr (IPC_FLAVOR == detail::ipc::flavor::LEGACY) { + brokers_.at({local_rank_, local_world_size_}).exchange_data( + reinterpret_cast(all_ipc_handles.data()), + reinterpret_cast(&ipc_handle), + sizeof(handle_t) + ); + } else if constexpr (IPC_FLAVOR == detail::ipc::flavor::VMM) { + brokers_.at({local_rank_, local_world_size_}).exchange_fds( + reinterpret_cast(all_ipc_handles.data()), + ipc_handle.handle_ + ); + } else { + throw std::runtime_error("Invalid IPC flavor"); + } + + + for (int i = 0; i < local_world_size_; i++) { + if (i == local_rank_) + raw_ptrs_[i] = raw_ptr; + else + detail::ipc::import_handle(&raw_ptrs_[i], all_ipc_handles[i], allocated_size_, local_world_size_); + } + } + + __host__ inline void initialize_multicast() { + using handle_t = detail::ipc::handle; + + detail::vmm::multicast_check(local_rank_); + detail::ipc::check_support(local_rank_); + detail::vmm::handle multicast_handle; + + if (local_rank_ == 0) { + + detail::vmm::multicast_create_handle( + &multicast_handle, + &multicast_allocated_size_, + allocated_size_, + local_world_size_ + ); + + + if (allocated_size_ != multicast_allocated_size_) + throw std::runtime_error("Multicast allocated size does not match memory allocated size"); + + + handle_t ipc_handle; + detail::ipc::export_handle(&ipc_handle, multicast_handle); + + + brokers_.at({local_rank_, local_world_size_}).broadcast_fd(nullptr, ipc_handle.handle_, 0); + } else { + + handle_t ipc_handle; + brokers_.at({local_rank_, local_world_size_}).broadcast_fd(&ipc_handle.handle_, -1, 0); + multicast_allocated_size_ = allocated_size_; + detail::ipc::import_handle(&multicast_handle, ipc_handle, multicast_allocated_size_, local_world_size_); + } + + + detail::vmm::multicast_bind_device(multicast_handle, local_rank_); + brokers_.at({local_rank_, local_world_size_}).sync(); + + + detail::vmm::handle memory_handle; + detail::vmm::vm_retrieve_handle(&memory_handle, raw_ptrs_[local_rank_]); + detail::vmm::multicast_bind_memory(multicast_handle, memory_handle, allocated_size_); + brokers_.at({local_rank_, local_world_size_}).sync(); + + + detail::vmm::vm_map(&multicast_ptr_, multicast_handle, multicast_allocated_size_); + detail::vmm::vm_set_access(multicast_ptr_, multicast_allocated_size_, local_world_size_); + + + detail::vmm::vm_free(multicast_handle); + detail::vmm::vm_free(memory_handle); + } + + __host__ inline void destroy() { + + if (multicast_ && multicast_ptr_) { + brokers_.at({local_rank_, local_world_size_}).sync(); + detail::vmm::handle multicast_handle; + detail::vmm::vm_retrieve_handle(&multicast_handle, multicast_ptr_); + detail::vmm::vm_unmap(multicast_ptr_, multicast_allocated_size_); + detail::vmm::multicast_unbind_device(multicast_handle, multicast_allocated_size_, local_rank_); + brokers_.at({local_rank_, local_world_size_}).sync(); + detail::vmm::vm_free(multicast_handle); + } + + + for (int i = 0; i < local_world_size_; i++) { + if (i != local_rank_ && i < raw_ptrs_.size()) { + if (ipc_flavor_ == detail::ipc::flavor::LEGACY) { + detail::ipc::free_handle(raw_ptrs_[i], allocated_size_); + } else if (ipc_flavor_ == detail::ipc::flavor::VMM) { + detail::ipc::free_handle(raw_ptrs_[i], allocated_size_); + } else { + throw std::runtime_error("Invalid IPC flavor"); + } + } + } + brokers_.at({local_rank_, local_world_size_}).sync(); + + + if (data_.defined()) + data_.reset(); + + + shape_.clear(); + dtype_ = at::ScalarType::Undefined; + raw_ptrs_.clear(); + allocated_size_ = 0; + local_rank_ = -1; + local_world_size_ = -1; + multicast_ = false; + multicast_ptr_ = nullptr; + multicast_allocated_size_ = 0; + } +}; + +} +} + +#define BIND_TK_PARALLEL_TENSOR(m) \ + pybind11::class_(m, "TKParallelTensor") \ + .def(pybind11::init(), \ + pybind11::arg("tensor"), \ + pybind11::arg("local_rank"), \ + pybind11::arg("local_world_size"), \ + pybind11::arg("multicast") = false) \ + .def(pybind11::init&, const at::ScalarType&, int, int, bool>(), \ + pybind11::arg("shape"), \ + pybind11::arg("dtype"), \ + pybind11::arg("local_rank"), \ + pybind11::arg("local_world_size"), \ + pybind11::arg("multicast") = false) \ + .def("data", &kittens::py::TKParallelTensor::data) \ + .def_readonly("data_", &kittens::py::TKParallelTensor::data_) \ + .def_readonly("local_rank_", &kittens::py::TKParallelTensor::local_rank_) \ + .def_readonly("local_world_size_", &kittens::py::TKParallelTensor::local_world_size_) diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/profiler.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/profiler.cuh new file mode 100644 index 0000000..6ecf712 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/profiler.cuh @@ -0,0 +1,148 @@ +#pragma once + +#include + +#ifdef KITTENS_SM10X + +#include +#include +#include + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +template +struct __align__(1024) TKProfiler { + static constexpr int TIMING_WIDTH = _TIMING_WIDTH; + int timings[TIMING_WIDTH]; + uint64_t start_time; + using timing_t = kittens::gl; + + + + + + __device__ inline void init() { + if (threadIdx.x < TIMING_WIDTH) + timings[threadIdx.x] = 0; + if (threadIdx.x == 0) { + uint64_t start_time_reg; + asm volatile("{mov.u64 %0, %%globaltimer;}" : "=l"(start_time_reg)); + uint32_t start_time_addr = static_cast(__cvta_generic_to_shared(&start_time)); + asm volatile("{st.shared.b64 [%0], %1;}" : : "r"(start_time_addr), "l"(start_time_reg)); + } + __syncthreads(); + } + + + + + + __device__ inline void record(int index) { + uint64_t current_time; + asm volatile("{mov.u64 %0, %%globaltimer;}" : "=l"(current_time)); + timings[index] = static_cast(current_time - start_time); + } + + + + + + + + __device__ inline void store_and_reset(const timing_t &timings_global, int i, int j) { + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(timings)); + uint64_t dst_ptr = reinterpret_cast(&timings_global[{i, j, 0}]); + asm volatile("{fence.proxy.async.shared::cta;}" ::: "memory"); + asm volatile("{cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;}" + :: "l"(dst_ptr), "r"(src_ptr), "n"(TIMING_WIDTH*sizeof(int)) : "memory"); + kittens::tma::store_commit_group(); + kittens::tma::store_async_read_wait(); + asm volatile("{st.bulk.weak [%0], %1, 0;}" :: "r"(src_ptr), "n"(TIMING_WIDTH*sizeof(int)) : "memory"); + } + + + + + + + + __host__ static inline timing_t allocate(int N, int M) { + const size_t timings_size = static_cast(N)*M*TIMING_WIDTH; + int *d_timings; + CUDACHECK(cudaMalloc(&d_timings, timings_size*sizeof(int))); + CUDACHECK(cudaMemset(d_timings, 0, timings_size*sizeof(int))); + return timing_t{d_timings, nullptr, static_cast(N), static_cast(M), nullptr}; + } + + + + + + __host__ static inline void deallocate(timing_t &timings_global) { + CUDACHECK(cudaFree(timings_global.raw_ptr)); + } + + + + + + + __host__ static inline void save(timing_t &timings_global, const std::string &filename) { + const int M = static_cast(timings_global.depth()); + const int N = static_cast(timings_global.rows()); + const size_t timings_size = static_cast(M)*N*TIMING_WIDTH; + + std::vector h_timings(timings_size); + CUDACHECK(cudaMemcpy(h_timings.data(), timings_global.raw_ptr, timings_size*sizeof(int), cudaMemcpyDeviceToHost)); + + std::ofstream outfile(filename); + if (outfile.is_open()) { + outfile << "# shape: (" << M << ", " << N << ", " << TIMING_WIDTH << ")\n"; + for (int i = 0; i < M; i++) { + for (int j = 0; j < N; j++) { + for (int k = 0; k < TIMING_WIDTH; k++) { + int idx = (i * N + j) * TIMING_WIDTH + k; + outfile << h_timings[idx]; + if (k < TIMING_WIDTH - 1) outfile << " "; + } + outfile << "\n"; + } + } + outfile.close(); + } + } +}; + +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/pyutils.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/pyutils.cuh new file mode 100644 index 0000000..4aad211 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/pyutils.cuh @@ -0,0 +1,94 @@ +#pragma once + +#include +#include + +namespace kittens { +namespace py { + +template struct from_object { + __host__ static T make(pybind11::object obj) { + return obj.cast(); + } +}; +template struct from_object { + __host__ static GL make(pybind11::object obj) { + + if (pybind11::hasattr(obj, "__class__") && + obj.attr("__class__").attr("__name__").cast() == "Tensor") { + + + if (!obj.attr("is_contiguous")().cast()) { + throw std::runtime_error("Tensor must be contiguous"); + } + if (obj.attr("device").attr("type").cast() == "cpu") { + throw std::runtime_error("Tensor must be on CUDA device"); + } + + + std::array shape = {1, 1, 1, 1}; + auto py_shape = obj.attr("shape").cast(); + size_t dims = py_shape.size(); + if (dims > 4) { + throw std::runtime_error("Expected Tensor.ndim <= 4"); + } + for (size_t i = 0; i < dims; ++i) { + shape[4 - dims + i] = pybind11::cast(py_shape[i]); + } + + + uint64_t data_ptr = obj.attr("data_ptr")().cast(); + + + return make_gl(data_ptr, shape[0], shape[1], shape[2], shape[3]); + } + throw std::runtime_error("Expected a torch.Tensor"); + } +}; + + +template struct first_arg; +template struct first_arg { using type = T; }; +template using first_arg_t = typename first_arg::type; + + +template +__host__ constexpr MT TGlobal::* member_cast(MT TBase::* p) { + return static_cast(p); +} + +template struct trait; +template using object = pybind11::object; +template struct trait { using member_type = MT; using type = T; }; +template concept has_dynamic_shared_memory = requires(T t) { { t.dynamic_shared_memory() } -> std::convertible_to; }; +template __host__ static void bind_kernel(auto m, auto name, MemberPtrs... member_ptrs) { + using TGlobal = first_arg_t; + ((void)member_cast(member_ptrs), ...); + m.def(name, [=](object... args, pybind11::kwargs kwargs) { + TGlobal __g__ {from_object::member_type>::make(args)...}; + cudaStream_t raw_stream = nullptr; + if (kwargs.contains("stream")) { + + uintptr_t stream_ptr = kwargs["stream"].attr("cuda_stream").cast(); + raw_stream = reinterpret_cast(stream_ptr); + } + if constexpr (has_dynamic_shared_memory) { + int __dynamic_shared_memory__ = (int)__g__.dynamic_shared_memory(); + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, __dynamic_shared_memory__); + kernel<<<__g__.grid(), __g__.block(), __dynamic_shared_memory__, raw_stream>>>(__g__); + } else { + kernel<<<__g__.grid(), __g__.block(), 0, raw_stream>>>(__g__); + } + }); +} +template __host__ static void bind_function(auto m, auto name, MemberPtrs... member_ptrs) { + m.def(name, [=](object... args) { + using TGlobal = first_arg_t; + ((void)member_cast(member_ptrs), ...); + TGlobal __g__ {from_object::member_type>::make(args)...}; + function(__g__); + }); +} + +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/torchutils.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/torchutils.cuh new file mode 100644 index 0000000..59763a5 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/torchutils.cuh @@ -0,0 +1,227 @@ +#pragma once + +#include +#include + +#include +#include "parallel_tensor.cuh" + +#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor") +#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") +#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) + +namespace kittens { +namespace py { + +template +concept has_min_blocks_per_sm = requires { std::integral_constant{}; }; + +template +__host__ consteval int min_blocks_per_sm() { + if constexpr(has_min_blocks_per_sm) + return Config::MIN_BLOCKS_PER_SM; + else + return 1; +} + +template +__global__ +__launch_bounds__(Config::NUM_THREADS, min_blocks_per_sm()) +void global_kernel(const __grid_constant__ Globals G) { + Kernel(G); +} + +template +__host__ static inline void tensor_check(const at::Tensor &t) { + TORCH_CHECK(t.is_cuda() || t.is_pinned(), "Tensor must be on CUDA device or pinned memory") + TORCH_CHECK(t.is_contiguous(), "Tensor must be contiguous") + TORCH_CHECK(t.dim() <= 4, "Expected Tensor.dim() <= 4"); + + if constexpr (!TypeCheck) { + return; + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Bool, "Tensor has invalid dtype (expected bool)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Char, "Tensor has invalid dtype (expected int8)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Byte, "Tensor has invalid dtype (expected uint8)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Short, "Tensor has invalid dtype (expected int16)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::UInt16, "Tensor has invalid dtype (expected uint16)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Int, "Tensor has invalid dtype (expected int32)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::UInt32, "Tensor has invalid dtype (expected uint32)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Long, "Tensor has invalid dtype (expected int64)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::UInt64, "Tensor has invalid dtype (expected uint64)"); +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Float8_e4m3fn, "Tensor has invalid dtype (expected fp8e4m3)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Float8_e5m2, "Tensor has invalid dtype (expected fp8e5m2)"); +#endif +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Float8_e8m0fnu || t.dtype() == at::ScalarType::Byte, "Tensor has invalid dtype (expected fp8e8m0)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Float4_e2m1fn_x2, "Tensor has invalid dtype (expected fp4_2)"); +#endif + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::BFloat16, "Tensor has invalid dtype (expected bfloat16)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Half, "Tensor has invalid dtype (expected float16)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Float, "Tensor has invalid dtype (expected float32)"); + } else if constexpr (std::is_same_v) { + TORCH_CHECK(t.dtype() == at::ScalarType::Double, "Tensor has invalid dtype (expected float64)"); + } else { + TORCH_CHECK(false, "Unsupported dtype"); + } +} + +template +__host__ static inline void parallel_tensor_check(const TKParallelTensor& t) { + tensor_check(t.data_); + TORCH_CHECK(t.data_.sizes().vec() == t.shape_, "Shape mismatch between TKParallelTensor and the underlying tensor"); + TORCH_CHECK(t.data_.dtype() == t.dtype_, "Dtype mismatch between TKParallelTensor and the underlying tensor"); + TORCH_CHECK(t.raw_ptrs_.size() == PGL::num_devices, "Number of devices mismatch between PGL and TKParallelTensor"); + TORCH_CHECK(t.local_rank_ == t.data_.device().index(), "Current tensor device index mismatch within TKParallelTensor"); + TORCH_CHECK(t.local_world_size_ == PGL::num_devices, "Number of devices mismatch between PGL and TKParallelTensor"); + TORCH_CHECK(t.multicast_ == PGL::multicast, "Multicast mismatch between PGL and TKParallelTensor"); + TORCH_CHECK(t.raw_ptrs_[t.local_rank_] == reinterpret_cast(t.data_.data_ptr()), "Current tensor data pointer not found in TKParallelTensor's raw_ptrs_"); +} + +template +__host__ static inline GL tensor_to_gl(const at::Tensor &t) { + tensor_check(t); + + std::array shape = {1, 1, 1, 1}; + for (int i = 0; i < static_cast(t.dim()); ++i) + shape[4 - t.dim() + i] = static_cast(t.size(i)); + + uint64_t data_ptr = reinterpret_cast(t.data_ptr()); + + return ::kittens::make_gl(data_ptr, shape[0], shape[1], shape[2], shape[3]); +} + +template +__host__ static inline GL tensor_to_gl(const at::Tensor &t, int B, int D, int R, int C) { + tensor_check(t); + + return ::kittens::make_gl(reinterpret_cast(t.data_ptr()), B, D, R, C); +} + +template +__host__ static inline PGL parallel_tensor_to_pgl(TKParallelTensor &t) { + parallel_tensor_check(t); + + std::array shape = {1, 1, 1, 1}; + for (int i = 0; i < static_cast(t.data_.dim()); ++i) { + shape[4 - t.data_.dim() + i] = static_cast(t.data_.size(i)); + } + + if constexpr (PGL::multicast) + return ::kittens::make_pgl( + reinterpret_cast(t.multicast_ptr_), reinterpret_cast(t.raw_ptrs_.data()), shape[0], shape[1], shape[2], shape[3]); + else + return ::kittens::make_pgl( + reinterpret_cast(t.raw_ptrs_.data()), shape[0], shape[1], shape[2], shape[3]); +} + +template +__host__ static inline PGL parallel_tensor_to_pgl(TKParallelTensor &t, int B, int D, int R, int C) { + parallel_tensor_check(t); + + if constexpr (PGL::multicast) + return ::kittens::make_pgl( + reinterpret_cast(t.multicast_ptr_), reinterpret_cast(t.raw_ptrs_.data()), B, D, R, C); + else + return ::kittens::make_pgl( + reinterpret_cast(t.raw_ptrs_.data()), B, D, R, C); +} + +template +__host__ static inline GL make_fake_gl(const int batch, const int depth, const int rows, const int cols) { + return ::kittens::make_gl(reinterpret_cast(nullptr), batch, depth, rows, cols); +} + +__host__ static inline void _device_check(const at::Tensor& first, const at::Tensor& second) { + TORCH_CHECK(first.device() == second.device(), "All tensors must be on the same device"); +} + +template +__host__ static inline void device_check(const T1& first, const Ts&... rest) { + (_device_check(first, rest), ...); +} + +__host__ static inline void _parallel_tensor_check(const TKParallelTensor& first, const TKParallelTensor& second) { + TORCH_CHECK(first.local_rank_ == second.local_rank_, "All parallel tensors must have the same local_rank"); + TORCH_CHECK(first.local_world_size_ == second.local_world_size_, "All parallel tensors must have the same local_world_size"); +} + +template +__host__ static inline void parallel_tensor_check(const T1& first, const Ts&... rest) { + (_parallel_tensor_check(first, rest), ...); +} + +template +concept static_grid = requires { Config::NUM_BLOCKS; }; + +template +concept static_block = requires { Config::NUM_THREADS; }; + +template +concept static_dynamic_shared_memory = requires { Config::DYNAMIC_SHARED_MEMORY; }; + +template +concept has_pdl_config = requires { { Config::USE_PDL } -> std::convertible_to; }; +template +inline constexpr bool use_pdl = false; +template requires has_pdl_config +inline constexpr bool use_pdl = Config::USE_PDL; + +template +__host__ static inline void launch_kernel(const Globals &G) { + dim3 grid; + if constexpr (static_grid) + grid = dim3{Config::NUM_BLOCKS, 1, 1}; + else + grid = G.grid(); + + dim3 block; + if constexpr (static_block) + block = dim3{Config::NUM_THREADS, 1, 1}; + else + block = G.block(); + + int dynamic_shared_memory; + if constexpr (static_dynamic_shared_memory) + dynamic_shared_memory = static_cast(Config::DYNAMIC_SHARED_MEMORY); + else + dynamic_shared_memory = G.dynamic_shared_memory(); + + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + +#if defined(KITTENS_SM90) + static_assert(Config::CLUSTER_SIZE <= 8, "Cluster size must be less than or equal to 8 for Hopper"); +#elif defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert(Config::CLUSTER_SIZE <= 16, "Cluster size must be less than or equal to 16 for Blackwell"); + if constexpr (Config::CLUSTER_SIZE > 8) + CUDACHECK(cudaFuncSetAttribute(global_kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1)); +#endif + CUDACHECK(cudaFuncSetAttribute(global_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dynamic_shared_memory)); + + if constexpr (Config::CLUSTER_SIZE <= 1) { + LaunchConfig> launch_config(grid, block, dynamic_shared_memory, stream); + CUDACHECK(cudaLaunchKernelEx(launch_config, global_kernel, G)); + } else { + LaunchConfig> launch_config(grid, block, dynamic_shared_memory, stream, Config::CLUSTER_SIZE); + CUDACHECK(cudaLaunchKernelEx(launch_config, global_kernel, G)); + } +} + +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/global/cgl.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/global/cgl.cuh new file mode 100644 index 0000000..c285239 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/global/cgl.cuh @@ -0,0 +1,50 @@ + + + + + +#pragma once + +#include "gl.cuh" + +namespace kittens { + + + +namespace ducks { +namespace cgl { +struct identifier {}; +} +} + + + + + + +template +struct cgl { + using identifier = ducks::cgl::identifier; + using component = _GL; + using T = component::T; + using T2 = component::T2; + using dtype = component::dtype; + component real, imag; +}; + +namespace ducks { +namespace cgl { + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v; +} +} + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/global/gl.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/global/gl.cuh new file mode 100644 index 0000000..16ed80b --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/global/gl.cuh @@ -0,0 +1,250 @@ + + + + + +#pragma once + +#include "../../common/common.cuh" +#include "../shared/shared.cuh" +#include "util.cuh" +#if (defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120)) && !defined(KITTENS_NO_HOST) +#include "tma.cuh" +#endif + +namespace kittens { + + + +struct dim { + static constexpr int BATCH = 0; + static constexpr int DEPTH = 1; + static constexpr int ROW = 2; + static constexpr int COL = 3; +}; + + + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +namespace ducks { +namespace tma { +namespace descriptor { +struct identifier {}; +template concept all = requires { + typename T::identifier; +} && std::is_same_v; +} +} +} +namespace detail { +namespace tma { +template struct descriptor_copy_helper {}; +template struct descriptor_copy_helper<_T> { static constexpr int value = _T::axis; using T = _T::T; }; +template struct descriptor_copy_helper<_T> { static constexpr int value = 2; using T = _T; }; +template struct descriptor_copy_helper<_T> { static constexpr int value = -1; using T = _T; }; +template using descriptor_copy_helper_t = descriptor_copy_helper::T; +template static constexpr int descriptor_copy_helper_v = descriptor_copy_helper::value; +} +} +namespace tma { +template struct descriptor { + using identifier = ducks::tma::descriptor::identifier; + using T = detail::tma::descriptor_copy_helper_t<_T>; + static_assert(ducks::st::all || ducks::sv::all || ducks::tma::descriptor::all, "Must be a shared TK type to generate a TMA descriptor."); + static constexpr int axis = ( + ducks::tma::descriptor::all<_T> ? detail::tma::descriptor_copy_helper_v<_T> : + (_axis != -9999) ? _axis : detail::tma::descriptor_copy_helper_v<_T>); + static_assert((kittens::ducks::st::all && axis >= 0 && axis <= 2) || (kittens::ducks::sv::all && axis == -1), "Internal template error detected."); +}; +} +#endif + +namespace detail { +template +struct descriptor_dict { +#ifndef KITTENS_NO_HOST + __host__ descriptor_dict() {} + template __host__ descriptor_dict(T _, int b, int d, int r, int c) {} +#endif + __host__ __device__ descriptor_dict(const descriptor_dict &other) {} +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + template __device__ const CUtensorMap* get() const { + static_assert( + std::is_same_v && std::is_same_v, + "SKILL ISSUE: Requested a TMA descriptor for a type not initialized in the global layout." + ); + return nullptr; + } +#endif +}; + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template +struct descriptor_dict<_T, Args...> { + static_assert(ducks::sv::all<_T> || ducks::st::all<_T> || ducks::tma::descriptor::all<_T>, "Must be a shared TK type to generate a TMA descriptor."); + using DESC = kittens::tma::descriptor<_T>; + CUtensorMap tma_desc; + descriptor_dict other_descs; +#ifndef KITTENS_NO_HOST + __host__ descriptor_dict() {} + __host__ descriptor_dict(typename DESC::T::dtype *data, int b, int d, int r, int c): other_descs(data, b, d, r, c) { + kittens::detail::tma::create_tensor_map(&tma_desc, data, b, d, r, c); + } +#endif + __host__ __device__ inline descriptor_dict(const descriptor_dict &other) : + tma_desc(other.tma_desc), other_descs(other.other_descs) {} + template __device__ inline const CUtensorMap* get() const { + if constexpr (std::is_same_v && DESC::axis == axis) { return &tma_desc; } + else { return other_descs.template get(); } + } +}; +#endif +} + + + +namespace ducks { +namespace gl { +struct identifier {}; +} +} + +template +struct gl { + static_assert(!std::is_same_v<_T, char>, "Bare 'char' is not permitted because its signedness is ISA-dependent."); +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert(!std::is_same_v<_T, fp4e2m1>, "For FP4 types, you must use a packed type (i.e., fp4e2m1_2 or fp4e2m1_4)."); +#endif + using identifier = ducks::gl::identifier; + + using T = base_types::packing<_T>::unpacked_type; + using T2 = base_types::packing<_T>::packed_type; + using dtype = T; + + T* raw_ptr; + + static constexpr int __b__ = b, __d__ = d, __r__ = r, __c__ = c; + + ducks::gl::make_dim_t batch_internal; + ducks::gl::make_dim_t depth_internal; + ducks::gl::make_dim_t rows_internal; + ducks::gl::make_dim_t cols_internal; + +#ifdef KITTENS_NO_HOST + __device__ constexpr int batch() const { if constexpr (b > 0) return b; else return batch_internal; } + __device__ constexpr int depth() const { if constexpr (d > 0) return d; else return depth_internal; } + __device__ constexpr int rows() const { if constexpr (r > 0) return r; else return rows_internal; } + __device__ constexpr int cols() const { if constexpr (c > 0) return c; else return cols_internal; } +#else + template __device__ __host__ static constexpr std::enable_if_t<(B > 0), int> batch() { return B; } + template __device__ __host__ std::enable_if_t<(B == -1), int> batch() const { return batch_internal; } + template __device__ __host__ static constexpr std::enable_if_t<(D > 0), int> depth() { return D; } + template __device__ __host__ std::enable_if_t<(D == -1), int> depth() const { return depth_internal; } + template __device__ __host__ static constexpr std::enable_if_t<(R > 0), int> rows() { return R; } + template __device__ __host__ std::enable_if_t<(R == -1), int> rows() const { return rows_internal; } + template __device__ __host__ static constexpr std::enable_if_t<(C > 0), int> cols() { return C; } + template __device__ __host__ std::enable_if_t<(C == -1), int> cols() const { return cols_internal; } +#endif + __device__ __host__ inline size_t numel() const { return static_cast(batch()) * depth() * rows() * cols(); } + + detail::descriptor_dict tma_descs; + + __host__ __device__ gl() = delete; +#ifndef KITTENS_NO_HOST + __host__ inline gl(T *_data, + ducks::gl::make_arg_t _batch, + ducks::gl::make_arg_t _depth, + ducks::gl::make_arg_t _rows, + ducks::gl::make_arg_t _cols) : + raw_ptr(_data), batch_internal(_batch), depth_internal(_depth), rows_internal(_rows), cols_internal(_cols) { + tma_descs = detail::descriptor_dict(raw_ptr, batch_internal, depth_internal, rows_internal, cols_internal); + } +#endif + __host__ __device__ inline gl(const gl &other) : + raw_ptr(other.raw_ptr), batch_internal(other.batch_internal), depth_internal(other.depth_internal), rows_internal(other.rows_internal), cols_internal(other.cols_internal), tma_descs(other.tma_descs) {} +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + template __device__ inline const CUtensorMap* get_tma() const { + return tma_descs.template get(); + } + template __device__ inline void prefetch_tma() const { + const CUtensorMap *tma_desc = tma_descs.template get(); + asm volatile ("{prefetch.tensormap [%0];}" :: "l"(reinterpret_cast(tma_desc)) : "memory"); + } +#endif + __device__ inline T& operator[](const coord &idx) const { + return raw_ptr[(((size_t)idx.b*depth() + idx.d)*rows() + idx.r)*cols() + idx.c]; + } + template __device__ inline size_t shape() const { + static_assert(axis==0 || axis==1 || axis==2 || axis==3, "Axis must be 0, 1, 2, or 3."); + if constexpr (axis==0) { return size_t(batch()); } + else if constexpr (axis==1) { return size_t(depth()); } + else if constexpr (axis==2) { return size_t(rows()); } + else if constexpr (axis==3) { return size_t(cols()); } + } + template __device__ inline size_t stride() const { + static_assert(axis==0 || axis==1 || axis==2 || axis==3, "Axis must be 0, 1, 2, or 3."); + if constexpr (axis==0) { return (size_t)depth()*rows()*cols(); } + else if constexpr (axis==1) { return (size_t)rows()*cols(); } + else if constexpr (axis==2) { return (size_t)cols(); } + else if constexpr (axis==3) { return 1; } + } +}; + +template using gl3 = gl<_T, 1, d, r, c, TMA_Types...>; +template using gl2 = gl<_T, 1, 1, r, c, TMA_Types...>; +template using gl1 = gl<_T, 1, 1, 1, c, TMA_Types...>; + +namespace ducks { +namespace gl { + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v; +} +} + +#ifndef KITTENS_NO_HOST + + + + + + + +template __host__ auto make_unsafe_gl_arg(int param) { + if constexpr (N > 0) { return nullptr; } + else { return param; } +} +template __host__ inline GL make_gl(uint64_t data, int b, int d, int r, int c) { + if constexpr (safe) { + if(GL::__b__ > 0 && b != GL::__b__) { + throw std::runtime_error("Batch dimension mismatch. Expected: " + std::to_string(GL::__b__) + ", Got: " + std::to_string(b)); + } + if(GL::__d__ > 0 && d != GL::__d__) { + throw std::runtime_error("Depth dimension mismatch. Expected: " + std::to_string(GL::__d__) + ", Got: " + std::to_string(d)); + } + if(GL::__r__ > 0 && r != GL::__r__) { + throw std::runtime_error("Row dimension mismatch. Expected: " + std::to_string(GL::__r__) + ", Got: " + std::to_string(r)); + } + if(GL::__c__ > 0 && c != GL::__c__) { + throw std::runtime_error("Column dimension mismatch. Expected: " + std::to_string(GL::__c__) + ", Got: " + std::to_string(c)); + } + } + return GL( + reinterpret_cast(data), + make_unsafe_gl_arg(b), + make_unsafe_gl_arg(d), + make_unsafe_gl_arg(r), + make_unsafe_gl_arg(c) + ); +} + +#endif + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/global/global.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/global/global.cuh new file mode 100644 index 0000000..c7380ce --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/global/global.cuh @@ -0,0 +1,10 @@ + + + + + +#pragma once + +#include "util.cuh" +#include "gl.cuh" +#include "cgl.cuh" diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/global/tma.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/global/tma.cuh new file mode 100644 index 0000000..b58e428 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/global/tma.cuh @@ -0,0 +1,337 @@ +#pragma once + +#include "../../common/common.cuh" +#include "../shared/shared.cuh" +#include "util.cuh" + +namespace kittens { +namespace detail { +namespace tma { + + + + + + + + + + + + + + + +template +__host__ static inline void create_tensor_map( + CUtensorMap *tma_map, const typename ST::dtype *src, int batch, int depth, int rows, int cols +) { + using dtype = typename ST::dtype; + static_assert(axis==0 || axis==1 || axis==2, "axis must be 0, 1, or 2"); +#ifdef KITTENS_SM10X + static_assert(!(std::is_same_v && axis != 2), "Axes 0 and 1 are not yet supported for FP4 type"); +#endif + + constexpr uint32_t tma_dim = ST::swizzle ? 5 : 4; + void *global_addr = (void*)(src); + + constexpr CUtensorMapDataType tma_format = ( + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_BFLOAT16 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_FLOAT16 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_FLOAT32 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_INT32 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : +#ifdef KITTENS_SM10X + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : +#endif + CUtensorMapDataType(-1) + ); + constexpr CUtensorMapInterleave tma_interleave = CU_TENSOR_MAP_INTERLEAVE_NONE; + constexpr CUtensorMapL2promotion tma_l2Promotion = CU_TENSOR_MAP_L2_PROMOTION_NONE; + constexpr CUtensorMapFloatOOBfill tma_oobFill = CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE; + constexpr CUtensorMapSwizzle tma_swizzle = ST::swizzle ? ( + ST::swizzle_bytes == 32 ? CU_TENSOR_MAP_SWIZZLE_32B : + ST::swizzle_bytes == 64 ? CU_TENSOR_MAP_SWIZZLE_64B : + ST::swizzle_bytes == 128 ? CU_TENSOR_MAP_SWIZZLE_128B : + CU_TENSOR_MAP_SWIZZLE_NONE + ) : CU_TENSOR_MAP_SWIZZLE_NONE; + + uint64_t gmem_shape [5] = {0, 0, 0, 0, 0}; + uint64_t gmem_stride[4] = {0, 0, 0, 0}; + uint32_t smem_shape [5] = {0, 0, 0, 0, 0}; + uint32_t smem_stride[5] = {1, 1, 1, 1, 1}; + + constexpr uint64_t shared_tile_height = ST::rows; + constexpr uint64_t shared_tile_width = ST::cols; + + + constexpr int swizzle_elements = ST::swizzle_bytes / sizeof(dtype); + + if constexpr (ST::swizzle) { + if constexpr (axis == 2) { + gmem_shape[0] = swizzle_elements; + gmem_shape[1] = (uint64_t)rows; + gmem_shape[2] = (uint64_t)(cols+swizzle_elements-1) / swizzle_elements; + gmem_shape[3] = (uint64_t)depth; + gmem_shape[4] = (uint64_t)batch; + + gmem_stride[0] = (uint64_t)cols * sizeof(dtype); + gmem_stride[1] = ST::swizzle_bytes; + gmem_stride[2] = (uint64_t)rows * cols * sizeof(dtype); + gmem_stride[3] = (uint64_t)depth * rows * cols * sizeof(dtype); + } + else if constexpr (axis == 1) { + gmem_shape[0] = swizzle_elements; + gmem_shape[1] = (uint64_t)depth; + gmem_shape[2] = (uint64_t)(cols+swizzle_elements-1) / swizzle_elements; + gmem_shape[3] = (uint64_t)rows; + gmem_shape[4] = (uint64_t)batch; + + gmem_stride[0] = (uint64_t)rows * cols * sizeof(dtype); + gmem_stride[1] = ST::swizzle_bytes; + gmem_stride[2] = (uint64_t)cols * sizeof(dtype); + gmem_stride[3] = (uint64_t)depth * rows * cols * sizeof(dtype); + + } + else { + gmem_shape[0] = swizzle_elements; + gmem_shape[1] = (uint64_t)batch; + gmem_shape[2] = (uint64_t)(cols+swizzle_elements-1) / swizzle_elements; + gmem_shape[3] = (uint64_t)rows; + gmem_shape[4] = (uint64_t)depth; + + gmem_stride[0] = (uint64_t)depth * rows * cols * sizeof(dtype); + gmem_stride[1] = ST::swizzle_bytes; + gmem_stride[2] = (uint64_t)cols * sizeof(dtype); + gmem_stride[3] = (uint64_t)rows * cols * sizeof(dtype); + } + smem_shape[0] = swizzle_elements; + smem_shape[1] = shared_tile_height; + smem_shape[2] = shared_tile_width / swizzle_elements; + smem_shape[3] = 1; + smem_shape[4] = 1; + } else { + static_assert(axis == 2, "For non-swizzled tiles, only axis 2 is supported."); + + gmem_shape[0] = (uint64_t)cols; + gmem_shape[1] = (uint64_t)rows; + gmem_shape[2] = (uint64_t)depth; + gmem_shape[3] = (uint64_t)batch; + + gmem_stride[0] = (uint64_t)cols * sizeof(dtype); + gmem_stride[1] = (uint64_t)rows * cols * sizeof(dtype); + gmem_stride[2] = (uint64_t)depth * rows * cols * sizeof(dtype); + + smem_shape[0] = shared_tile_width; + smem_shape[1] = shared_tile_height; + smem_shape[2] = 1; + smem_shape[3] = 1; + } + + + assert((reinterpret_cast(global_addr) & 0b1111) == 0); + + assert(gmem_stride[0] % 16 == 0); + assert(gmem_stride[1] % 16 == 0); + assert(gmem_stride[2] % 16 == 0); + assert(gmem_stride[3] % 16 == 0); + + assert(smem_shape[0] <= 256); + assert(smem_shape[1] <= 256); + assert(smem_shape[2] <= 256); + + assert((smem_shape[0]*sizeof(dtype)) % 16 == 0); + + assert(smem_stride[0] <= 8); + assert(smem_stride[1] <= 8); + assert(smem_stride[2] <= 8); + assert(smem_stride[3] <= 8); + assert(smem_stride[4] <= 8); + + assert(smem_stride[0] == 1); + + if constexpr (tma_interleave == CU_TENSOR_MAP_INTERLEAVE_NONE && tma_swizzle != CU_TENSOR_MAP_SWIZZLE_NONE) { + assert(smem_shape[0] * sizeof(dtype) <= ST::swizzle_bytes); + } + + const uint64_t *gmem_shape_ptr = &gmem_shape[0]; + const uint64_t *gmem_stride_ptr = &gmem_stride[0]; + const uint32_t *smem_shape_ptr = &smem_shape[0]; + const uint32_t *smem_stride_ptr = &smem_stride[0]; + + CUresult result = cuTensorMapEncodeTiled( + tma_map, + tma_format, + tma_dim, + global_addr, + gmem_shape_ptr, + gmem_stride_ptr, + smem_shape_ptr, + smem_stride_ptr, + tma_interleave, + tma_swizzle, + tma_l2Promotion, + tma_oobFill); + + const char *error_string; + CUresult res = cuGetErrorString(result, &error_string); + if (result != CUDA_SUCCESS) { + std::string error_msg = format_tma_error( + "tile", error_string, + batch, depth, rows, cols, + tma_map, tma_format, tma_dim, global_addr, + gmem_shape_ptr, gmem_stride_ptr, + smem_shape_ptr, smem_stride_ptr, + 5, 4, 5, 5, + tma_interleave, tma_swizzle, tma_l2Promotion, tma_oobFill, + "ST::rows: " + std::to_string(ST::rows) + "\n ST::cols: " + std::to_string(ST::cols) + ); + throw std::runtime_error(error_msg); + } +} + + + + + + + + + + + + +template +__host__ static inline CUtensorMap* allocate_and_create_tensor_map(const typename ST::dtype *src, int batch, int depth, int rows, int cols) { + CUtensorMap *tma_map_d; + cudaMalloc(&tma_map_d, sizeof(CUtensorMap)); + CUtensorMap tma_map_host; + create_tensor_map(&tma_map_host, src, batch, depth, rows, cols); + cudaMemcpy(tma_map_d, &tma_map_host, sizeof(CUtensorMap), cudaMemcpyHostToDevice); + return tma_map_d; +} + + + + + + + + + + + + + + + +template +__host__ static inline void create_tensor_map(CUtensorMap *tma_map, const typename SV::dtype *src, int batch, int depth, int rows, int cols) { + using dtype = typename SV::dtype; + static_assert(axis == -1, "for vector TMA, row axis must be -1 as it's unused"); + static_assert(SV::length <= 256 || (SV::length*sizeof(dtype)) % 128 == 0); + + + + constexpr uint32_t tma_dim = 4; + void *global_addr = (void*)(src); + + constexpr CUtensorMapDataType tma_format = ( + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_BFLOAT16 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_FLOAT16 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_FLOAT32 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_INT32 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : +#ifdef KITTENS_SM10X + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : + std::is_same_v ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : +#endif + CUtensorMapDataType(-1) + ); + constexpr CUtensorMapInterleave tma_interleave = CU_TENSOR_MAP_INTERLEAVE_NONE; + constexpr CUtensorMapL2promotion tma_l2Promotion = CU_TENSOR_MAP_L2_PROMOTION_NONE; + constexpr CUtensorMapFloatOOBfill tma_oobFill = CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE; + constexpr CUtensorMapSwizzle swizzle = CU_TENSOR_MAP_SWIZZLE_NONE; + + constexpr uint64_t dim1 = sv_tma_dim1; + + + uint64_t gmem_shape [4] = {(uint64_t)cols, (uint64_t)rows, (uint64_t)depth, (uint64_t)batch}; + uint64_t gmem_stride[3] = {(uint64_t)cols*sizeof(dtype), (uint64_t)cols*rows*sizeof(dtype), (uint64_t)cols*rows*depth*sizeof(dtype)}; + uint32_t smem_shape [4] = {(uint32_t)dim1, 1, 1, 1}; + uint32_t smem_stride[4] = {1, 1, 1, 1}; + + + assert((reinterpret_cast(global_addr) & 0b1111) == 0); + + assert(smem_shape[0] <= 256); + + const uint64_t *gmem_shape_ptr = &gmem_shape[0]; + const uint64_t *gmem_stride_ptr = &gmem_stride[0]; + const uint32_t *smem_shape_ptr = &smem_shape[0]; + const uint32_t *smem_stride_ptr = &smem_stride[0]; + + CUresult result = cuTensorMapEncodeTiled( + tma_map, + tma_format, + tma_dim, + global_addr, + gmem_shape_ptr, + gmem_stride_ptr, + smem_shape_ptr, + smem_stride_ptr, + tma_interleave, + swizzle, + tma_l2Promotion, + tma_oobFill + ); + + const char *error_string; + CUresult res = cuGetErrorString(result, &error_string); + if (result != CUDA_SUCCESS) { + std::string error_msg = format_tma_error( + "vector", error_string, + batch, depth, rows, cols, + tma_map, tma_format, tma_dim, global_addr, + gmem_shape_ptr, gmem_stride_ptr, + smem_shape_ptr, smem_stride_ptr, + 4, 3, 4, 4, + tma_interleave, swizzle, tma_l2Promotion, tma_oobFill, + "SV::length: " + std::to_string(SV::length) + ); + throw std::runtime_error(error_msg); + } +}; + + + + + + + + + + + + + +template +__host__ static inline CUtensorMap* allocate_and_create_tensor_map(const typename SV::dtype *src, int batch, int depth, int rows, int cols) { + CUtensorMap *tma_map_d; + cudaMalloc(&tma_map_d, sizeof(CUtensorMap)); + CUtensorMap tma_map_host; + create_tensor_map(&tma_map_host, src, batch, depth, rows, cols); + cudaMemcpy(tma_map_d, &tma_map_host, sizeof(CUtensorMap), cudaMemcpyHostToDevice); + return tma_map_d; +} + +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/global/util.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/global/util.cuh new file mode 100644 index 0000000..7dd5236 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/global/util.cuh @@ -0,0 +1,180 @@ +#pragma once + +namespace kittens { +namespace ducks { +namespace gl { + +template concept cdim = (d > 0); +template concept rdim = (d == -1); +template struct compiled_dim { + static_assert(cdim<_v>, "Invalid compile-time dimension value"); + static constexpr size_t v = _v; + __host__ __device__ inline compiled_dim(const std::nullptr_t &_) {} + __host__ __device__ inline constexpr operator size_t() const { return v; } +}; +struct runtime_dim { + size_t v; + __host__ __device__ inline runtime_dim(const size_t &_v) : v(_v) {} + __host__ __device__ inline operator size_t() const { return v; } +}; +template using make_dim_t = std::conditional_t, runtime_dim, compiled_dim>; +template using make_arg_t = std::conditional_t, size_t, std::nullptr_t>; +} +} + +namespace detail { +template concept tile = ducks::st::all || ducks::rt::all || ducks::cst::all || ducks::crt::all; +template concept vec = ducks::sv::all || ducks::rv::all || ducks::csv::all || ducks::crv::all; + +#if (defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120)) && !defined(KITTENS_NO_HOST) +namespace tma { + +__host__ static inline std::string format_tma_error( + const char* error_type, + const char* error_string, + int batch, int depth, int rows, int cols, + CUtensorMap* tma_map, + CUtensorMapDataType tma_format, + uint32_t tma_dim, + void* global_addr, + const uint64_t* gmem_shape, + const uint64_t* gmem_stride, + const uint32_t* smem_shape, + const uint32_t* smem_stride, + size_t gmem_shape_size, + size_t gmem_stride_size, + size_t smem_shape_size, + size_t smem_stride_size, + CUtensorMapInterleave tma_interleave, + CUtensorMapSwizzle tma_swizzle, + CUtensorMapL2promotion tma_l2Promotion, + CUtensorMapFloatOOBfill tma_oobFill, + const std::string& extra_info = "" +) { + std::string msg; + msg += std::string("Error in TMA descriptor creation (") + + (error_type ? error_type : "Unknown") + "): " + + (error_string ? error_string : "Unknown CUDA error") + "\n"; + msg += "Parameters:\n"; + msg += " batch: " + std::to_string(batch) + "\n"; + msg += " depth: " + std::to_string(depth) + "\n"; + msg += " rows: " + std::to_string(rows) + "\n"; + msg += " cols: " + std::to_string(cols) + "\n"; + if (!extra_info.empty()) msg += " " + extra_info + "\n"; + msg += "cuTensorMapEncodeTiled arguments:\n"; + msg += " tma_map: " + std::to_string(reinterpret_cast(tma_map)) + "\n"; + msg += " tma_format: " + std::to_string(tma_format) + "\n"; + msg += " tma_dim: " + std::to_string(tma_dim) + "\n"; + msg += " global_addr: " + std::to_string(reinterpret_cast(global_addr)) + "\n"; + cudaPointerAttributes attributes; + cudaError_t err = cudaPointerGetAttributes(&attributes, global_addr); + msg += " global_addr memory type: "; + if (err == cudaSuccess) { + if (attributes.type == cudaMemoryTypeDevice) msg += "valid device memory\n"; + else if (attributes.type == cudaMemoryTypeHost) msg += "host memory (invalid for TMA)\n"; + else if (attributes.type == cudaMemoryTypeManaged) msg += "managed memory\n"; + else msg += "unknown memory type\n"; + } else { + msg += "unable to determine (error: " + std::string(cudaGetErrorString(err)) + ")\n"; + } + msg += " gmem_shape: " + std::to_string(reinterpret_cast(gmem_shape)) + " ["; + for (size_t i = 0; i < gmem_shape_size; ++i) { + msg += std::to_string(gmem_shape[i]); + if (i < gmem_shape_size - 1) msg += ", "; + } + msg += "]\n"; + msg += " gmem_stride: " + std::to_string(reinterpret_cast(gmem_stride)) + " ["; + for (size_t i = 0; i < gmem_stride_size; ++i) { + msg += std::to_string(gmem_stride[i]); + if (i < gmem_stride_size - 1) msg += ", "; + } + msg += "]\n"; + msg += " smem_shape: " + std::to_string(reinterpret_cast(smem_shape)) + " ["; + for (size_t i = 0; i < smem_shape_size; ++i) { + msg += std::to_string(smem_shape[i]); + if (i < smem_shape_size - 1) msg += ", "; + } + msg += "]\n"; + msg += " smem_stride: " + std::to_string(reinterpret_cast(smem_stride)) + " ["; + for (size_t i = 0; i < smem_stride_size; ++i) { + msg += std::to_string(smem_stride[i]); + if (i < smem_stride_size - 1) msg += ", "; + } + msg += "]\n"; + msg += " tma_interleave: " + std::to_string(tma_interleave) + "\n"; + msg += " tma_swizzle: " + std::to_string(tma_swizzle) + "\n"; + msg += " tma_l2Promotion: " + std::to_string(tma_l2Promotion) + "\n"; + msg += " tma_oobFill: " + std::to_string(tma_oobFill) + "\n"; + return msg; +} + +} +#endif +} + +namespace ducks { +namespace coord { +struct identifier {}; +} +} +template struct coord { + using identifier = ducks::coord::identifier; + using BASE = _T; + static_assert(std::is_same_v || detail::tile || detail::vec); + int b, d, r, c; + __device__ inline coord(int _b, int _d, int _r, int _c) : b(_b), d(_d), r(_r), c(_c) {} + __device__ inline coord( int _d, int _r, int _c) : b( 0), d(_d), r(_r), c(_c) {} + __device__ inline coord( int _r, int _c) : b( 0), d( 0), r(_r), c(_c) {} + __device__ inline coord( int _c) : b( 0), d( 0), r( 0), c(_c) {} + __device__ inline coord( ) : b( 0), d( 0), r( 0), c( 0) {} + template __device__ inline coord(const coord &other) : b(other.b), d(other.d), r(other.r), c(other.c) {} + __device__ inline coord(const int4 &other) : b(other.x), d(other.y), r(other.z), c(other.w) {} + __device__ inline operator int4() const { return int4(b, d, r, c); } + template __device__ inline coord unit_coord() const { + if constexpr (detail::tile) { + static_assert(row_axis != col_axis, "row and column axes must be different"); + static_assert(row_axis >= 0 && row_axis <= 3, "row axis must be between 0 and 3"); + static_assert(col_axis >= 0 && col_axis <= 3, "column axis must be between 0 and 3"); + static_assert(col_axis == 3, "for now, column axis must be 3"); + return coord( + row_axis == 0 ? b*BASE::rows : b, + row_axis == 1 ? d*BASE::rows : d, + row_axis == 2 ? r*BASE::rows : r, + c*BASE::cols + ); + } + else if constexpr (detail::vec) { + static_assert(row_axis == -1, "row axis must be be -1 for a vector coordinate to be converted to a unit coordinate"); + static_assert(col_axis >= 0 && col_axis <= 3, "column axis must be between 0 and 3"); + static_assert(col_axis == 3, "for now, column axis must be 3"); + return coord(b, d, r, c*BASE::length); + } + else { + return coord(*this); + } + } + template __device__ inline int dim() const { + static_assert(axis >= 0 && axis <= 3, "axis must be between 0 and 3"); + if constexpr (axis == 0) { return b; } + else if constexpr (axis == 1) { return d; } + else if constexpr (axis == 2) { return r; } + else { return c; } + } +}; +namespace ducks { +namespace coord { + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v; +template concept tile = all && (std::is_same_v || detail::tile); +template concept vec = all && (std::is_same_v || detail::vec); +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/crt.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/crt.cuh new file mode 100644 index 0000000..fdfa55f --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/crt.cuh @@ -0,0 +1,95 @@ + + + + + +#pragma once + +#include "rt.cuh" +#include "crv.cuh" + +namespace kittens { + +namespace ducks { +namespace crt { + + + + + + +struct identifier {}; + + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v && ducks::rt::all; + + + + + + +template +concept row_layout = all && std::is_same_v; + + + + + + + + +template +concept col_layout = all && std::is_same_v; +} +} + + + + + + + + + + + + + + +template +struct crt { + using identifier = ducks::crt::identifier; + using component = rt<_T, _rows, _cols, _layout>; + using layout = component::layout; + using T = component::T; + using T2 = component::T2; + using dtype = component::dtype; + + static constexpr int rows = component::rows; + static constexpr int cols = component::cols; + static constexpr int height = component::height; + static constexpr int width = component::width; + + + component real; + component imag; + + using row_vec = crv::row_vec_layout>; + using col_vec = crv::col_vec_layout>; +}; + +template using crt_fl = crt; +template using crt_bf = crt; +template using crt_hf = crt; + + + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/crv.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/crv.cuh new file mode 100644 index 0000000..e3432b0 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/crv.cuh @@ -0,0 +1,85 @@ + + + + + +#pragma once + +#include "../../common/common.cuh" +#include "rv.cuh" + +namespace kittens { + + + + +namespace ducks { + + + + + +namespace crv { + + + + + + +struct identifier {}; + + + + + + + +template +concept all = requires { + typename T::identifier; +} && std::is_same_v; + +template concept naive_layout = all && std::is_same_v; +template concept align_layout = all && std::is_same_v; +template concept ortho_layout = all && std::is_same_v; +template concept tile_layout = align_layout || ortho_layout; +} +} + + + + + + + + + + + + + + +template +struct crv { + using identifier = ducks::crv::identifier; + using component = rv<_T, _length, _layout>; + using layout = component::layout; + + using T = component::T; + using T2 = component::T2; + using dtype = component::dtype; + + static constexpr int length = component::length; + static constexpr int tiles = component::tiles; + + + component real; + component imag; +}; + + +template using crv_fl = crv; +template using crv_bf = crv; +template using crv_hf = crv; + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/register.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/register.cuh new file mode 100644 index 0000000..35e7b5f --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/register.cuh @@ -0,0 +1,14 @@ + + + + + +#pragma once + +#include "rv_layout.cuh" +#include "rt_base.cuh" +#include "rv.cuh" +#include "rt.cuh" + +#include "crv.cuh" +#include "crt.cuh" diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/rt.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/rt.cuh new file mode 100644 index 0000000..7ba92fe --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/rt.cuh @@ -0,0 +1,274 @@ + + + + + +#pragma once + +#include "../../common/common.cuh" + +#include "rt_layout.cuh" +#include "rt_base.cuh" +#include "rv.cuh" + +namespace kittens { + + + + +namespace ducks { + + + + + +namespace rt { + + + + + + +struct identifier {}; + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v; + + + + + + + + +template +concept row_layout = all && std::is_same_v; + + + + + + + + +template +concept col_layout = all && std::is_same_v; +} +} + + + + + + + + + + + + + + + + +template +struct rt { + static_assert(!std::is_same_v<_T, char>, "Bare 'char' is not permitted because its signedness is ISA-dependent."); + using identifier = ducks::rt::identifier; + using layout = _layout; + static_assert(kittens::ducks::base_types::T1<_T>); + using T = kittens::base_types::packing<_T>::unpacked_type; + using T2 = kittens::base_types::packing<_T>::packed_type; + using dtype = T2; + + static constexpr int rows = _rows; + static_assert(rows % rt_base::tile_size_row == 0, "Rows must be divisible by the tile size"); + static constexpr int cols = _cols; + static_assert(cols % rt_base::tile_size_col == 0, "Columns must be divisible by the tile size"); + static constexpr int height = rows / rt_base::tile_size_row; + static constexpr int width = cols / rt_base::tile_size_col; + static constexpr int tile_size_row = rt_base::tile_size_row; + static constexpr int tile_size_col = rt_base::tile_size_col; + static constexpr int num_elements = rt_base::num_elements * width * height; + static constexpr int elements_per_thread = rt_base::elements_per_thread * width * height; + static constexpr int packed_per_thread = rt_base::packed_per_thread * width * height; + static constexpr int packed_per_tile = rt_base::packed_per_thread; + + rt_base tiles[height][width]; + + using row_vec = rv::row_vec_layout>; + using col_vec = rv::col_vec_layout>; + + __device__ inline void operator=(const T &value) { + T2 value2 = base_types::packing::pack(value); + #pragma unroll + for(int i = 0; i < height; i++) { + #pragma unroll + for(int j = 0; j < width; j++) { + #pragma unroll + for(int k = 0; k < packed_per_tile; k++) { + tiles[i][j].data[k] = value2; + } + } + } + } + template + __device__ inline void operator=(const rt &other) { + using U2 = base_types::packing::packed_type; + #pragma unroll + for(int i = 0; i < height; i++) { + #pragma unroll + for(int j = 0; j < width; j++) { + #pragma unroll + for(int k = 0; k < packed_per_tile; k++) { + tiles[i][j].data[k] = base_types::convertor::convert(other.tiles[i][j].data[k]); + } + } + } + } +}; + + + + + + +template using rt_fl = rt; +template using rt_bf = rt; +template using rt_hf = rt; +template using rt_int8 = rt; +template using rt_uint8 = rt; +template using rt_int = rt; +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template using rt_fp8e4m3 = rt; +template using rt_fp8e5m2 = rt; +#endif +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template using rt_fp8e8m0 = rt; +template using rt_fp4e2m1_2 = rt; +#endif + + + + + + +template +__device__ constexpr const char* get_rt_type_name() { + if constexpr (std::is_same_v) { + return "rt_fl"; + } else if constexpr (std::is_same_v) { + return "rt_hf"; + } else if constexpr (std::is_same_v) { + return "rt_bf"; +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + return "rt_fp4_e2m1_2"; + } else if constexpr (std::is_same_v) { + return "rt_fp8_e8m0"; +#endif +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + return "rt_fp8_e4m3"; + } else if constexpr (std::is_same_v) { + return "rt_fp8_e5m2"; +#endif + } else { + return "rt_unknown"; + } +} + + + + + + + + + + +template +__device__ inline void print(const RT& tile) { + if (laneid() == 0) { + printf("Block %d, Warp %d: Register Tile %dx%d (Type: %s<%d,%d>) - Distributed View:\n", + blockIdx.x, threadIdx.x / WARP_THREADS, RT::rows, RT::cols, + get_rt_type_name(), RT::rows, RT::cols); + printf("Each thread holds %d elements (%d packed)\n", + RT::elements_per_thread, RT::packed_per_thread); + printf("\n"); + } + __syncwarp(); + + + for (int tid = 0; tid < WARP_THREADS; tid++) { + if (laneid() == tid) { + printf("Thread %2d: ", tid); + + + for (int i = 0; i < RT::height; i++) { + for (int j = 0; j < RT::width; j++) { + printf("Subtile[%d][%d]: ", i, j); + for (int k = 0; k < RT::packed_per_tile && k < 4; k++) { + auto packed_val = tile.tiles[i][j].data[k]; + + if constexpr (std::is_same_v) { + + if constexpr (std::is_same_v) { + printf("%.3f ", packed_val); + } else if constexpr (std::is_same_v) { + printf("%.3f ", __half2float(packed_val)); + } else if constexpr (std::is_same_v) { + printf("%.3f ", __bfloat162float(packed_val)); +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + printf("%.3f ", (float)packed_val); +#endif + } else { + printf("%.3f ", (float)packed_val); + } + } else { + + if constexpr (std::is_same_v) { + printf("[%.3f, %.3f] ", packed_val.x, packed_val.y); + } else if constexpr (std::is_same_v) { + + printf("[%.3f, %.3f] ", __bfloat162float(packed_val.x), __bfloat162float(packed_val.y)); +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + + __nv_fp8_e8m0 *vals = reinterpret_cast<__nv_fp8_e8m0*>(const_cast(&packed_val)); + printf("[%.3f,%.3f,%.3f,%.3f] ", + (float)vals[0], (float)vals[1], (float)vals[2], (float)vals[3]); + } else if constexpr (std::is_same_v) { + + uint8_t *vals = reinterpret_cast(const_cast(&packed_val)); + printf("[%.3f,%.3f,%.3f,%.3f] ", (float)fp4e2m1(vals[0] & 0xF), (float)fp4e2m1((vals[0] >> 4) & 0xF), (float)fp4e2m1(vals[1] & 0xF), (float)fp4e2m1((vals[1] >> 4) & 0xF)); +#endif + } else { + + printf("0x%x ", *(uint32_t*)&packed_val); + } + } + } + if (RT::packed_per_tile > 4) printf("... "); + } + } + printf("\n"); + } + __syncwarp(); + } + + if (laneid() == 0) { + printf("\n"); + } + __syncwarp(); +} + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/rt_base.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/rt_base.cuh new file mode 100644 index 0000000..4a36591 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/rt_base.cuh @@ -0,0 +1,124 @@ + + + + + +#pragma once + +#include "../../common/common.cuh" +#include "rt_layout.cuh" +#include "rv_layout.cuh" + +namespace kittens { + + + +namespace ducks { + + + + + +namespace rt_base { + + + + + + +struct identifier {}; +} +} + + + + + + + + + + + + +template struct rt_base { + using identifier = ducks::rt_base::identifier; + using layout = _layout; + static_assert(kittens::ducks::base_types::T1<_T>); + using T = kittens::base_types::packing<_T>::unpacked_type; + using T2 = kittens::base_types::packing<_T>::packed_type; + using dtype = T2; + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert( + std::is_same_v || std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + || std::is_same_v || std::is_same_v +#endif + , + "rt_base was provided an unsupported type." + ); +#else + static_assert( + std::is_same_v || std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v || std::is_same_v, + "rt_base was provided an unsupported type." + ); +#endif + + static constexpr int tile_size_row = kittens::TILE_ROW_DIM; + static constexpr int tile_size_col = kittens::TILE_COL_DIM; + static constexpr int rows = tile_size_row; + static constexpr int cols = tile_size_col; + static constexpr int num_elements = rows*cols; + static constexpr int elements_per_thread = num_elements / 32; + + static constexpr int packed_per_thread = (elements_per_thread / base_types::packing::num()) ; + static constexpr int registers_per_thread = packed_per_thread * sizeof(dtype) / 4; + + using row_vec_layout = std::conditional_t, ducks::rv_layout::align, ducks::rv_layout::ortho>; + using col_vec_layout = std::conditional_t, ducks::rv_layout::ortho, ducks::rv_layout::align>; + + dtype data[packed_per_thread]; +}; + + + + + + + +namespace ducks { +namespace rt_base { + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v; +} +} + + + +template using rt_base_fl = rt_base; +template using rt_base_bf = rt_base; +template using rt_base_hf = rt_base; +template using rt_base_int8 = rt_base; +template using rt_base_uint8 = rt_base; +template using rt_base_int = rt_base; +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template using rt_base_fp8e4m3 = rt_base; +template using rt_base_fp8e5m2 = rt_base; +#endif +#ifdef KITTENS_SM10X +template using rt_base_fp8e8m0 = rt_base; +template using rt_base_fp4e2m1_2 = rt_base; +#endif +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/rt_layout.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/rt_layout.cuh new file mode 100644 index 0000000..251b09a --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/rt_layout.cuh @@ -0,0 +1,40 @@ + + + + + +#pragma once + +namespace kittens { +namespace ducks { + + + + + +namespace rt_layout { + + + + +struct row {}; + + + +struct col {}; + + + + +template +concept all = std::is_same_v || std::is_same_v; + + + + +template struct transpose { using type = col; }; +template<> struct transpose { using type = row; }; + +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/rv.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/rv.cuh new file mode 100644 index 0000000..c384d7b --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/rv.cuh @@ -0,0 +1,230 @@ + + + + + +#pragma once + +#include "../../common/common.cuh" +#include "rv_layout.cuh" + +namespace kittens { + + + + +namespace ducks { + + + + + +namespace rv { + + + + + + +struct identifier {}; + + + + + + + +template +concept all = requires { + typename T::identifier; +} && std::is_same_v; + +template concept naive_layout = all && std::is_same_v; +template concept align_layout = all && std::is_same_v; +template concept ortho_layout = all && std::is_same_v; +template concept tile_layout = align_layout || ortho_layout; +} +} + + + + + + + + + + + + + +template +struct rv { + static_assert(!std::is_same_v<_T, char>, "Bare 'char' is not permitted because its signedness is ISA-dependent."); + using identifier = ducks::rv::identifier; + static_assert(kittens::ducks::base_types::T1<_T>); + using layout = _layout; + static constexpr bool is_naive = std::is_same_v; + using T = kittens::base_types::packing<_T>::unpacked_type; + using T2 = kittens::base_types::packing<_T>::packed_type; + using dtype = std::conditional_t; + + static constexpr int length = _length; + static_assert(length % kittens::TILE_ROW_DIM == 0, "Length must be divisible by the tile dimension"); + static constexpr int tiles = _length / kittens::TILE_ROW_DIM; + static constexpr int inner_dim = layout::inner_dim; + static constexpr int outer_dim = is_naive ? (tiles+1)/2 : tiles; + #if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert(!std::is_same_v && !std::is_same_v, "Unsupported type for fp8"); + #endif + + dtype data[outer_dim][inner_dim]; + + __device__ inline dtype* operator[](size_t idx) { return &data[idx][0]; } + __device__ inline const dtype* operator[](size_t idx) const { return &data[idx][0]; } + __device__ inline dtype& operator[](int2 outin) { return data[outin.x][outin.y]; } + __device__ inline const dtype& operator[](int2 outin) const { return data[outin.x][outin.y]; } + + __device__ inline void operator=(const T &value) { + dtype value2; + if constexpr(is_naive) { + value2 = value; + } else { + value2 = base_types::packing::pack(value); + } + #pragma unroll + for(int i = 0; i < outer_dim; i++) { + #pragma unroll + for(int j = 0; j < inner_dim; j++) { + data[i][j] = value2; + } + } + } + template + __device__ inline void operator=(const rv &other) { + using U2 = base_types::packing::packed_type; + #pragma unroll + for(int i = 0; i < outer_dim; i++) { + #pragma unroll + for(int j = 0; j < inner_dim; j++) { + data[i][j] = base_types::convertor::convert(other.data[i][j]); + } + } + } +}; + +template using rv_fl = rv; +template using rv_bf = rv; +template using rv_hf = rv; + + + + + + + + + + + +template +__device__ void print(const RV &vec) { + if (laneid() == 0) { + printf("Block %d, Warp %d: Register Vector %d (Type: %s, Layout: %s) - Distributed View:\n", + blockIdx.x, threadIdx.x / WARP_THREADS, RV::length, + std::is_same_v ? "float" : + std::is_same_v ? "bf16" : + std::is_same_v ? "half" : +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + std::is_same_v ? "fp8e8m0" : +#endif +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + std::is_same_v ? "fp8e4m3" : + std::is_same_v ? "fp8e5m2" : "unknown", +#endif + RV::is_naive ? "naive" : "tile"); + printf("Each thread holds %dx%d elements\n", RV::outer_dim, RV::inner_dim); + printf("\n"); + } + __syncwarp(); + + + for (int tid = 0; tid < WARP_THREADS; tid++) { + if (laneid() == tid) { + printf("Thread %2d: ", tid); + + + for (int i = 0; i < RV::outer_dim; i++) { + printf("Outer[%d]: ", i); + for (int j = 0; j < RV::inner_dim; j++) { + auto value = vec.data[i][j]; + + if constexpr (std::is_same_v) { + + if constexpr (std::is_same_v) { + printf("%.3f ", value); + } else if constexpr (std::is_same_v) { + printf("%.3f ", __half2float(value)); + } else if constexpr (std::is_same_v) { + printf("%.3f ", __bfloat162float(value)); +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + printf("%.3f ", (float)value); +#endif +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + printf("%.3f ", (float)value); + } else if constexpr (std::is_same_v) { + printf("%.3f ", (float)value); +#endif + } else { + printf("%.3f ", (float)value); + } + } else { + + if constexpr (std::is_same_v) { + printf("[%.3f, %.3f] ", value.x, value.y); + } else if constexpr (std::is_same_v) { + + printf("[%.3f, %.3f] ", __bfloat162float(value.x), __bfloat162float(value.y)); + } else if constexpr (std::is_same_v) { + + printf("[%.3f, %.3f] ", __half2float(value.x), __half2float(value.y)); +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + + __nv_fp8_e8m0 *vals = reinterpret_cast<__nv_fp8_e8m0*>(const_cast(&value)); + printf("[%.3f,%.3f,%.3f,%.3f] ", + (float)vals[0], (float)vals[1], (float)vals[2], (float)vals[3]); +#endif +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + + __nv_fp8_e4m3 *vals = reinterpret_cast<__nv_fp8_e4m3*>(const_cast(&value)); + printf("[%.3f,%.3f,%.3f,%.3f] ", + (float)vals[0], (float)vals[1], (float)vals[2], (float)vals[3]); + } else if constexpr (std::is_same_v) { + + __nv_fp8_e5m2 *vals = reinterpret_cast<__nv_fp8_e5m2*>(const_cast(&value)); + printf("[%.3f,%.3f,%.3f,%.3f] ", + (float)vals[0], (float)vals[1], (float)vals[2], (float)vals[3]); +#endif + } else { + + printf("0x%x ", *(uint32_t*)&value); + } + } + } + printf(" "); + } + printf("\n"); + } + __syncwarp(); + } + + if (laneid() == 0) { + printf("\n"); + } +} + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/rv_layout.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/rv_layout.cuh new file mode 100644 index 0000000..75ac192 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/rv_layout.cuh @@ -0,0 +1,38 @@ + + + + + +#pragma once + +namespace kittens { +namespace ducks { + + + + + +namespace rv_layout { + + + + +struct align { constexpr static int inner_dim = 2; }; + + + +struct ortho { constexpr static int inner_dim = 1; }; + + + +struct naive { constexpr static int inner_dim = 1; }; + + + + +template +concept all = std::is_same_v || std::is_same_v || std::is_same_v; + +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/cst.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/cst.cuh new file mode 100644 index 0000000..b4f867f --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/cst.cuh @@ -0,0 +1,78 @@ + + + + + +#pragma once + +#include "st.cuh" +#include "csv.cuh" + +namespace kittens { + +namespace ducks { +namespace cst { + + + + + + +struct identifier {}; + + + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v && ducks::st::all; + +} +} + + + + + + + + + + + + + +template +struct cst { + using identifier = ducks::cst::identifier; + using component = st<_T, _rows, _cols>; + using T = component::T; + using T2 = component::T2; + using dtype = component::dtype; + + static constexpr int rows = component::rows; + static constexpr int cols = component::cols; + + + + + component real; + component imag; + + + using col_vec = csv; + using row_vec = csv; +}; + + + +template using cst_bf = cst; +template using cst_hf = cst; +template using cst_fl = cst; + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/csv.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/csv.cuh new file mode 100644 index 0000000..a961774 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/csv.cuh @@ -0,0 +1,72 @@ + + + + + +#pragma once + +#include "sv.cuh" + +namespace kittens { + +namespace ducks { +namespace csv { + + + + + + +struct identifier {}; + + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v && ducks::sv::all; + +} +} + + + + + + + + + + + + +template +struct csv { + using identifier = ducks::csv::identifier; + using component = sv<_T, _length>; + using T = component::T; + using T2 = component::T2; + using dtype = component::dtype; + + static constexpr int length = component::length; + static constexpr int tiles = component::tiles; + + + + + component real; + component imag; +}; + + + + +template using csv_bf = csv; +template using csv_hf = csv; +template using csv_fl = csv; + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/descriptor.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/descriptor.cuh new file mode 100644 index 0000000..d5b955f --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/descriptor.cuh @@ -0,0 +1,142 @@ + + + + + +#pragma once + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + +#include "../../common/common.cuh" +#include "st.cuh" +#include "cst.cuh" + +namespace kittens { +namespace ducks { +namespace st_descriptor { +struct identifier {}; +} +} + +namespace detail { + + +__device__ static inline uint64_t matrix_descriptor_encode(uint64_t x) { return (((x) & 0x3FFFF) >> 0x4); } + +template +__device__ static inline uint64_t matrix_descriptor_raw( + T *addr, + uint32_t leading_dim_offset, + uint32_t stride_dim_offset, + uint32_t swizzle_mode +) { +#ifdef KITTENS_SM10X + + return matrix_descriptor_encode(reinterpret_cast(addr)) | + (1llu << 46) | +#else + return matrix_descriptor_encode(reinterpret_cast(addr)) | +#endif + matrix_descriptor_encode((uint64_t)leading_dim_offset) << 16 | + matrix_descriptor_encode((uint64_t)stride_dim_offset) << 32 | + (uint64_t)swizzle_mode << 62; +} + +} + +template +struct st_descriptor { + using identifier = ducks::st_descriptor::identifier; + using ST = _ST; + using T = typename ST::T; + static constexpr int rows = ST::rows; + static constexpr int cols = ST::cols; + static constexpr bool swizzle = ST::swizzle; + static_assert(swizzle, "Non-swizzled descriptor is not supported yet."); + uint64_t base_desc; + __device__ inline st_descriptor(const ST &tile) { + + if constexpr (MN_major) { + if constexpr (ST::swizzle_bytes == 128) + base_desc = detail::matrix_descriptor_raw(&tile.data[0], 2048*ST::rows/TILE_ROW_DIM, 1024, 1); + else if constexpr (ST::swizzle_bytes == 64) + base_desc = detail::matrix_descriptor_raw(&tile.data[0], 1024*ST::rows/TILE_ROW_DIM, 512, 2); + else + base_desc = detail::matrix_descriptor_raw(&tile.data[0], 512*ST::rows/TILE_ROW_DIM, 256, 3); + } + else { + if constexpr (ST::swizzle_bytes == 128) + base_desc = detail::matrix_descriptor_raw(&tile.data[0], 16 , 1024, 1); + else if constexpr (ST::swizzle_bytes == 64) + base_desc = detail::matrix_descriptor_raw(&tile.data[0], 16 , 512, 2); + else + base_desc = detail::matrix_descriptor_raw(&tile.data[0], 16 , 256, 3); + } + } + __device__ inline st_descriptor(const st_descriptor &other) : base_desc(other.base_desc) {} + __device__ inline uint64_t chunk_descriptor(int chunk_idx) { + + + + + + + if constexpr (MN_major) { + if constexpr (ST::swizzle_bytes == 128) { + return base_desc + detail::matrix_descriptor_encode(chunk_idx*2048); + } + else if constexpr (ST::swizzle_bytes == 64) { + return base_desc + detail::matrix_descriptor_encode(chunk_idx*1024); + } + else { + return base_desc + detail::matrix_descriptor_encode(chunk_idx*512); + } + } + else { + if constexpr (ST::swizzle_bytes == 128) { + + return base_desc + detail::matrix_descriptor_encode((chunk_idx%4)*32 + (chunk_idx/4)*(ST::rows/TILE_ROW_DIM)*2048); + } + else if constexpr (ST::swizzle_bytes == 64) { + + return base_desc + detail::matrix_descriptor_encode((chunk_idx%2)*32 + (chunk_idx/2)*(ST::rows/TILE_ROW_DIM)*1024); + } + else { + + return base_desc + detail::matrix_descriptor_encode(chunk_idx*(ST::rows/TILE_ROW_DIM)*512); + } + } + } +#ifdef KITTENS_SM103 + __device__ inline uint64_t chunk_descriptor_k96(int chunk_idx) { + static_assert(!MN_major, "K96 chunk descriptors are only supported for K-major tcgen05 operands."); + static_assert(std::is_same_v, "K96 chunk descriptors are only supported for packed FP4."); + static_assert(ST::swizzle_bytes == 128, "K96 chunk descriptors require 128B swizzle mode."); + const int start_byte = chunk_idx*48; + const int edge_gap = 128 - start_byte%128; + const int cont_byte = start_byte + (edge_gap < 48 ? edge_gap : 32); + + const uint64_t start_desc = base_desc + detail::matrix_descriptor_encode(start_byte%128 + (start_byte/128)*(ST::rows/TILE_ROW_DIM)*2048); + const uint64_t cont_desc = base_desc + detail::matrix_descriptor_encode(cont_byte%128 + (cont_byte/128)*(ST::rows/TILE_ROW_DIM)*2048); + return (start_desc & ~(0x3FFFull << 16)) | ((cont_desc & 0x3FFFull) << 16) | (1ull << 52); + } +#endif +}; + +namespace ducks { +namespace st_descriptor { + +template concept input = ducks::st::all || (requires {typename T::identifier;} && std::is_same_v); +template concept complex_input = ducks::cst::all; +namespace detail { +template struct st_getter { using type = typename T::ST; }; +template struct st_getter { using type = T; }; +template struct st_getter { using type = T::component; }; +template using get_st = typename st_getter::type; +} +} +} + +} + +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/shared.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/shared.cuh new file mode 100644 index 0000000..3755973 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/shared.cuh @@ -0,0 +1,14 @@ + + + + + +#pragma once + +#include "sv.cuh" +#include "st.cuh" + +#include "csv.cuh" +#include "cst.cuh" + +#include "descriptor.cuh" diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/st.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/st.cuh new file mode 100644 index 0000000..438299e --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/st.cuh @@ -0,0 +1,658 @@ + + + + + +#pragma once + +#include "../../common/common.cuh" +#include "sv.cuh" + + + + +namespace kittens { +namespace ducks { + + + + + +namespace st { + + + + + + + +struct identifier {}; + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v; +} +} + + +template< + typename ST, + int _subtile_height, + int _subtile_width +> +struct st_subtile; + + + + + + + + +template +struct KITTENS_DEFAULT_ALIGN st { + static_assert(!std::is_same_v<_T, char>, "Bare 'char' is not permitted because its signedness is ISA-dependent."); +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static_assert(!std::is_same_v<_T, fp4e2m1>, "For FP4 types, you must use a packed type (i.e., fp4e2m1_2 or fp4e2m1_4)."); +#endif + using identifier = ducks::st::identifier; + using T = base_types::packing<_T>::unpacked_type; + using T2 = base_types::packing<_T>::packed_type; + using dtype = T; + + static constexpr bool swizzle = _swizzle; + + + static constexpr int underlying_rows = _rows; + static constexpr int underlying_cols = _cols; + static constexpr int underlying_num_elements = underlying_rows * underlying_cols; + + static constexpr int rows = _rows; + static constexpr int cols = _cols; + static constexpr int num_elements = rows * cols; + + static_assert(!swizzle || (rows % kittens::TILE_ROW_DIM == 0), "Rows must be divisible by the tile dimension"); + static_assert((swizzle && (cols % kittens::TILE_COL_DIM == 0)) || (!swizzle && (cols % kittens::BASE_TILE_DIM == 0)), "Cols must be divisible by the tile dimension"); + +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + static_assert(base_types::packing::num() == 1 || std::is_same_v); +#else + static_assert(base_types::packing::num() == 1); +#endif + + + static_assert(_swizzle_bytes == 0 || _swizzle_bytes == 32 || _swizzle_bytes == 64 || _swizzle_bytes == 128); + static constexpr int swizzle_bytes = _swizzle_bytes > 0 ? _swizzle_bytes : ( + sizeof(dtype) == 1 ? ( + (cols/kittens::TILE_COL_DIM)%4 == 0 ? 128 : + (cols/kittens::TILE_COL_DIM)%2 == 0 ? 64 : 32 + ) : + sizeof(dtype) == 2 ? ( + (cols/kittens::TILE_COL_DIM)%4 == 0 ? 128 : + (cols/kittens::TILE_COL_DIM)%2 == 0 ? 64 : 32 + ) : + sizeof(dtype) == 4 ? ( + (cols/kittens::TILE_COL_DIM)%2 == 0 ? 128 : 64 + ) : -1 + ); + + dtype data[rows*cols]; + + __device__ static inline T* idx(T *ptr, int2 coord) { + int r = coord.x, c = coord.y; + if constexpr (swizzle) { + static constexpr int swizzle_repeat = swizzle_bytes * 8; + static constexpr int subtile_cols = swizzle_bytes / sizeof(T); + const int outer_idx = c/subtile_cols; + const uint64_t addr = (uint64_t)(&ptr[outer_idx*rows*subtile_cols + r*subtile_cols + c%subtile_cols]); + const int swizzle = ((addr % swizzle_repeat) >> 7) << 4; + return (T*)(addr ^ swizzle); + } else { + return &ptr[r*cols + c]; + } + } + __device__ static inline uint32_t idx(uint32_t ptr, int2 coord) { + int r = coord.x, c = coord.y; + if constexpr (swizzle) { + static constexpr int swizzle_repeat = swizzle_bytes * 8; + static constexpr int subtile_cols = swizzle_bytes / sizeof(T); + const int outer_idx = c/subtile_cols; + const uint32_t addr = ptr + sizeof(T)*(outer_idx*rows*subtile_cols + r*subtile_cols + c%subtile_cols); + const int swizzle = ((addr % swizzle_repeat) >> 7) << 4; + return (addr ^ swizzle); + } else { + return ptr + sizeof(T)*(r*cols + c); + } + } + + + + + + + __device__ inline dtype& operator[](const int2 &rowcol) { + return *idx(data, rowcol); + } + __device__ inline const dtype& operator[](const int2 &rowcol) const { + return *(const dtype*)idx((dtype*)data, rowcol); + } + __device__ inline dtype& operator[](int idx) { + return data[idx]; + } + __device__ inline const dtype& operator[](int idx) const { + return data[idx]; + } + + template + __device__ inline st_subtile, subtile_rows, subtile_cols> subtile(int2 rowcol); + + + + + + template + __device__ inline st<_T, _rows, subtile_cols, _swizzle, swizzle_bytes > &subtile(int idx) { + static_assert(swizzle_bytes > 0, "Parent shared tile must have an explicit swizzle_bytes."); + constexpr int swizzle_elements = swizzle_bytes / sizeof(T); + static_assert(subtile_cols >= 0 && subtile_cols % swizzle_elements == 0); + return *reinterpret_cast *>( + &data[rows*swizzle_elements*(subtile_cols/swizzle_elements)*idx] + ); + } + + + using col_vec = sv; + using row_vec = sv; +}; + + + + + + + + + + + + + +template< + typename _ST, + int _subtile_rows, + int _subtile_cols +> +struct st_subtile { + using identifier = ducks::st::identifier; + using ST = _ST; + using T = ST::T; + using T2 = ST::T2; + using dtype = T; + + static constexpr bool swizzle = ST::swizzle; + + static constexpr int underlying_rows = ST::underlying_rows; + static_assert(underlying_rows % kittens::TILE_ROW_DIM == 0, "Underlying rows must be divisible by the tile dimension"); + static constexpr int underlying_cols = ST::underlying_cols; + static_assert(underlying_cols % kittens::TILE_COL_DIM == 0, "Underlying cols must be divisible by the tile dimension"); + static constexpr int underlying_num_elements = ST::underlying_num_elements; + + static constexpr int rows = _subtile_rows; + static_assert(rows % kittens::TILE_ROW_DIM == 0, "Rows must be divisible by the tile dimension"); + static constexpr int cols = _subtile_cols; + static_assert(cols % kittens::TILE_COL_DIM == 0, "Cols must be divisible by the tile dimension"); + static constexpr int num_elements = rows * cols; + + static constexpr int swizzle_bytes = ST::swizzle_bytes; + + dtype *data; + int row_offset, col_offset; + + __device__ st_subtile(ST &src, int2 rowcol) { + data = &src.data[0]; + row_offset = rowcol.x * rows; + col_offset = rowcol.y * cols; + } + + __device__ inline T* idx(T *ptr, const int2 coord) const { + int r = coord.x+row_offset, c = coord.y+col_offset; + if constexpr (swizzle) { + static constexpr int swizzle_repeat = swizzle_bytes * 8; + static constexpr int subtile_cols = swizzle_bytes / sizeof(T); + const int outer_idx = c/subtile_cols; + const uint64_t addr = (uint64_t)(&ptr[outer_idx*underlying_rows*subtile_cols + r*subtile_cols + c%subtile_cols]); + const int swizzle = ((addr % swizzle_repeat) >> 7) << 4; + return (T*)(addr ^ swizzle); + } else { + return &ptr[r*cols + c]; + } + } + __device__ inline uint32_t idx(uint32_t ptr, const int2 coord) const { + int r = coord.x+row_offset, c = coord.y+col_offset; + if constexpr(swizzle) { + static constexpr int swizzle_repeat = swizzle_bytes * 8; + static constexpr int subtile_cols = swizzle_bytes / sizeof(T); + const int outer_idx = c/subtile_cols; + const uint32_t addr = ptr + sizeof(T)*(outer_idx*underlying_rows*subtile_cols + r*subtile_cols + c%subtile_cols); + const int swizzle = ((addr % swizzle_repeat) >> 7) << 4; + return (addr ^ swizzle); + } else { + return ptr + sizeof(T)*(r*cols + c); + } + } + + + + + + + __device__ inline dtype& operator[](const int2 &rowcol) { + return *idx(data, rowcol); + } + __device__ inline const dtype& operator[](const int2 &rowcol) const { + return *(const dtype*)idx((dtype*)data, rowcol); + } + + + + + + using col_vec = sv; + using row_vec = sv; + + __device__ inline void operator=(const dtype &value) { + #pragma unroll + for(int i = kittens::laneid(); i < num_elements; i += WARP_THREADS) { + data[i] = value; + } + } +}; + +template +template +__device__ inline st_subtile, subtile_rows, subtile_cols> +st<_T, _rows, _cols, _swizzle, _swizzle_bytes>::subtile(int2 rowcol) +{ + + using ST_t = st<_T, _rows, _cols>; + using dtype = typename ST_t::dtype; + + + static_assert(subtile_rows > 0 && subtile_cols > 0, "Subtile dimensions must be positive."); + static_assert(subtile_rows % kittens::TILE_ROW_DIM == 0, + "Subtile rows must be divisible by the base tile row dimension."); + static_assert(subtile_cols % kittens::TILE_COL_DIM == 0, + "Subtile cols must be divisible by the base tile col dimension."); + + + static_assert(ST_t::rows % subtile_rows == 0, + "Parent tile rows must be divisible by subtile rows."); + static_assert(ST_t::cols % subtile_cols == 0, + "Parent tile cols must be divisible by subtile cols."); + + + + static_assert(ST_t::rows == ST_t::underlying_rows && ST_t::cols == ST_t::underlying_cols, + "Cannot create a subtile from an object that appears to be a subtile view (rows/cols mismatch underlying)."); + + + + + return st_subtile(*this, rowcol); +} + + + +template +using st_bf = st; +template +using st_hf = st; +template +using st_fl = st; +template +using st_int8 = st; +template +using st_uint8 = st; +template +using st_int = st; +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template +using st_fp8e4m3 = st; +template +using st_fp8e5m2 = st; +#endif +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template +using st_fp8e8m0 = st; +template +using st_fp4e2m1_2 = st; +#endif + + + + + + +template +__device__ constexpr const char* get_tile_type_name() { + if constexpr (std::is_same_v) { + return "st_fl"; + } else if constexpr (std::is_same_v) { + return "st_hf"; + } else if constexpr (std::is_same_v) { + return "st_bf"; +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + return "st_fp4_e2m1_2"; + } else if constexpr (std::is_same_v) { + return "st_fp8_e8m0"; +#endif +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + return "st_fp8_e4m3"; + } else if constexpr (std::is_same_v) { + return "st_fp8_e5m2"; +#endif + } else { + return "st_unknown"; + } +} + +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + + + + + + + + +template +__device__ inline void print_fp4(const ST& tile) { + if (std::is_same_v) { + + constexpr int cols = ST::cols * 2; + printf("Block %d: Shared Tile %dx%d (Type: %s<%d,%d>):\n", blockIdx.x, ST::rows, cols, get_tile_type_name(), ST::rows, cols); + + + printf(" "); + for (int c = 0; c < cols; c++) { + printf("%8d ", c); + } + printf("\n"); + + + printf(" "); + for (int c = 0; c < cols; c++) { + printf("--------+"); + } + printf("\n"); + + + for (int r = 0; r < ST::rows; r++) { + printf("%3d |", r); + for (int c = 0; c < cols; c += 2) { + uint8_t *vals = reinterpret_cast(const_cast(&tile[{r,c/2}])); + + + float f1 = static_cast(fp4e2m1(vals[0] & 0xF)); + float f2 = static_cast(fp4e2m1((vals[0] >> 4) & 0xF)); + + printf("%8.3f %8.3f ", f1, f2); + + } + printf("\n"); + } + printf("\n"); + } else { + printf("Type must be FP4 in this function\n"); + } +} +#endif + + + + + + + + + +template +__device__ inline void print(const ST& tile) { + printf("Block %d: Shared Tile %dx%d (Type: %s<%d,%d>):\n", blockIdx.x, ST::rows, ST::cols, get_tile_type_name(), ST::rows, ST::cols); + + + printf(" "); + for (int c = 0; c < ST::cols; c++) { + printf("%8d ", c); + } + printf("\n"); + + + printf(" "); + for (int c = 0; c < ST::cols; c++) { + printf("--------+"); + } + printf("\n"); + + + for (int r = 0; r < ST::rows; r++) { + printf("%3d |", r); + for (int c = 0; c < ST::cols; c++) { + if constexpr (std::is_same_v) { + printf("%8.3f ", tile[{r,c}]); + } else if constexpr (std::is_same_v) { + printf("%8.3f ", __bfloat162float(tile[{r,c}])); + } else if constexpr (std::is_integral_v) { + printf("%8d ", (int)tile[{r,c}]); +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + printf("%8.3f ", static_cast(tile[{r,c}])); + } else if constexpr (std::is_same_v) { + printf("%8.3f ", static_cast(tile[{r,c}])); +#endif +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + printf("%8.3f ", static_cast(tile[{r,c}])); +#endif + } else { + printf("%8.3f ", (float)tile[{r,c}]); + } + } + printf("\n"); + } + printf("\n"); +} + + + + + + + + + +template +__device__ inline void print(const st_subtile& subtile) { + printf("Block %d: Shared Subtile %dx%d (offset: [%d,%d], Type: %s<%d,%d> from %s<%d,%d>):\n", + blockIdx.x, subtile.rows, subtile.cols, + subtile.row_offset, subtile.col_offset, + get_tile_type_name(), subtile.rows, subtile.cols, + get_tile_type_name(), ST::rows, ST::cols); + + + printf(" "); + for (int c = 0; c < subtile.cols; c++) { + printf("%8d ", c); + } + printf("\n"); + + + printf(" "); + for (int c = 0; c < subtile.cols; c++) { + printf("--------+"); + } + printf("\n"); + + + for (int r = 0; r < subtile.rows; r++) { + printf("%3d |", r); + for (int c = 0; c < subtile.cols; c++) { + if constexpr (std::is_same_v) { + printf("%8.3f ", subtile[{r,c}]); + } else if constexpr (std::is_same_v) { + printf("%8.3f ", __bfloat162float(subtile[{r,c}])); + } else if constexpr (std::is_integral_v) { + printf("%8d ", (int)subtile[{r,c}]); + } else { + printf("%8.3f ", (float)subtile[{r,c}]); + } + } + printf("\n"); + } + printf("\n"); +} + + + + + + + + + +template +__device__ inline void fill_value(ST& tile, float value) { + printf("Filling Tile %dx%d with %f:\n", ST::rows, ST::cols, value); + + + for (int r = 0; r < ST::rows; r++) { + for (int c = 0; c < ST::cols; c++) { + if constexpr (std::is_same_v) { + tile[{r,c}] = value; + } else if constexpr (std::is_same_v) { + tile[{r,c}] = __float2bfloat16(value); +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + tile.data[r*ST::cols + c] = fp4e2m1(value); + } else if constexpr (std::is_same_v) { + tile.data[r*ST::cols + c] = fp8e8m0(value); +#endif + } else { + tile[{r,c}] = value; + } + } + } +} + + + + + + + + + +template +__device__ inline void fill_identity(ST& tile) { + printf("Filling Tile %dx%d with identity:\n", ST::rows, ST::cols); + + + for (int r = 0; r < ST::rows; r++) { + for (int c = 0; c < ST::cols; c++) { + if constexpr (std::is_same_v) { + + tile[{r,c}] = 1.0f; + } else if constexpr (std::is_same_v) { + + tile[{r,c}] = __float2bfloat16(1.0f); +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + if(r == c){ + + tile.data[r*ST::cols + c] = std::bit_cast(uint8_t(0xFF)); + } else { + tile.data[r*ST::cols + c] = std::bit_cast(uint8_t(0x00)); + } +#endif + } else { + + } + } + } +} + + + + + + + + + + +template +__device__ inline void print_bits(const ST& tile, bool unswizzle = false) { + printf("Block %d: Shared Tile %dx%d (Type: %s<%d,%d>):\n", blockIdx.x, ST::rows, ST::cols, get_tile_type_name(), ST::rows, ST::cols); + + + printf(" "); + for (int c = 0; c < ST::cols; c++) { + printf("%11d ", c); + } + printf("\n"); + + + printf(" "); + for (int c = 0; c < ST::cols; c++) { + printf("-----------+"); + } + printf("\n"); + + + for (int r = 0; r < ST::rows; r++) { + printf("%3d |", r); + for (int c = 0; c < ST::cols; c++) { + if constexpr (std::is_same_v) { + printf("%8.3f ", tile[{r,c}]); + } else if constexpr (std::is_same_v) { + printf("%8.3f ", __bfloat162float(tile[{r,c}])); + + +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v || std::is_same_v) { + + + uint8_t bits; + if(unswizzle){ + bits = *reinterpret_cast(&tile.data[r*ST::cols + c]); + } else { + bits = *reinterpret_cast(&tile[{r,c}]); + } + + printf("0b"); + + for (int bit = 7; bit >= 0; bit--) { + printf("%d", (bits >> bit) & 0x1); + if (bit % 4 == 0 && bit > 0) printf("_"); + } + printf(" "); +#endif + } else { + printf("%8.3f ", (float)tile[{r,c}]); + } + } + printf("\n"); + } + printf("\n"); +} + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/sv.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/sv.cuh new file mode 100644 index 0000000..1bf6120 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/sv.cuh @@ -0,0 +1,155 @@ + + + + + +#pragma once + +#include "../../common/common.cuh" + +namespace kittens { + + + +namespace ducks { + + + + + +namespace sv { + + + + + + +struct identifier {}; + + + + + + + +template +concept all = requires { + typename T::identifier; +} && std::is_same_v; +} +} + + + + + + + + + + + +template +struct KITTENS_DEFAULT_ALIGN sv { + static_assert(!std::is_same_v<_T, char>, "Bare 'char' is not permitted because its signedness is ISA-dependent."); + using identifier = ducks::sv::identifier; + using T = base_types::packing<_T>::unpacked_type; + using T2 = base_types::packing<_T>::packed_type; + using dtype = T; + + static constexpr int length = _length; + static_assert(length % TILE_ROW_DIM == 0, "Length must be divisible by the tile dimension"); + static constexpr int tiles = length / TILE_ROW_DIM; + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + static constexpr int num_alloc_elements = ((length * sizeof(dtype) + 127) / 128) * (128 / sizeof(dtype)); +#else + static constexpr int num_alloc_elements = length; +#endif + dtype data[num_alloc_elements]; + + __device__ static inline T* idx(T *ptr, int idx) { + return ptr[idx]; + } + + __device__ inline dtype& operator[](size_t idx) { return data[idx]; } + __device__ inline const dtype& operator[](size_t idx) const { return data[idx]; } + + template __device__ inline sv<_T, sub_length> &subvec(int idx) { + return *(sv*)&data[idx * sub_length]; + } + template __device__ inline const sv<_T, sub_length> &subvec(int idx) const { + return *(sv*)&data[idx * sub_length]; + } + + __device__ inline void operator=(const dtype &value) { + #pragma unroll + for(int i = kittens::laneid(); i < length; i += WARP_THREADS) { + data[i] = value; + } + } +}; + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +namespace detail { +namespace tma { + + + +template struct find_vector_divider { + static constexpr int value = (SV::length % (16*D) == 0 && (SV::length < 256 || ((16*D)*sizeof(typename SV::dtype)) % 128 == 0)) ? + 16*D : find_vector_divider::value; +}; +template struct find_vector_divider { static constexpr int value = 16; }; +template constexpr int sv_tma_dim1 = find_vector_divider::value; +template constexpr int sv_tma_dim2 = (SV::length / sv_tma_dim1); +} +} +#endif + + + + +template using sv_bf = sv; +template using sv_hf = sv; +template using sv_fl = sv; +template using sv_int8 = sv; +template using sv_uint8 = sv; +template using sv_int = sv; +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template using sv_fp8e4m3 = sv; +template using sv_fp8e5m2 = sv; +#endif +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) +template using sv_fp8e8m0 = sv; +template using sv_fp4e2m1_2 = sv; +#endif + + + +template +__device__ inline void print(const SV& sv) { + printf("Shared Vector %d:\n", SV::length); + for(int i = 0; i < SV::length; i++) { + if constexpr (std::is_same_v) { + printf("%f ", __bfloat162float(sv[i])); + } else if constexpr (std::is_same_v) { + printf("%f ", __half2float(sv[i])); + } else if constexpr (std::is_same_v) { + printf("%f ", sv[i]); +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + printf("%f ", static_cast(sv[i])); +#endif +#if defined(KITTENS_SM10X) || defined(KITTENS_SM120) + } else if constexpr (std::is_same_v) { + printf("%f ", static_cast(sv[i])); +#endif + } else { + printf("%d ", (int)(sv[i])); + } + } + printf("\n"); +} + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/system/ipc.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/system/ipc.cuh new file mode 100644 index 0000000..91463fc --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/system/ipc.cuh @@ -0,0 +1,143 @@ +#pragma once + +#include "../../common/common.cuh" +#include "vmm.cuh" + +namespace kittens { +namespace ducks { +namespace ipc { +namespace handle { + +struct identifier {}; + +template concept all = requires { + typename T::identifier; +} && std::is_same_v; + +} +} +} + +namespace detail { +namespace ipc { + +enum flavor { + LEGACY = 0, + VMM = 1 +}; + +template +struct handle; + +template<> +struct handle { + using identifier = ducks::ipc::handle::identifier; + static constexpr flavor flavor_ = flavor::LEGACY; + cudaIpcMemHandle_t handle_ {}; +}; + +template<> +struct handle { + using identifier = ducks::ipc::handle::identifier; + static constexpr flavor flavor_ = flavor::VMM; + int handle_; +}; + +__host__ inline static void check_support(const int device_id) { + CUdevice device; + CUCHECK(cuDeviceGet(&device, device_id)); + + int ipc_supported = 0; + CUDACHECK(cudaDeviceGetAttribute(&ipc_supported, cudaDevAttrIpcEventSupport, device_id)); + int ipc_handle_supported = 0; + CUCHECK(cuDeviceGetAttribute(&ipc_handle_supported, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR_SUPPORTED, device)); + + if (!ipc_supported || !ipc_handle_supported) + throw std::runtime_error("CUDA IPC is not supported on this device"); +} + +template +__host__ inline static void export_handle( + IPC_HANDLE *ipc_handle, + void *ptr +) { + if constexpr (IPC_HANDLE::flavor_ == flavor::LEGACY) { + CUDACHECK(cudaIpcGetMemHandle(&ipc_handle->handle_, ptr)); + } else if constexpr (IPC_HANDLE::flavor_ == flavor::VMM) { + CUmemGenericAllocationHandle memory_handle; + detail::vmm::vm_retrieve_handle(&memory_handle, ptr); + + CUCHECK(cuMemExportToShareableHandle(&ipc_handle->handle_, memory_handle, detail::vmm::HANDLE_TYPE, 0)); + detail::vmm::vm_free(memory_handle); + } else { + throw std::runtime_error("Invalid IPC handle type"); + } +} + +template +__host__ inline static void export_handle( + IPC_HANDLE *ipc_handle, + CUmemGenericAllocationHandle &memory_handle +) { + if constexpr (IPC_HANDLE::flavor_ == flavor::VMM) { + CUCHECK(cuMemExportToShareableHandle(&ipc_handle->handle_, memory_handle, detail::vmm::HANDLE_TYPE, 0)); + } else { + throw std::runtime_error("Invalid IPC handle type"); + } +} + +template +__host__ inline static void import_handle ( + void **ptr, + IPC_HANDLE &ipc_handle, + const size_t size, + int local_world_size +) { + if constexpr (IPC_HANDLE::flavor_ == flavor::LEGACY) { + CUDACHECK(cudaIpcOpenMemHandle(ptr, ipc_handle.handle_, cudaIpcMemLazyEnablePeerAccess)); + } else if constexpr (IPC_HANDLE::flavor_ == flavor::VMM) { + CUmemGenericAllocationHandle memory_handle; + CUCHECK(cuMemImportFromShareableHandle(&memory_handle, reinterpret_cast(static_cast(ipc_handle.handle_)), detail::vmm::HANDLE_TYPE)); + detail::vmm::vm_map(ptr, memory_handle, size); + detail::vmm::vm_set_access(*ptr, size, local_world_size); + detail::vmm::vm_free(memory_handle); + close(ipc_handle.handle_); + ipc_handle.handle_ = -1; + } else { + throw std::runtime_error("Invalid IPC handle type"); + } +} + +template +__host__ inline static void import_handle ( + CUmemGenericAllocationHandle *memory_handle, + IPC_HANDLE &ipc_handle, + const size_t size, + int local_world_size +) { + if constexpr (IPC_HANDLE::flavor_ == flavor::VMM) { + CUCHECK(cuMemImportFromShareableHandle(memory_handle, reinterpret_cast(static_cast(ipc_handle.handle_)), detail::vmm::HANDLE_TYPE)); + close(ipc_handle.handle_); + ipc_handle.handle_ = -1; + } else { + throw std::runtime_error("Invalid IPC handle type"); + } +} + +template +__host__ inline static void free_handle( + void *ptr, + const size_t size +) { + if constexpr (_flavor == flavor::LEGACY) { + CUDACHECK(cudaIpcCloseMemHandle(ptr)); + } else if constexpr (_flavor == flavor::VMM) { + detail::vmm::vm_unmap(ptr, size); + } else { + throw std::runtime_error("Invalid IPC handle type"); + } +} + +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/system/pgl.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/system/pgl.cuh new file mode 100644 index 0000000..d3a0492 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/system/pgl.cuh @@ -0,0 +1,182 @@ + + + + + +#pragma once + +#include "../../common/common.cuh" +#include "../shared/shared.cuh" +#include "../global/global.cuh" + +namespace kittens { + + + +namespace ducks { +namespace pgl { + +struct identifier {}; + + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v; + +} +} + + + + + + + + + +template +struct pgl { + using identifier = ducks::pgl::identifier; + using GL = _GL; + using T = GL::dtype; + static_assert(!std::is_same_v, "Bare 'char' is not permitted because its signedness is ISA-dependent."); + using dtype = T; + + static constexpr int num_devices = NUM_DEVICES; + static constexpr bool multicast = MULTICAST; + + T *mc_ptr; + GL gls[NUM_DEVICES]; + + detail::descriptor_dict tma_descs; + + __host__ __device__ const GL &operator[](int idx) const { return gls[idx]; } + __device__ inline T* mc_ptr_at(const coord &idx) const { + static_assert(MULTICAST, "Multicast is not enabled for this PGL."); + const GL &gl = gls[0]; + return &mc_ptr[((idx.b * static_cast(gl.depth()) + idx.d) * gl.rows() + idx.r) * gl.cols() + idx.c]; + } + + __host__ __device__ pgl() = delete; +#ifndef KITTENS_NO_HOST + __host__ inline pgl(T **_data, + ducks::gl::make_arg_t _batch, + ducks::gl::make_arg_t _depth, + ducks::gl::make_arg_t _rows, + ducks::gl::make_arg_t _cols) : + pgl(std::make_index_sequence{}, _data, _batch, _depth, _rows, _cols) { } + __host__ inline pgl(T *_mc_ptr, + T **_data, + ducks::gl::make_arg_t _batch, + ducks::gl::make_arg_t _depth, + ducks::gl::make_arg_t _rows, + ducks::gl::make_arg_t _cols) : + pgl(std::make_index_sequence{}, _mc_ptr, _data, _batch, _depth, _rows, _cols) { } + template + __host__ inline pgl(std::index_sequence, + T **_data, + ducks::gl::make_arg_t _batch, + ducks::gl::make_arg_t _depth, + ducks::gl::make_arg_t _rows, + ducks::gl::make_arg_t _cols) : + mc_ptr(nullptr), gls{GL(_data[I], _batch, _depth, _rows, _cols)...} { + static_assert(!MULTICAST, "Multicast pointer not passed to multicast-enabled PGL."); + } + template + __host__ inline pgl(std::index_sequence, + T *_mc_ptr, + T **_data, + ducks::gl::make_arg_t _batch, + ducks::gl::make_arg_t _depth, + ducks::gl::make_arg_t _rows, + ducks::gl::make_arg_t _cols) : + mc_ptr(_mc_ptr), gls{GL(_data[I], _batch, _depth, _rows, _cols)...} { + static_assert(MULTICAST, "Multicast pointer passed to multicast-disabled PGL."); + tma_descs = detail::descriptor_dict( + mc_ptr, gls[0].batch_internal, gls[0].depth_internal, gls[0].rows_internal, gls[0].cols_internal); + } +#endif + + template + __device__ inline const CUtensorMap* get_tma() const { + return tma_descs.template get(); + } + + __host__ __device__ inline auto batch() const { return gls[0].batch(); } + __host__ __device__ inline auto depth() const { return gls[0].depth(); } + __host__ __device__ inline auto rows() const { return gls[0].rows(); } + __host__ __device__ inline auto cols() const { return gls[0].cols(); } + __host__ __device__ inline auto numel() const { return gls[0].numel(); } + + template __device__ inline size_t shape() const { return gls[0].template shape(); } + template __device__ inline size_t stride() const { return gls[0].template stride(); } +}; + +#ifndef KITTENS_NO_HOST + +template __host__ inline PGL make_pgl( + uint64_t *data, int b, int d, int r, int c +) { + if constexpr (safe) { + if (PGL::GL::__b__ > 0 && b != PGL::GL::__b__) { + throw std::runtime_error("Batch dimension mismatch. Expected: " + std::to_string(PGL::GL::__b__) + ", Got: " + std::to_string(b)); + } + if (PGL::GL::__d__ > 0 && d != PGL::GL::__d__) { + throw std::runtime_error("Depth dimension mismatch. Expected: " + std::to_string(PGL::GL::__d__) + ", Got: " + std::to_string(d)); + } + if (PGL::GL::__r__ > 0 && r != PGL::GL::__r__) { + throw std::runtime_error("Row dimension mismatch. Expected: " + std::to_string(PGL::GL::__r__) + ", Got: " + std::to_string(r)); + } + if (PGL::GL::__c__ > 0 && c != PGL::GL::__c__) { + throw std::runtime_error("Column dimension mismatch. Expected: " + std::to_string(PGL::GL::__c__) + ", Got: " + std::to_string(c)); + } + } + return PGL( + reinterpret_cast(data), + make_unsafe_gl_arg(b), + make_unsafe_gl_arg(d), + make_unsafe_gl_arg(r), + make_unsafe_gl_arg(c) + ); +} + +template __host__ inline PGL make_pgl( + uint64_t mc_ptr, uint64_t *data, int b, int d, int r, int c +) { + if constexpr (safe) { + if (PGL::GL::__b__ > 0 && b != PGL::GL::__b__) { + throw std::runtime_error("Batch dimension mismatch. Expected: " + std::to_string(PGL::GL::__b__) + ", Got: " + std::to_string(b)); + } + if (PGL::GL::__d__ > 0 && d != PGL::GL::__d__) { + throw std::runtime_error("Depth dimension mismatch. Expected: " + std::to_string(PGL::GL::__d__) + ", Got: " + std::to_string(d)); + } + if (PGL::GL::__r__ > 0 && r != PGL::GL::__r__) { + throw std::runtime_error("Row dimension mismatch. Expected: " + std::to_string(PGL::GL::__r__) + ", Got: " + std::to_string(r)); + } + if (PGL::GL::__c__ > 0 && c != PGL::GL::__c__) { + throw std::runtime_error("Column dimension mismatch. Expected: " + std::to_string(PGL::GL::__c__) + ", Got: " + std::to_string(c)); + } + } + return PGL( + reinterpret_cast(mc_ptr), + reinterpret_cast(data), + make_unsafe_gl_arg(b), + make_unsafe_gl_arg(d), + make_unsafe_gl_arg(r), + make_unsafe_gl_arg(c) + ); +} + +#endif + + +template +using barrier_t = pgl, NUM_DEVICES, true>; + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/system/system.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/system/system.cuh new file mode 100644 index 0000000..067ca81 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/system/system.cuh @@ -0,0 +1,14 @@ + + + + + +#pragma once + +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +#ifndef KITTENS_NO_HOST +#include "ipc.cuh" +#include "vmm.cuh" +#endif +#include "pgl.cuh" +#endif diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/system/vmm.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/system/vmm.cuh new file mode 100644 index 0000000..b11cd85 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/system/vmm.cuh @@ -0,0 +1,176 @@ +#pragma once + +#include "../../common/common.cuh" + +namespace kittens { +namespace detail { +namespace vmm { + + + +static constexpr CUmemAllocationHandleType HANDLE_TYPE = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR; + +typedef CUmemGenericAllocationHandle handle; + +__host__ inline static void vm_alloc( + CUmemGenericAllocationHandle *handle, + size_t *allocated_size, + const size_t size, + const int device_id +) { + CUmemAllocationProp prop = {}; + prop.location.id = device_id; + prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + prop.requestedHandleTypes = HANDLE_TYPE; + prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; + + size_t granularity; + CUCHECK(cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED)); + *allocated_size = (size + granularity - 1) / granularity * granularity; + + CUCHECK(cuMemCreate(handle, *allocated_size, &prop, 0)); +} + +__host__ inline static void vm_map( + void **ptr, + const CUmemGenericAllocationHandle &handle, + const size_t size +) { + CUdeviceptr device_ptr; + CUCHECK(cuMemAddressReserve(&device_ptr, size, 0, 0, 0)); + CUCHECK(cuMemMap(device_ptr, size, 0, handle, 0)); + *ptr = (void *)device_ptr; +} + +__host__ inline static void vm_set_access( + void *ptr, + const size_t size, + const int num_devices +) { + std::vector descs(num_devices); + for (int i = 0; i < num_devices; i++) { + descs[i].location.id = i; + descs[i].location.type = CU_MEM_LOCATION_TYPE_DEVICE; + descs[i].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; + } + CUCHECK(cuMemSetAccess(reinterpret_cast(ptr), size, descs.data(), num_devices)); +} + +__host__ inline static void vm_retrieve_handle( + CUmemGenericAllocationHandle *handle, + void *ptr +) { + + CUCHECK(cuMemRetainAllocationHandle(handle, ptr)); +} + +__host__ inline static void vm_unmap( + void *ptr, + const size_t size +) { + CUCHECK(cuMemUnmap(reinterpret_cast(ptr), size)); + CUCHECK(cuMemAddressFree(reinterpret_cast(ptr), size)); +} + +__host__ inline static void vm_free(CUmemGenericAllocationHandle &handle) { + + + CUCHECK(cuMemRelease(handle)); +} + +__host__ inline static void vm_alloc_map_set_access( + void **ptr, + size_t *allocated_size, + const size_t size, + const int device_id, + const int num_devices +) { + CUmemGenericAllocationHandle handle; + vm_alloc(&handle, allocated_size, size, device_id); + vm_map(ptr, handle, *allocated_size); + vm_set_access(*ptr, *allocated_size, num_devices); + vm_free(handle); +} + +__host__ inline static void multicast_check(const int device_id) { + CUdevice device; + CUCHECK(cuDeviceGet(&device, device_id)); + + int multicast_supported; + CUresult result = cuDeviceGetAttribute( + &multicast_supported, + CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, + device + ); + + if (!multicast_supported) + throw std::runtime_error("Device does not support multicast"); +} + +__host__ inline static void multicast_create_handle( + CUmemGenericAllocationHandle *handle, + size_t *allocated_size, + const size_t size, + const int num_devices +) { + if (num_devices <= 1) + throw std::runtime_error("Multicast requires at least 2 devices"); + + CUmulticastObjectProp prop = {}; + prop.numDevices = num_devices; + prop.handleTypes = HANDLE_TYPE; + + size_t granularity; + CUCHECK(cuMulticastGetGranularity(&granularity, &prop, CU_MULTICAST_GRANULARITY_RECOMMENDED)); + *allocated_size = (size + granularity - 1) / granularity * granularity; + prop.size = *allocated_size; + + + cuMulticastCreate(handle, &prop); +} + +__host__ inline static void multicast_bind_device( + const CUmemGenericAllocationHandle &handle, + const int device_id +) { + + CUdevice device; + CUCHECK(cuDeviceGet(&device, device_id)); + CUCHECK(cuMulticastAddDevice(handle, device)); +} + +__host__ inline static void multicast_bind_memory( + const CUmemGenericAllocationHandle &multicast_handle, + const CUmemGenericAllocationHandle &memory_handle, + const size_t size +) { + + CUCHECK(cuMulticastBindMem(multicast_handle, 0, memory_handle, 0, size, 0)); +} + +__host__ inline static void multicast_bind_address( + const CUmemGenericAllocationHandle &multicast_handle, + void *ptr, + const size_t size +) { + + CUmemGenericAllocationHandle memory_handle; + vm_retrieve_handle(&memory_handle, ptr); + multicast_bind_memory(multicast_handle, memory_handle, size); + vm_free(memory_handle); +} + +__host__ inline static void multicast_unbind_device( + const CUmemGenericAllocationHandle &handle, + const size_t size, + const int device_id +) { + + CUdevice device; + CUCHECK(cuDeviceGet(&device, device_id)); + CUCHECK(cuMulticastUnbind(handle, device, 0, size)); +} + +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/tensor/tensor.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/tensor/tensor.cuh new file mode 100644 index 0000000..052967c --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/tensor/tensor.cuh @@ -0,0 +1,152 @@ + + + + + +#pragma once + +#include "tt.cuh" + + +namespace kittens { +namespace ducks { + + + + + +namespace tensor_allocator { + + + +struct identifier {}; + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v; +} +} + +template struct tensor_allocator { + static_assert(_nblocks_per_sm == 1 || _nblocks_per_sm == 2, "nblocks_per_sm must be 1 or 2"); + static_assert(_ncta == 1 || _ncta == 2, "ncta must be 1 or 2"); + + using identifier = ducks::tensor_allocator::identifier; + + static constexpr int nblocks_per_sm = _nblocks_per_sm; + static constexpr int cols =((MAX_TENSOR_COLS/nblocks_per_sm) / 32) * 32; + static_assert(cols>0 && cols%32==0, "cols must be a multiple of 32"); + static constexpr int ncta = _ncta; + static constexpr bool managed = _managed; + + uint32_t addr; + + template __device__ inline void check_bounds() { + static_assert(col_offset >= 0 && col_offset + TT::cols <= cols, "Tile allocation extends out of bounds of the tensor allocator!"); + } + + __device__ inline void set_addr(uint32_t _addr) { + addr = _addr; + } + + __device__ inline uint32_t get_addr(int superlane, int col_offset) const { + return addr + ((superlane*16) << 16) + col_offset; + } + + __device__ inline uint32_t get_addr(int col_offset) const { + return addr + col_offset; + } + + __device__ inline void provision(uint32_t &shared_addr) { + + + + + if constexpr (ncta == 1) { + asm volatile( + "tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;\n" + :: "l"(reinterpret_cast(&shared_addr)), "n"(cols) + ); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;\n"); + } + else { + asm volatile( + "tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;\n" + :: "l"(reinterpret_cast(&shared_addr)), "n"(cols) + ); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;\n"); + } + } + + __device__ inline tensor_allocator() { + if constexpr (managed) { + __shared__ uint32_t shared_addr; + if (warpid() == 0) provision(shared_addr); + asm volatile("tcgen05.fence::before_thread_sync;\n"); + asm volatile("bar.sync 0;\n"); + asm volatile("tcgen05.fence::after_thread_sync;\n"); + set_addr(shared_addr); + } + } + + template __device__ inline auto allocate(int superlane, int col_offset) { +#ifndef NDEBUG + int allocate_cols = std::is_same_v ? TT::cols/4 : TT::cols; + if(col_offset + allocate_cols > cols) { + if(laneid() == 0) printf("Tile allocation extends out of bounds of the tensor allocator! col_offset: %d, TT::cols: %d, allocator cols: %d\n", col_offset, TT::cols, cols); + asm volatile("trap;"); + } + if(superlane < 0 || superlane > 1) { + printf("Superlane must be 0 or 1! superlane: %d\n", superlane); + asm volatile("trap;"); + } +#endif + return TT(get_addr(superlane, col_offset)); + } + + template __device__ inline auto allocate(int col_offset) { +#ifndef NDEBUG + int allocate_cols = std::is_same_v ? TT::cols/4 : TT::cols; + if(col_offset + allocate_cols > cols) { + if(laneid() == 0) printf("Tile allocation extends out of bounds of the tensor allocator! col_offset: %d, TT::cols: %d, allocator cols: %d\n", col_offset, TT::cols, cols); + asm volatile("trap;"); + } +#endif + return TT(get_addr(0, col_offset)); + } + + __device__ inline void deprovision() { + + + + if constexpr (ncta == 1) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;\n" + :: "r"(addr), "n"(cols) + ); + } else { + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;\n" + :: "r"(addr), "n"(cols) + ); + } + } + + __device__ inline ~tensor_allocator() { + if constexpr (managed) { + if (warpid() == 0) { + if (ncta == 2) { + asm volatile ("barrier.cluster.arrive.release.aligned;\n"); + asm volatile ("barrier.cluster.wait.acquire.aligned;\n"); + } + deprovision(); + } + } + } +}; + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/tensor/tt.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/tensor/tt.cuh new file mode 100644 index 0000000..7d37e42 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/tensor/tt.cuh @@ -0,0 +1,135 @@ + + + + + +#pragma once + +#include "../../common/common.cuh" + + + + +namespace kittens { + +constexpr int MAX_TENSOR_ROWS = 128; +constexpr int MAX_TENSOR_COLS = 512; + +namespace ducks { + + + + + +namespace tt { + + + +struct identifier {}; + + + + + + + +template concept all = requires { + typename T::identifier; +} && std::is_same_v; +template concept half = all && T::rows == MAX_TENSOR_ROWS / 2; +template concept full = all && T::rows == MAX_TENSOR_ROWS; +} +} + + + + + + + + +template +struct tt { + static_assert(!std::is_same_v<_T, char>, "Bare 'char' is not permitted because its signedness is ISA-dependent."); + using identifier = ducks::tt::identifier; + using T = base_types::packing<_T>::unpacked_type; + using T2 = base_types::packing<_T>::packed_type; + using dtype = T; + + static constexpr int rows = _rows; + static constexpr int cols = _cols; + + static_assert(rows / (4 / sizeof(T)) <= MAX_TENSOR_ROWS, "Row dimension must be less than or equal to MAX_TENSOR_ROWS"); + static_assert(cols / (4 / sizeof(T)) <= MAX_TENSOR_COLS, "Column dimension must be less than or equal to MAX_TENSOR_COLS"); + static_assert(rows % kittens::BASE_TILE_DIM == 0, "Row dimension must be divisible by the 16"); + static_assert(cols % kittens::BASE_TILE_DIM == 0, "Column dimension must be divisible by the 16"); + + uint32_t addr; + + __device__ inline tt() : addr(0) {} + __device__ inline tt(uint32_t addr) : addr(addr) {} + + template __device__ inline TT subtile(int row_offset, int col_offset) const { +#ifndef NDEBUG + if(row_offset < 0 || row_offset+TT::rows > rows || col_offset < 0 || col_offset+TT::cols > cols) { + printf("Subtile out of bounds! full tile rows: %d, full tile cols: %d, subtile rows: %d, subtile cols: %d, row_offset: %d, col_offset: %d\n", rows, cols, TT::rows, TT::cols, row_offset, col_offset); + asm volatile("trap;"); + } +#endif + return TT(addr + (row_offset<<16) + col_offset/(4/(uint32_t)sizeof(T))); + } + template __device__ inline TT subtile(int col_offset) const { + return TT(addr + col_offset/(4/(uint32_t)sizeof(T))); + } + template __device__ inline uint32_t chunk_addr(int chunk) const { + if constexpr (transpose) { + static_assert(!transpose, "Transposed TMEM chunk addressing is not supported due to Blackwell tensor-core limitations."); + } + else { + if constexpr (std::is_same_v || std::is_same_v) { + return addr + (16 * chunk / (4/(uint32_t)sizeof(T))); + } + else if constexpr (std::is_same_v || std::is_same_v || std::is_same_v || std::is_same_v) { + return addr + (32 * chunk / (4/(uint32_t)sizeof(T))); + } + else { + static_assert(sizeof(T) == 999, "Currently unsupported type for input to an mma."); + } + } + } +}; + + + +template using tt_bf = tt; +template using tt_hf = tt; +template using tt_fl = tt; +template using tt_int8 = tt; +template using tt_uint8 = tt; +template using tt_int = tt; +template using tt_fp8e4m3 = tt; +template using tt_fp8e5m2 = tt; +template using tt_fp8e8m0 = tt; +template using tt_fp4e2m1_2 = tt; +template using half_tt_bf = tt; +template using half_tt_hf = tt; +template using half_tt_fl = tt; +template using half_tt_int8 = tt; +template using half_tt_uint8 = tt; +template using half_tt_int = tt; +template using half_tt_fp8e4m3 = tt; +template using half_tt_fp8e5m2 = tt; +template using half_tt_fp8e8m0 = tt; +template using half_tt_fp4e2m1_2 = tt; +template using full_tt_bf = tt; +template using full_tt_hf = tt; +template using full_tt_fl = tt; +template using full_tt_int8 = tt; +template using full_tt_uint8 = tt; +template using full_tt_int = tt; +template using full_tt_fp8e4m3 = tt; +template using full_tt_fp8e5m2 = tt; +template using full_tt_fp8e8m0 = tt; +template using full_tt_fp4e2m1_2 = tt; + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/types.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/types.cuh new file mode 100644 index 0000000..d665854 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/types.cuh @@ -0,0 +1,67 @@ + + + + + +#pragma once + +#include "register/register.cuh" +#include "shared/shared.cuh" +#include "global/global.cuh" +#if defined(KITTENS_SM90) || defined(KITTENS_SM10X) || defined(KITTENS_SM120) +#include "system/system.cuh" +#endif +#ifdef KITTENS_SM10X +#include "tensor/tensor.cuh" +#endif + + + +namespace kittens { + + + + + + + + + + + + + + + +template +using row_vec = T::row_vec; + + + + + + + + + + + + + + + +template +using col_vec = T::col_vec; + + + + +using row_l = ducks::rt_layout::row; +using col_l = ducks::rt_layout::col; + + +using align_l = ducks::rv_layout::align; +using ortho_l = ducks::rv_layout::ortho; +using naive_l = ducks::rv_layout::naive; + +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/common/common.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/common/common.cuh new file mode 100644 index 0000000..d036194 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/common/common.cuh @@ -0,0 +1,4 @@ +#pragma once + +#include "util.cuh" +#include "templates.cuh" diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/common/templates.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/common/templates.cuh new file mode 100644 index 0000000..bc237a3 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/common/templates.cuh @@ -0,0 +1,110 @@ +#pragma once + +namespace kittens { +namespace prototype { + +struct empty { + __device__ empty& operator*() { return *this; } + __device__ const empty& operator*() const { return *this; } +}; + + +template struct __align__(ALIGN) padder { + T val; + __device__ padder(T _val) : val(_val) {} + __device__ T& operator*() { return val; } + __device__ const T& operator*() const { return val; } +}; + + +template concept kittens_layout = requires { + typename T::globals; + typename T::input_block; +}; +template concept store_kittens_layout = kittens_layout && requires { + typename T::output_block; +}; + + +namespace detail { + + +#define FLAG_GETTER(flag_name, default_value) \ +template concept has_##flag_name##_flag = requires { T::flag_name; }; \ +template constexpr int flag_name##_v = default_value; \ +template constexpr int flag_name##_v = T::flag_name; + + +FLAG_GETTER(FORCE_ALIGN, 1024) + +FLAG_GETTER(DEBUG, 0) + +FLAG_GETTER(NUM_BLOCKS, 1) + +FLAG_GETTER(NUM_CONSUMER_WARPS, 8) + +FLAG_GETTER(NUM_PRODUCER_WARPS, 4) + +FLAG_GETTER(MAX_SHARED_MEMORY, ((kittens::MAX_SHARED_MEMORY/NUM_BLOCKS_v)-512)) + +FLAG_GETTER(INPUT_PIPE_STAGES, 1) + +FLAG_GETTER(OUTPUT_PIPE_STAGES, 1) + +FLAG_GETTER(CONSUMER_BARRIER_ARRIVALS, NUM_CONSUMER_WARPS_v) + +FLAG_GETTER(PRODUCER_BARRIER_ARRIVALS, NUM_PRODUCER_WARPS_v) + +FLAG_GETTER(CLUSTER_BLOCKS, 1) + +template constexpr int NUM_WARPS_v = NUM_CONSUMER_WARPS_v + NUM_PRODUCER_WARPS_v; +template constexpr int NUM_THREADS_v = NUM_WARPS_v * 32; +template constexpr int NUM_CONSUMER_WARPGROUPS_v = NUM_CONSUMER_WARPS_v / kittens::WARPGROUP_WARPS; + + +#define BLOCK_GETTER(block_name) \ +template concept has_##block_name##_block = kittens_layout && requires { typename T::block_name; }; \ +template struct block_name##_getter { using type = empty; }; \ +template struct block_name##_getter { using type = typename T::block_name; }; \ +template using block_name##_t = typename block_name##_getter::type; + + +BLOCK_GETTER(input_block) + +BLOCK_GETTER(output_block) + +BLOCK_GETTER(scratch_block) + +BLOCK_GETTER(finish_block) + +BLOCK_GETTER(common_state) + +BLOCK_GETTER(producer_state) + +BLOCK_GETTER(consumer_state) + +} + + +template struct complete_kittens_layout : T { + + using globals_t = T::globals; + + using input_block_t = typename detail::input_block_t; + using output_block_t = typename detail::output_block_t; + using scratch_block_t = typename detail::scratch_block_t; + using finish_block_t = typename detail::finish_block_t; + + using common_state_t = typename detail::common_state_t; + using producer_state_t = typename detail::producer_state_t; + using consumer_state_t = typename detail::consumer_state_t; + + + constexpr static int FORCE_ALIGN = detail::FORCE_ALIGN_v; + using input_alloc_block_t = padder; + using output_alloc_block_t = padder; + using scratch_alloc_block_t = padder; +}; + +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/common/util.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/common/util.cuh new file mode 100644 index 0000000..ee82de5 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/common/util.cuh @@ -0,0 +1,13 @@ +#pragma once + +#include + +namespace kittens { +namespace prototype { + +__device__ static inline int get_task_iter(int ti) { + return ti*gridDim.x*gridDim.y*gridDim.z + blockIdx.z*gridDim.y*gridDim.x + blockIdx.y*gridDim.x+blockIdx.x; +} + +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/interpreter/interpreter.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/interpreter/interpreter.cuh new file mode 100644 index 0000000..7635035 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/interpreter/interpreter.cuh @@ -0,0 +1,440 @@ +#pragma once + +#include +#include "../common/common.cuh" +#include "templates.cuh" + +namespace kittens { +namespace prototype { +namespace interpreter { + +static constexpr int NUM_PRODUCER_WARPS = 4; +static constexpr int NUM_CONSUMER_WARPS = 8; +static constexpr int NUM_WARPS = NUM_PRODUCER_WARPS + NUM_CONSUMER_WARPS; +static constexpr int NUM_BLOCKS = 1; + +template concept kernel_template = requires { + typename T::layout; + typename T::producer; + typename T::consumer; + T::common_setup; + T::producer::setup; + T::producer::load; + T::consumer::setup; + T::consumer::compute; + T::consumer::finish; +} && kittens_layout; +template concept has_store = requires { + T::producer::store; +}; + +template __device__ inline void run_op_producer(const typename Op::config::globals &globals, persistent_state &ps) { + static_assert(kernel_template, "interpreter producer kernel template parameter does not satisfy concept requirements"); + using L = typename Op::layout; + using CKL = complete_kittens_layout; + using common_state = typename CKL::common_state_t; + using producer_state = typename CKL::producer_state_t; + using consumer_state = typename CKL::consumer_state_t; + using input_block = typename CKL::input_block_t; + using output_block = typename CKL::output_block_t; + using scratch_block = typename CKL::scratch_block_t; + using finish_block = typename CKL::finish_block_t; + using input_alloc_block = typename CKL::input_alloc_block_t; + using output_alloc_block = typename CKL::output_alloc_block_t; + using scratch_alloc_block = typename CKL::scratch_alloc_block_t; + constexpr int MAX_BLOCK_SHARED_MEMORY = detail::MAX_SHARED_MEMORY_v; + constexpr int INPUT_PIPE_STAGES = detail::INPUT_PIPE_STAGES_v; + static_assert(INPUT_PIPE_STAGES >= 1 && INPUT_PIPE_STAGES <= 8, "Invalid number of input pipe stages"); + if constexpr (has_store) { + constexpr int OUTPUT_PIPE_STAGES = detail::OUTPUT_PIPE_STAGES_v; + static_assert(OUTPUT_PIPE_STAGES >= 1 && OUTPUT_PIPE_STAGES <= 8, "Invalid number of output pipe stages"); + static_assert( + INPUT_PIPE_STAGES*sizeof(input_alloc_block) + + OUTPUT_PIPE_STAGES*sizeof(output_alloc_block) + + sizeof(scratch_alloc_block) + <= MAX_BLOCK_SHARED_MEMORY, "Shared memory usage exceeds limits" + ); + + shared_allocator alloc(ps.shmem); + input_alloc_block (&input_smem) [INPUT_PIPE_STAGES] = alloc.allocate(); + output_alloc_block (&output_smem) [OUTPUT_PIPE_STAGES] = alloc.allocate(); + + + scratch_alloc_block (*scratch_smem) = reinterpret_cast((((uint64_t)(ps.shmem) + (MAX_BLOCK_SHARED_MEMORY - sizeof(scratch_alloc_block)))/1024)*1024); + + finish_block (*finish_smem) = reinterpret_cast((((uint64_t)(scratch_smem) - sizeof(finish_block))/1024)*1024); + static_assert(sizeof(scratch_alloc_block) + sizeof(finish_block) <= MAX_BLOCK_SHARED_MEMORY, "Finish block is too large for shared memory."); + + int danger_stage = ps.max_finish_offset/sizeof(input_alloc_block); + ps.max_finish_offset = (uint64_t)(finish_smem) - (uint64_t)(ps.shmem); + common_state common; + + using producers = group; + producer_state p_state; + int num_iters = 0; +#ifdef KITTENS_SM10X + common_setup_args unif{common, ps.task_iter, num_iters, globals, **scratch_smem, ps.tensor_alloc, ps.instruction}; +#else + common_setup_args unif{common, ps.task_iter, num_iters, globals, **scratch_smem, ps.instruction}; +#endif +#ifdef KITTENS_TIMINGS + unif.timings = ps.timings; +#endif + Op::common_setup(unif); + if(num_iters < 0) return; + int input_ring = 0; + int output_ring = 0; + int load_iter, store_iter = 0; + Op::producer::setup({p_state, unif}); + for(load_iter = 0; load_iter(ps.semaphore_bitfield, input_ring)); + if(load_iter == danger_stage) wait(*ps.finish_finished, (ps.task_iter%2)^1); + update_phasebit<1>(ps.semaphore_bitfield, input_ring); + Op::producer::load({p_state, *input_smem[input_ring], ps.inputs_arrived[input_ring], load_iter, unif}); + input_ring=ring_advance(input_ring); + } + while(load_iter(ps.semaphore_bitfield, output_ring))) { + update_phasebit<0>(ps.semaphore_bitfield, output_ring); + Op::producer::store({p_state, *output_smem[output_ring], ps.outputs_finished[output_ring], store_iter, unif}); + output_ring=ring_advance(output_ring); + store_iter++; + } + + if(load_iter(ps.semaphore_bitfield, input_ring))) { + update_phasebit<1>(ps.semaphore_bitfield, input_ring); + Op::producer::load({p_state, *input_smem[input_ring], ps.inputs_arrived[input_ring], load_iter, unif}); + input_ring=ring_advance(input_ring); + load_iter++; + } + __nanosleep(10); + } + producers::sync(14); + } + else { + static_assert( + INPUT_PIPE_STAGES*sizeof(input_alloc_block) + sizeof(scratch_alloc_block) + <= MAX_SHARED_MEMORY-1024, "Shared memory usage exceeds limits" + ); + + shared_allocator alloc(ps.shmem); + input_alloc_block (&input_smem) [INPUT_PIPE_STAGES] = alloc.allocate(); + + + scratch_alloc_block (*scratch_smem) = reinterpret_cast((((uint64_t)(ps.shmem) + (MAX_BLOCK_SHARED_MEMORY - sizeof(scratch_alloc_block)))/1024)*1024); + + finish_block (*finish_smem) = reinterpret_cast((((uint64_t)(scratch_smem) - sizeof(finish_block))/1024)*1024); + static_assert(sizeof(scratch_alloc_block) + sizeof(finish_block) <= MAX_BLOCK_SHARED_MEMORY, "Finish block is too large for shared memory."); + + int danger_stage = ps.max_finish_offset/sizeof(input_alloc_block); + ps.max_finish_offset = (uint64_t)(finish_smem) - (uint64_t)(ps.shmem); + common_state common; + + using producers = group; + producer_state p_state; + int num_iters = -1; +#ifdef KITTENS_SM10X + common_setup_args unif{common, ps.task_iter, num_iters, globals, **scratch_smem, ps.tensor_alloc, ps.instruction}; +#else + common_setup_args unif{common, ps.task_iter, num_iters, globals, **scratch_smem, ps.instruction}; +#endif +#ifdef KITTENS_TIMINGS + unif.timings = ps.timings; +#endif + Op::common_setup(unif); + if(num_iters < 0) return; + int input_ring = 0; + int load_iter; + Op::producer::setup({p_state, unif}); + for(load_iter = 0; load_iter < INPUT_PIPE_STAGES && load_iter(ps.semaphore_bitfield, input_ring)); + if(load_iter == danger_stage) wait(*ps.finish_finished, (ps.task_iter%2)^1); + update_phasebit<1>(ps.semaphore_bitfield, input_ring); + Op::producer::load({p_state, *input_smem[input_ring], ps.inputs_arrived[input_ring], load_iter, unif}); + input_ring=ring_advance(input_ring); + } + for(; load_iter(ps.semaphore_bitfield, input_ring)); + update_phasebit<1>(ps.semaphore_bitfield, input_ring); + Op::producer::load({p_state, *input_smem[input_ring], ps.inputs_arrived[input_ring], load_iter, unif}); + input_ring=ring_advance(input_ring); + } + producers::sync(14); + } +} +template __device__ inline void run_op_consumer(const typename Op::config::globals &globals, persistent_state &ps) { + static_assert(kernel_template, "interpreter consumer kernel template parameter does not satisfy concept requirements"); + using L = typename Op::layout; + using CKL = complete_kittens_layout; + using common_state = typename CKL::common_state_t; + using producer_state = typename CKL::producer_state_t; + using consumer_state = typename CKL::consumer_state_t; + using input_block = typename CKL::input_block_t; + using output_block = typename CKL::output_block_t; + using scratch_block = typename CKL::scratch_block_t; + using finish_block = typename CKL::finish_block_t; + using input_alloc_block = typename CKL::input_alloc_block_t; + using output_alloc_block = typename CKL::output_alloc_block_t; + using scratch_alloc_block = typename CKL::scratch_alloc_block_t; + constexpr int MAX_BLOCK_SHARED_MEMORY = detail::MAX_SHARED_MEMORY_v; + constexpr int INPUT_PIPE_STAGES = detail::INPUT_PIPE_STAGES_v; + static_assert(INPUT_PIPE_STAGES >= 1 && INPUT_PIPE_STAGES <= 8, "Invalid number of input pipe stages"); + if constexpr (has_store) { + constexpr int OUTPUT_PIPE_STAGES = detail::OUTPUT_PIPE_STAGES_v; + static_assert(OUTPUT_PIPE_STAGES >= 1 && OUTPUT_PIPE_STAGES <= 8, "Invalid number of output pipe stages"); + static_assert( + INPUT_PIPE_STAGES*sizeof(input_alloc_block) + + OUTPUT_PIPE_STAGES*sizeof(output_alloc_block) + + sizeof(scratch_alloc_block) + <= MAX_BLOCK_SHARED_MEMORY, "Shared memory usage exceeds limits" + ); + + shared_allocator alloc(ps.shmem); + input_alloc_block (&input_smem) [INPUT_PIPE_STAGES] = alloc.allocate(); + output_alloc_block (&output_smem) [OUTPUT_PIPE_STAGES] = alloc.allocate(); + + + scratch_alloc_block (*scratch_smem) = reinterpret_cast((((uint64_t)(ps.shmem) + (MAX_BLOCK_SHARED_MEMORY - sizeof(scratch_alloc_block)))/1024)*1024); + + finish_block (*finish_smem) = reinterpret_cast((((uint64_t)(scratch_smem) - sizeof(finish_block))/1024)*1024); + static_assert(sizeof(scratch_alloc_block) + sizeof(finish_block) <= MAX_BLOCK_SHARED_MEMORY, "Finish block is too large for shared memory."); + + common_state common; + + using consumers = group; + consumer_state c_state; + int num_iters = 0; +#ifdef KITTENS_SM10X + common_setup_args unif{common, ps.task_iter, num_iters, globals, **scratch_smem, ps.tensor_alloc, ps.instruction}; +#else + common_setup_args unif{common, ps.task_iter, num_iters, globals, **scratch_smem, ps.instruction}; +#endif +#ifdef KITTENS_TIMINGS + unif.timings = ps.timings; +#endif + Op::common_setup(unif); + if(num_iters < 0) return; + int input_ring = 0; + int output_ring = 0; + Op::consumer::setup({c_state, unif}); +#ifdef CONSUMER_UNROLL + #pragma unroll CONSUMER_UNROLL_VALUE +#endif + for(int it = 0; it < num_iters; it++) { + wait(ps.outputs_finished[output_ring], get_phasebit<1>(ps.semaphore_bitfield, output_ring)); + update_phasebit<1>(ps.semaphore_bitfield, output_ring); + wait(ps.inputs_arrived[input_ring], get_phasebit<0>(ps.semaphore_bitfield, input_ring)); + update_phasebit<0>(ps.semaphore_bitfield, input_ring); + Op::consumer::compute({c_state, *input_smem[input_ring], *output_smem[output_ring], ps.inputs_finished[input_ring], ps.outputs_arrived[output_ring], it, unif}); + input_ring=ring_advance(input_ring); + output_ring=ring_advance(output_ring); + } + + #pragma unroll + for(int i = 0; i < OUTPUT_PIPE_STAGES; i++) { + wait(ps.outputs_finished[i], get_phasebit<1>(ps.semaphore_bitfield, i)); + } + + consumers::sync(13); + Op::consumer::finish({c_state, *finish_smem, *ps.finish_finished, unif}); + consumers::sync(13); + } + else { + static_assert( + INPUT_PIPE_STAGES*sizeof(input_alloc_block) + sizeof(scratch_alloc_block) + <= MAX_SHARED_MEMORY-1024, "Shared memory usage exceeds limits" + ); + + shared_allocator alloc(ps.shmem); + input_alloc_block (&input_smem) [INPUT_PIPE_STAGES] = alloc.allocate(); + + + scratch_alloc_block (*scratch_smem) = reinterpret_cast((((uint64_t)(ps.shmem) + (MAX_BLOCK_SHARED_MEMORY - sizeof(scratch_alloc_block)))/1024)*1024); + + finish_block (*finish_smem) = reinterpret_cast((((uint64_t)(scratch_smem) - sizeof(finish_block))/1024)*1024); + static_assert(sizeof(scratch_alloc_block) + sizeof(finish_block) <= MAX_BLOCK_SHARED_MEMORY, "Finish block is too large for shared memory."); + + common_state common; + + using consumers = group; + consumer_state c_state; + int num_iters = -1; +#ifdef KITTENS_SM10X + common_setup_args unif{common, ps.task_iter, num_iters, globals, **scratch_smem, ps.tensor_alloc, ps.instruction}; +#else + common_setup_args unif{common, ps.task_iter, num_iters, globals, **scratch_smem, ps.instruction}; +#endif +#ifdef KITTENS_TIMINGS + unif.timings = ps.timings; +#endif + Op::common_setup(unif); + if(num_iters < 0) return; + int input_ring = 0; + Op::consumer::setup({c_state, unif}); +#ifdef CONSUMER_UNROLL + #pragma unroll CONSUMER_UNROLL_VALUE +#endif + for(int it = 0; it < num_iters; it++) { + wait(ps.inputs_arrived[input_ring], get_phasebit<0>(ps.semaphore_bitfield, input_ring)); + update_phasebit<0>(ps.semaphore_bitfield, input_ring); + Op::consumer::compute({c_state, *input_smem[input_ring], ps.inputs_finished[input_ring], it, unif}); + input_ring=ring_advance(input_ring); + } + consumers::sync(13); + Op::consumer::finish({c_state, *finish_smem, *ps.finish_finished, unif}); + consumers::sync(13); + } +} + +template +struct dispatch_producer { + __device__ static inline void run(int opcode, const typename config::globals &globals, persistent_state &ps) {} +}; +template +struct dispatch_producer { + __device__ static inline void run(int opcode, const typename config::globals &globals, persistent_state &ps) { + if(opcode == op::opcode) run_op_producer(globals, ps); + else dispatch_producer::run(opcode, globals, ps); + } +}; +template +struct dispatch_consumer { + __device__ static inline void run(int opcode, const typename config::globals &globals, persistent_state &ps) {} +}; +template +struct dispatch_consumer { + __device__ static inline void run(int opcode, const typename config::globals &globals, persistent_state &ps) { + if(opcode == op::opcode) run_op_consumer(globals, ps); + else dispatch_consumer::run(opcode, globals, ps); + } +}; + +template __device__ void inline load_instructions(int *instructions, int task_iter, const typename config::globals &globals, kittens::semaphore &bar) { + if(laneid() == 0) { + constexpr int ints_per_instruction = globals.instructions.cols(); + constexpr int bytes = ints_per_instruction*sizeof(int); + tma::expect_bytes(bar, bytes); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(&bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(instructions)); + uint64_t src_ptr = (uint64_t)(&globals.instructions[kittens::coord<>{(int)(blockIdx.x), task_iter, 0}]); + asm volatile ( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1], %3, [%2];" + : + : "r"(dst_ptr), "l"(src_ptr), "r"(mbar_ptr), "n"(bytes) + : "memory" + ); + } +} + +template +__launch_bounds__((NUM_CONSUMER_WARPS+NUM_PRODUCER_WARPS)*WARP_THREADS, 1) +__cluster_dims__(detail::CLUSTER_BLOCKS_v) +__global__ void kernel(const __grid_constant__ typename config::globals globals) { +#ifdef KITTENS_TIMINGS + uint64_t kernel_start = clock64(); + __shared__ uint64_t timings[2][64]; + if(threadIdx.x < 64) { + timings[0][threadIdx.x] = 0; + timings[1][threadIdx.x] = 0; + } +#endif +#ifdef KITTENS_SM10X + constexpr int NCTA_TENSOR_ALLOC = detail::CLUSTER_BLOCKS_v > 1 ? 2 : 1; + tensor_allocator<1, NCTA_TENSOR_ALLOC> tensor_alloc{}; +#endif + __shared__ __align__(16) int instructions[2][globals.instructions.cols()]; + __shared__ kittens::semaphore inputs_arrived[8], inputs_finished[8], outputs_arrived[8], outputs_finished[8], finish_finished, instruction_arrived[2], instruction_finished[2]; + if(warpid() == 0) { + init_semaphore(instruction_arrived[0], 0, 1); + init_semaphore(instruction_arrived[1], 0, 1); + init_semaphore(instruction_finished[0], NUM_WARPS, 0); + init_semaphore(instruction_finished[1], NUM_WARPS, 0); + } + extern __shared__ int __shm[]; + persistent_state ps{ + 0, + &__shm[0], + 1000000, + &inputs_arrived[0], + &inputs_finished[0], + &outputs_arrived[0], + &outputs_finished[0], + &finish_finished, + 0xFFFF0000, +#ifdef KITTENS_SM10X + tensor_alloc, +#endif +#ifdef KITTENS_TIMINGS + &timings[0][0], +#endif + nullptr + }; + + if(warpid() < 8) { + init_semaphore(inputs_arrived[warpid()], NUM_PRODUCER_WARPS, 0); + init_semaphore(inputs_finished[warpid()], NUM_CONSUMER_WARPS, 0); + init_semaphore(outputs_arrived[warpid()], NUM_CONSUMER_WARPS, 0); + init_semaphore(outputs_finished[warpid()], NUM_PRODUCER_WARPS, 0); + } + if(warpid() == 0) init_semaphore(finish_finished, NUM_CONSUMER_WARPS, 0); + if(detail::CLUSTER_BLOCKS_v == 1) group::sync(15); + else everyone::tma::cluster::sync(); + if(warpid() < NUM_CONSUMER_WARPS) { + warpgroup::increase_registers<232>(); + for(ps.task_iter = 0; ps.task_iter < globals.instructions.rows(); ps.task_iter++) { + ps.instruction = &instructions[ps.task_iter%2][0]; +#ifdef KITTENS_TIMINGS + ps.timings = &timings[ps.task_iter%2][0]; +#endif + wait(instruction_arrived[ps.task_iter%2], ((ps.task_iter/2)%2)); + int opcode = ps.instruction[0]; + if(opcode == 0) break; + dispatch_consumer::run(opcode, globals, ps); +#ifdef KITTENS_TIMINGS + if(threadIdx.x < 64) { + if(ps.timings[threadIdx.x] != 0) { + globals.timings[kittens::coord<>{(int)(blockIdx.x), ps.task_iter, (int)(threadIdx.x)}] = (uint32_t)(ps.timings[threadIdx.x]); + + ps.timings[threadIdx.x] = 0; + } + } +#endif + if(laneid() == 0) arrive(instruction_finished[ps.task_iter%2]); + } + } + else { + warpgroup::decrease_registers<40>(); + if(warpgroup::warpid() == 2) { + load_instructions(&instructions[0][0], 0, globals, instruction_arrived[0]); + } + ps.instruction = &instructions[1][0]; + for(ps.task_iter = 0; ps.task_iter < globals.instructions.rows(); ps.task_iter++) { + if(warpgroup::warpid() == 2 && ps.task_iter+1 < globals.instructions.rows()) { + wait(instruction_finished[(ps.task_iter+1)%2], ((ps.task_iter+3)/2)%2); + __syncwarp(); + load_instructions(ps.instruction, ps.task_iter+1, globals, instruction_arrived[(ps.task_iter+1)%2]); + } + ps.instruction = &instructions[ps.task_iter%2][0]; +#ifdef KITTENS_TIMINGS + ps.timings = &timings[ps.task_iter%2][0]; +#endif + wait(instruction_arrived[ps.task_iter%2], ((ps.task_iter/2)%2)); + int opcode = ps.instruction[0]; + if(opcode == 0) break; + dispatch_producer::run(opcode, globals, ps); + if(laneid() == 0) arrive(instruction_finished[ps.task_iter%2]); + } + } + if(detail::CLUSTER_BLOCKS_v > 1) everyone::tma::cluster::sync(); + else group::sync(15); +} +template +__host__ void run(typename config::globals &globals) { + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, globals.dynamic_shared_memory()); + kernel<<>>(globals); +} + +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/interpreter/templates.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/interpreter/templates.cuh new file mode 100644 index 0000000..b17c2f1 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/interpreter/templates.cuh @@ -0,0 +1,178 @@ +#pragma once + +#include "../common/common.cuh" + +namespace kittens { +namespace prototype { +namespace interpreter { + +struct persistent_state { + int task_iter; + int *shmem; + int max_finish_offset; + kittens::semaphore *inputs_arrived, *outputs_arrived, *inputs_finished, *outputs_finished, *finish_finished; + uint32_t semaphore_bitfield; +#ifdef KITTENS_SM10X + tensor_allocator<1, 1> &tensor_alloc; +#endif +#ifdef KITTENS_TIMINGS + uint64_t *timings; +#endif + int *instruction; +}; + + +template struct uniform_args { + using CKL = complete_kittens_layout; + typename CKL::common_state_t & common; + int & task_iter; + int & num_iters; + const typename CKL::globals_t & globals; + typename CKL::scratch_block_t & scratch; +#ifdef KITTENS_SM10X + tensor_allocator<1, 1> &tensor_alloc; +#endif +#ifdef KITTENS_TIMINGS + uint64_t *timings; +#endif + int *instruction; + __device__ uniform_args( + typename CKL::common_state_t & _common, + int & _task_iter, + int & _num_iters, + const typename CKL::globals_t& _globals, + typename CKL::scratch_block_t& _scratch, +#ifdef KITTENS_SM10X + tensor_allocator<1, 1> &_tensor_alloc, +#endif + int * _instruction + ) : common(_common), + task_iter(_task_iter), + num_iters(_num_iters), + globals(_globals), + scratch(_scratch), +#ifdef KITTENS_SM10X + tensor_alloc(_tensor_alloc), +#endif + instruction(_instruction) {} + __device__ uniform_args(uniform_args &_args) : + common(_args.common), + task_iter(_args.task_iter), + num_iters(_args.num_iters), + globals(_args.globals), + scratch(_args.scratch), +#ifdef KITTENS_SM10X + tensor_alloc(_args.tensor_alloc), +#endif +#ifdef KITTENS_TIMINGS + timings(_args.timings), +#endif + instruction(_args.instruction) {} +}; + + +template using common_setup_args = uniform_args; + + +template struct producer_setup_args : uniform_args { + using CKL = complete_kittens_layout; + typename CKL::producer_state_t & state; + __device__ producer_setup_args( + typename CKL::producer_state_t& _state, + uniform_args &_args + ) : uniform_args(_args), state(_state) {} +}; + + +template struct producer_load_args : uniform_args { + using CKL = complete_kittens_layout; + typename CKL::producer_state_t & state; + typename CKL::input_block_t & input; + kittens::semaphore & inputs_arrived; + int iter; + __device__ producer_load_args( + typename CKL::producer_state_t& _state, + typename CKL::input_block_t& _input, + semaphore& _inputs_arrived, + int _iter, + uniform_args &_args + ) : uniform_args(_args), input(_input), state(_state), inputs_arrived(_inputs_arrived), iter(_iter) {} +}; + + +template struct producer_store_args : uniform_args { + using CKL = complete_kittens_layout; + typename CKL::producer_state_t & state; + typename CKL::output_block_t & output; + kittens::semaphore & outputs_finished; + int iter; + __device__ producer_store_args( + typename CKL::producer_state_t& _state, + typename CKL::output_block_t& _output, + semaphore& _outputs_finished, + int _iter, + uniform_args &_args + ) : uniform_args(_args), output(_output), state(_state), outputs_finished(_outputs_finished), iter(_iter) {} +}; + + +template struct consumer_setup_args : uniform_args { + using CKL = complete_kittens_layout; + typename CKL::consumer_state_t & state; + __device__ consumer_setup_args( + typename CKL::consumer_state_t& _state, + uniform_args &_args + ) : uniform_args(_args), state(_state) {} +}; + + +template struct consumer_compute_args : uniform_args { + using CKL = complete_kittens_layout; + typename CKL::consumer_state_t & state; + typename CKL::input_block_t & input; + kittens::semaphore & inputs_finished; + int iter; + __device__ consumer_compute_args( + typename CKL::consumer_state_t& _state, + typename CKL::input_block_t& _input, + semaphore& _inputs_finished, + int _iter, + uniform_args &_args + ) : uniform_args(_args), input(_input), state(_state), inputs_finished(_inputs_finished), iter(_iter) {} +}; +template struct consumer_compute_args : uniform_args { + using CKL = complete_kittens_layout; + typename CKL::consumer_state_t & state; + typename CKL::input_block_t & input; + typename CKL::output_block_t & output; + kittens::semaphore & inputs_finished; + kittens::semaphore & outputs_arrived; + int iter; + __device__ consumer_compute_args( + typename CKL::consumer_state_t& _state, + typename CKL::input_block_t& _input, + typename CKL::output_block_t& _output, + semaphore& _inputs_finished, + semaphore& _outputs_arrived, + int _iter, + uniform_args &_args + ) : uniform_args(_args), input(_input), output(_output), state(_state), inputs_finished(_inputs_finished), outputs_arrived(_outputs_arrived), iter(_iter) {} +}; + + +template struct consumer_finish_args : uniform_args { + using CKL = complete_kittens_layout; + typename CKL::consumer_state_t & state; + typename CKL::finish_block_t & finish; + kittens::semaphore & finish_finished; + __device__ consumer_finish_args( + typename CKL::consumer_state_t& _state, + typename CKL::finish_block_t& _finish, + semaphore& _finish_finished, + uniform_args &_args + ) : uniform_args(_args), finish(_finish), state(_state), finish_finished(_finish_finished) {} +}; + +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/lcf/lcf.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/lcf/lcf.cuh new file mode 100644 index 0000000..b70c80d --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/lcf/lcf.cuh @@ -0,0 +1,188 @@ +#pragma once + +#include +#include "../common/common.cuh" +#include "templates.cuh" + +namespace kittens { +namespace prototype { +namespace lcf { + +template concept kernel_template = requires { + typename lcft::layout; + typename lcft::producer; + typename lcft::consumer; + lcft::common_setup; + lcft::producer::setup; + lcft::producer::load; + lcft::consumer::setup; + lcft::consumer::compute; + lcft::consumer::finish; +} && kittens_layout; + +template +__global__ __launch_bounds__(detail::NUM_THREADS_v, detail::NUM_BLOCKS_v) +#ifdef KITTENS_SM10X +__cluster_dims__(detail::CLUSTER_BLOCKS_v) +#endif +void kernel(const __grid_constant__ typename lcft::layout::globals globals) { + static_assert(kernel_template, "lcf kernel template parameter does not satisfy concept requirements"); + using L = typename lcft::layout; + using CKL = complete_kittens_layout; + using common_state = typename CKL::common_state_t; + using producer_state = typename CKL::producer_state_t; + using consumer_state = typename CKL::consumer_state_t; + using input_block = typename CKL::input_block_t; + using scratch_block = typename CKL::scratch_block_t; + using finish_block = typename CKL::finish_block_t; + using input_alloc_block = typename CKL::input_alloc_block_t; + using scratch_alloc_block = typename CKL::scratch_alloc_block_t; + constexpr int MAX_SHARED_MEMORY = detail::MAX_SHARED_MEMORY_v; + constexpr int INPUT_PIPE_STAGES = detail::INPUT_PIPE_STAGES_v; + static_assert(INPUT_PIPE_STAGES >= 1 && INPUT_PIPE_STAGES <= 16, "Invalid number of input pipe stages"); + static_assert( + INPUT_PIPE_STAGES*sizeof(input_alloc_block) + sizeof(scratch_alloc_block) + <= MAX_SHARED_MEMORY-1024, "Shared memory usage exceeds limits" + ); + constexpr int NUM_CONSUMER_WARPS = detail::NUM_CONSUMER_WARPS_v; + constexpr int NUM_PRODUCER_WARPS = detail::NUM_PRODUCER_WARPS_v; + +#ifdef KITTENS_SM10X + constexpr int NCTA_TENSOR_ALLOC = detail::CLUSTER_BLOCKS_v > 1 ? 2 : 1; + tensor_allocator, NCTA_TENSOR_ALLOC> tensor_alloc{}; +#endif + + extern __shared__ int __shm[]; + shared_allocator alloc(&__shm[0]); + scratch_alloc_block (&scratch_smem) = alloc.allocate(); + input_alloc_block (&input_smem) [INPUT_PIPE_STAGES] = alloc.allocate(); + + + constexpr int FINISH_BLOCK_OFFSET = (MAX_SHARED_MEMORY-1024)/detail::NUM_BLOCKS_v - sizeof(finish_block); + static_assert(FINISH_BLOCK_OFFSET >= 0, "Finish block is too large for shared memory."); + constexpr int NON_FINISH_BLOCK_SPACE = FINISH_BLOCK_OFFSET - 1024 - sizeof(scratch_alloc_block); + constexpr int SAFE_STAGES_BETWEEN_BLOCKS = NON_FINISH_BLOCK_SPACE/sizeof(input_alloc_block)((((uint64_t)&__shm[0] + FINISH_BLOCK_OFFSET)/1024)*1024); + + if constexpr (detail::DEBUG_v) { + if(threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0) { + printf("DEBUG REPORT FOR PRODUCER TEMPLATE KERNEL:\n"); + printf(" BLOCK INFORMATION\n"); + printf(" gridDim.x: %d\n", gridDim.x); + printf(" gridDim.y: %d\n", gridDim.y); + printf(" gridDim.z: %d\n", gridDim.z); + printf(" blockDim.x: %d\n", blockDim.x); + printf(" blockDim.y: %d\n", blockDim.y); + printf(" blockDim.z: %d\n", blockDim.z); + printf(" num_blocks per SM: %d\n", detail::NUM_BLOCKS_v); + printf(" num_threads per SM: %d\n", detail::NUM_THREADS_v); + printf(" num_warps per SM: %d\n", detail::NUM_WARPS_v); + printf(" num_consumer_warpgroups: %d\n", detail::NUM_CONSUMER_WARPGROUPS_v); + printf(" num_consumer_warps: %d\n", detail::NUM_CONSUMER_WARPS_v); + printf(" num_producer_warps: %d\n", detail::NUM_PRODUCER_WARPS_v); + printf(" PIPELINE INFORMATION\n"); + printf(" input_pipe_stages: %d\n", INPUT_PIPE_STAGES); + printf(" safe_stages_between_blocks: %d\n", SAFE_STAGES_BETWEEN_BLOCKS); + printf(" SHARED MEMORY INFORMATION\n"); + printf(" input_smem block size: %llu\n", sizeof(input_block)); + printf(" input_smem block size (aligned): %llu\n", sizeof(input_alloc_block)); + printf(" input_smem: %p\n", (void*)&input_smem); + printf(" input_smem size: %llu\n", INPUT_PIPE_STAGES*sizeof(input_alloc_block)); + printf(" scratch_smem block size: %llu\n", sizeof(scratch_block)); + printf(" scratch_smem block size (aligned): %llu\n", sizeof(scratch_alloc_block)); + printf(" scratch_smem: %p\n", (void*)&scratch_smem); + printf(" finish_smem: %p\n", (void*)finish_smem); + printf(" finish_smem size: %llu\n", sizeof(finish_block)); + printf(" dynamic shared memory usage: %llu\n", sizeof(scratch_alloc_block) + uint64_t(&scratch_smem) - uint64_t(&__shm[0])); + } + everyone::sync(15); + } + + + __shared__ kittens::semaphore inputs_arrived[INPUT_PIPE_STAGES], inputs_finished[INPUT_PIPE_STAGES]; + __shared__ kittens::semaphore finish_finished; + uint32_t semaphore_bitfield = 0xFFFF0000; + common_state common; + + if(warpid() >= NUM_CONSUMER_WARPS) { + using producers = group; + if (warpid() == NUM_CONSUMER_WARPS) { + for(int i = 0; i < INPUT_PIPE_STAGES; i++) { + init_semaphore(inputs_arrived[i], detail::PRODUCER_BARRIER_ARRIVALS_v, 0); + init_semaphore(inputs_finished[i], detail::CONSUMER_BARRIER_ARRIVALS_v, 0); + } + init_semaphore(finish_finished, detail::CONSUMER_BARRIER_ARRIVALS_v, 0); + } + + if constexpr (detail::CLUSTER_BLOCKS_v > 1) everyone::tma::cluster::sync(); + else everyone::sync(15); + producer_state p_state; + for(int task_iter = 0; true; task_iter++) { + int num_iters = -1; +#ifdef KITTENS_SM10X + common_setup_args unif{common, task_iter, num_iters, globals, *scratch_smem, tensor_alloc}; +#else + common_setup_args unif{common, task_iter, num_iters, globals, *scratch_smem}; +#endif + lcft::common_setup(unif); + if(num_iters < 0) break; + int input_ring = 0; + int load_iter; + lcft::producer::setup({p_state, unif}); + for(load_iter = 0; load_iter < SAFE_STAGES_BETWEEN_BLOCKS && load_iter(semaphore_bitfield, input_ring)); + update_phasebit<1>(semaphore_bitfield, input_ring); + lcft::producer::load({p_state, *input_smem[input_ring], inputs_arrived[input_ring], load_iter, unif}); + input_ring=ring_advance(input_ring); + } + wait(finish_finished, (task_iter%2)^1); + for(; load_iter(semaphore_bitfield, input_ring)); + update_phasebit<1>(semaphore_bitfield, input_ring); + lcft::producer::load({p_state, *input_smem[input_ring], inputs_arrived[input_ring], load_iter, unif}); + input_ring=ring_advance(input_ring); + } + producers::sync(13); + } + } + else { + using consumers = group; + + if constexpr (detail::CLUSTER_BLOCKS_v > 1) everyone::tma::cluster::sync(); + else everyone::sync(15); + consumer_state c_state; + for(int task_iter = 0; true; task_iter++) { + int num_iters = -1; +#ifdef KITTENS_SM10X + common_setup_args unif{common, task_iter, num_iters, globals, *scratch_smem, tensor_alloc}; +#else + common_setup_args unif{common, task_iter, num_iters, globals, *scratch_smem}; +#endif + lcft::common_setup(unif); + if(num_iters < 0) break; + int input_ring = 0; + lcft::consumer::setup({c_state, unif}); +#ifdef CONSUMER_UNROLL + #pragma unroll CONSUMER_UNROLL_VALUE +#endif + for(int it = 0; it < num_iters; it++) { + wait(inputs_arrived[input_ring], get_phasebit<0>(semaphore_bitfield, input_ring)); + update_phasebit<0>(semaphore_bitfield, input_ring); + lcft::consumer::compute({c_state, *input_smem[input_ring], inputs_finished[input_ring], it, unif}); + input_ring=ring_advance(input_ring); + } + consumers::sync(14); + lcft::consumer::finish({c_state, *finish_smem, finish_finished, unif}); + consumers::sync(14); + } + } + + if constexpr (detail::CLUSTER_BLOCKS_v > 1) everyone::tma::cluster::sync(); +#ifdef KITTENS_SM10X + else everyone::sync(15); +#endif +} + +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/lcf/templates.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/lcf/templates.cuh new file mode 100644 index 0000000..e64d3bb --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/lcf/templates.cuh @@ -0,0 +1,121 @@ +#pragma once + +#include "../common/common.cuh" + +namespace kittens { +namespace prototype { +namespace lcf { + + +template struct uniform_args { + using CKL = complete_kittens_layout; + typename CKL::common_state_t & common; + int task_iter; + int & num_iters; + const typename CKL::globals_t & globals; + typename CKL::scratch_block_t & scratch; +#ifdef KITTENS_SM10X + kittens::tensor_allocator<512 / detail::NUM_BLOCKS_v, detail::CLUSTER_BLOCKS_v == 1 ? 1 : 2> & tt; +#endif + __device__ uniform_args( + typename CKL::common_state_t& _common, + int _task_iter, + int & _num_iters, + const typename CKL::globals_t& _globals, + typename CKL::scratch_block_t& _scratch +#ifdef KITTENS_SM10X + , kittens::tensor_allocator<512 / detail::NUM_BLOCKS_v, detail::CLUSTER_BLOCKS_v == 1 ? 1 : 2> & _tt +#endif + ) : common(_common), + task_iter(_task_iter), + num_iters(_num_iters), + globals(_globals), + scratch(_scratch) +#ifdef KITTENS_SM10X + , tt(_tt) +#endif + {} + __device__ uniform_args(uniform_args &_args) : + common(_args.common), + task_iter(_args.task_iter), + num_iters(_args.num_iters), + globals(_args.globals), + scratch(_args.scratch) +#ifdef KITTENS_SM10X + , tt(_args.tt) +#endif + {} +}; + + +template using common_setup_args = uniform_args; + + +template struct producer_setup_args : uniform_args { + using CKL = complete_kittens_layout; + typename CKL::producer_state_t & state; + __device__ producer_setup_args( + typename CKL::producer_state_t& _state, + uniform_args &_args + ) : uniform_args(_args), state(_state) {} +}; + + +template struct producer_load_args : uniform_args { + using CKL = complete_kittens_layout; + typename CKL::producer_state_t & state; + typename CKL::input_block_t & input; + kittens::semaphore & inputs_arrived; + int iter; + __device__ producer_load_args( + typename CKL::producer_state_t& _state, + typename CKL::input_block_t& _input, + semaphore& _inputs_arrived, + int _iter, + uniform_args &_args + ) : uniform_args(_args), input(_input), state(_state), inputs_arrived(_inputs_arrived), iter(_iter) {} +}; + + +template struct consumer_setup_args : uniform_args { + using CKL = complete_kittens_layout; + typename CKL::consumer_state_t & state; + __device__ consumer_setup_args( + typename CKL::consumer_state_t& _state, + uniform_args &_args + ) : uniform_args(_args), state(_state) {} +}; + + +template struct consumer_compute_args : uniform_args { + using CKL = complete_kittens_layout; + typename CKL::consumer_state_t & state; + typename CKL::input_block_t & input; + kittens::semaphore & inputs_finished; + int iter; + __device__ consumer_compute_args( + typename CKL::consumer_state_t& _state, + typename CKL::input_block_t& _input, + semaphore& _inputs_finished, + int _iter, + uniform_args &_args + ) : uniform_args(_args), input(_input), state(_state), inputs_finished(_inputs_finished), iter(_iter) {} +}; + + +template struct consumer_finish_args : uniform_args { + using CKL = complete_kittens_layout; + typename CKL::consumer_state_t & state; + typename CKL::finish_block_t & finish; + kittens::semaphore & finish_finished; + __device__ consumer_finish_args( + typename CKL::consumer_state_t& _state, + typename CKL::finish_block_t& _finish, + semaphore& _finish_finished, + uniform_args &_args + ) : uniform_args(_args), finish(_finish), state(_state), finish_finished(_finish_finished) {} +}; + +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/lcsc/lcsc.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/lcsc/lcsc.cuh new file mode 100644 index 0000000..ca0fdb9 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/lcsc/lcsc.cuh @@ -0,0 +1,117 @@ +#pragma once + +#include + +namespace kittens { +namespace prototype { +namespace lcsc { + +template +concept has_min_blocks_per_sm = requires { std::integral_constant{}; }; + +template +__host__ consteval int min_blocks_per_sm() { + if constexpr(has_min_blocks_per_sm) + return config::MIN_BLOCKS_PER_SM; + else + return 1; +} + +template +__launch_bounds__(config::NUM_THREADS, min_blocks_per_sm()) +__global__ void main_kernel(const __grid_constant__ globals G) { + extern __shared__ int __shm[]; + tma_swizzle_allocator allocator((int*)&__shm[0]); + + const int num_comm_sms = G.num_comm_sms; + const int num_comp_sms = gridDim.x - num_comm_sms; + + if (blockIdx.x >= num_comp_sms) { + + static_assert(sizeof(typename lcsct::comm_smem) <= config::DYNAMIC_SHARED_MEMORY, "Comm shared memory size exceeds dynamic shared memory limit"); + typename lcsct::comm_regs regs = { + .worker_id = static_cast(blockIdx.x) - num_comp_sms, + .comm_workers = num_comm_sms, + .comp_workers = num_comp_sms, + }; + typename lcsct::comm_smem (&smem) = allocator.allocate(); + __shared__ typename lcsct::comm_sem sem; + + lcsct::comm_setup(G, sem, smem, regs); + + for (regs.task_id = regs.task_offset; regs.task_id < regs.num_tasks; regs.task_id += regs.task_stride) + lcsct::communicator(G, sem, smem, regs); + } else { + + static_assert(sizeof(typename lcsct::comp_smem) <= config::DYNAMIC_SHARED_MEMORY, "Comp shared memory size exceeds dynamic shared memory limit"); + typename lcsct::comp_regs regs = { + .worker_id = static_cast(blockIdx.x), + .comm_workers = num_comm_sms, + .comp_workers = num_comp_sms, + }; + typename lcsct::comp_smem (&smem) = allocator.allocate(); + __shared__ typename lcsct::comp_sem sem; + + int warpgroup_id = warpgroup::groupid(); + int warp_id = warpgroup::warpid(); + int lane_id = warp::laneid(); + lcsct::comp_setup(G, sem, smem, regs); + + if (warpgroup_id == config::NUM_WARPGROUPS - 1) { + warpgroup::decrease_registers(); + + if (warp_id == 0 && lane_id == 0) { + + for (regs.task_id = regs.task_offset; regs.task_id < regs.num_tasks; regs.task_id += regs.task_stride) + lcsct::loader(G, sem, smem, regs); + } else if (warp_id == 1 && lane_id == 0) { + + for (regs.task_id = regs.task_offset; regs.task_id < regs.num_tasks; regs.task_id += regs.task_stride) + lcsct::storer(G, sem, smem, regs); + } + } else { + + warpgroup::increase_registers(); + + for (regs.task_id = regs.task_offset; regs.task_id < regs.num_tasks; regs.task_id += regs.task_stride) + lcsct::consumer(G, sem, smem, regs); + } + } +} + +template +concept static_grid = requires { config::NUM_BLOCKS; }; + +template +concept static_block = requires { config::NUM_THREADS; }; + +template +concept static_dynamic_shared_memory = requires { config::DYNAMIC_SHARED_MEMORY; }; + +template +__host__ static inline void launch_kernel(const globals &G, cudaStream_t stream = 0) { + dim3 grid; + if constexpr (static_grid) + grid = dim3{config::NUM_BLOCKS, 1, 1}; + else + grid = G.grid(); + + dim3 block; + if constexpr (static_block) + block = dim3{config::NUM_THREADS, 1, 1}; + else + block = G.block(); + + int dynamic_shared_memory; + if constexpr (static_dynamic_shared_memory) + dynamic_shared_memory = static_cast(config::DYNAMIC_SHARED_MEMORY); + else + dynamic_shared_memory = G.dynamic_shared_memory(); + + CUDACHECK(cudaFuncSetAttribute(main_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dynamic_shared_memory)); + main_kernel<<>>(G); +} + +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/lcsf/lcsf.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/lcsf/lcsf.cuh new file mode 100644 index 0000000..e22008b --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/lcsf/lcsf.cuh @@ -0,0 +1,229 @@ +#pragma once + +#include +#include "../common/common.cuh" +#include "templates.cuh" + +namespace kittens { +namespace prototype { +namespace lcsf { + +template concept kernel_template = requires { + typename lcsft::layout; + typename lcsft::producer; + typename lcsft::consumer; + lcsft::common_setup; + lcsft::producer::setup; + lcsft::producer::load; + lcsft::producer::store; + lcsft::consumer::setup; + lcsft::consumer::compute; + lcsft::consumer::finish; +} && kittens_layout; + + +template +__global__ __launch_bounds__(detail::NUM_THREADS_v, detail::NUM_BLOCKS_v) +void kernel(const __grid_constant__ typename lcsft::layout::globals globals) { + static_assert(kernel_template, "lcsf kernel template parameter does not satisfy concept requirements"); + using L = typename lcsft::layout; + using CKL = complete_kittens_layout; + using common_state = typename CKL::common_state_t; + using producer_state = typename CKL::producer_state_t; + using consumer_state = typename CKL::consumer_state_t; + using input_block = typename CKL::input_block_t; + using output_block = typename CKL::output_block_t; + using scratch_block = typename CKL::scratch_block_t; + using finish_block = typename CKL::finish_block_t; + using input_alloc_block = typename CKL::input_alloc_block_t; + using output_alloc_block = typename CKL::output_alloc_block_t; + using scratch_alloc_block = typename CKL::scratch_alloc_block_t; + constexpr int MAX_BLOCK_SHARED_MEMORY = detail::MAX_SHARED_MEMORY_v; + constexpr int OUTPUT_PIPE_STAGES = detail::OUTPUT_PIPE_STAGES_v; + static_assert(OUTPUT_PIPE_STAGES >= 1 && OUTPUT_PIPE_STAGES <= 16, "Invalid number of output pipe stages"); + constexpr int INPUT_PIPE_STAGES = detail::INPUT_PIPE_STAGES_v; + static_assert(INPUT_PIPE_STAGES >= 1 && INPUT_PIPE_STAGES <= 16, "Invalid number of input pipe stages"); + static_assert( + INPUT_PIPE_STAGES*sizeof(input_alloc_block) + + OUTPUT_PIPE_STAGES*sizeof(output_alloc_block) + + sizeof(scratch_alloc_block) + <= MAX_BLOCK_SHARED_MEMORY, "Shared memory usage exceeds limits" + ); + constexpr int NUM_CONSUMER_WARPS = detail::NUM_CONSUMER_WARPS_v; + constexpr int NUM_PRODUCER_WARPS = detail::NUM_PRODUCER_WARPS_v; + +#ifdef KITTENS_SM10X + constexpr int NCTA_TENSOR_ALLOC = detail::CLUSTER_BLOCKS_v > 1 ? 2 : 1; + tensor_allocator, NCTA_TENSOR_ALLOC> tensor_alloc{}; +#endif + + extern __shared__ int __shm[]; + shared_allocator alloc(&__shm[0]); + scratch_alloc_block (&scratch_smem) = alloc.allocate(); + input_alloc_block (&input_smem) [INPUT_PIPE_STAGES] = alloc.allocate(); + output_alloc_block (&output_smem) [OUTPUT_PIPE_STAGES] = alloc.allocate(); + + + constexpr int FINISH_BLOCK_OFFSET = (MAX_SHARED_MEMORY-1024)/detail::NUM_BLOCKS_v - sizeof(finish_block); + static_assert(sizeof(scratch_alloc_block) + sizeof(finish_block) <= MAX_BLOCK_SHARED_MEMORY, "Finish block is too large for shared memory."); + constexpr int NON_FINISH_BLOCK_SPACE = FINISH_BLOCK_OFFSET - 1024 - sizeof(scratch_alloc_block); + constexpr int SAFE_STAGES_BETWEEN_BLOCKS = NON_FINISH_BLOCK_SPACE/sizeof(input_alloc_block)((((uint64_t)&__shm[0] + FINISH_BLOCK_OFFSET)/1024)*1024); + + if constexpr (detail::DEBUG_v) { + if(threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0) { + printf("DEBUG REPORT FOR PRODUCER TEMPLATE KERNEL:\n"); + printf(" BLOCK INFORMATION\n"); + printf(" gridDim.x: %d\n", gridDim.x); + printf(" gridDim.y: %d\n", gridDim.y); + printf(" gridDim.z: %d\n", gridDim.z); + printf(" blockDim.x: %d\n", blockDim.x); + printf(" blockDim.y: %d\n", blockDim.y); + printf(" blockDim.z: %d\n", blockDim.z); + printf(" num_blocks per SM: %d\n", detail::NUM_BLOCKS_v); + printf(" num_threads per SM: %d\n", detail::NUM_THREADS_v); + printf(" num_warps per SM: %d\n", detail::NUM_WARPS_v); + printf(" num_consumer_warpgroups: %d\n", detail::NUM_CONSUMER_WARPGROUPS_v); + printf(" num_consumer_warps: %d\n", detail::NUM_CONSUMER_WARPS_v); + printf(" num_producer_warps: %d\n", detail::NUM_PRODUCER_WARPS_v); + printf(" PIPELINE INFORMATION\n"); + printf(" input_pipe_stages: %d\n", INPUT_PIPE_STAGES); + printf(" output_pipe_stages: %d\n", OUTPUT_PIPE_STAGES); + printf(" safe_stages_between_blocks: %d\n", SAFE_STAGES_BETWEEN_BLOCKS); + printf(" SHARED MEMORY INFORMATION\n"); + printf(" input_smem block size: %llu\n", sizeof(input_block)); + printf(" input_smem block size (aligned): %llu\n", sizeof(input_alloc_block)); + printf(" input_smem: %p\n", (void*)&input_smem); + printf(" input_smem size: %llu\n", INPUT_PIPE_STAGES*sizeof(input_alloc_block)); + printf(" output_smem block size: %llu\n", sizeof(output_block)); + printf(" output_smem block size (aligned): %llu\n", sizeof(output_alloc_block)); + printf(" output_smem: %p\n", (void*)&output_smem); + printf(" output_smem size: %llu\n", OUTPUT_PIPE_STAGES*sizeof(output_alloc_block)); + printf(" scratch_smem block size: %llu\n", sizeof(scratch_block)); + printf(" scratch_smem block size (aligned): %llu\n", sizeof(scratch_alloc_block)); + printf(" scratch_smem: %p\n", (void*)&scratch_smem); + printf(" finish_smem: %p\n", (void*)finish_smem); + printf(" finish_smem size: %llu\n", sizeof(finish_block)); + printf(" dynamic shared memory usage: %llu\n", sizeof(output_block) + uint64_t(&(*output_smem[OUTPUT_PIPE_STAGES-1])) - uint64_t(&__shm[0])); + } + everyone::sync(15); + } + + + __shared__ kittens::semaphore inputs_arrived[INPUT_PIPE_STAGES], inputs_finished[INPUT_PIPE_STAGES]; + __shared__ kittens::semaphore outputs_arrived[OUTPUT_PIPE_STAGES], outputs_finished[OUTPUT_PIPE_STAGES]; + __shared__ kittens::semaphore finish_finished; + uint32_t semaphore_bitfield = 0xFFFF0000; + common_state common; + + if (threadIdx.x == 0) { + for(int i = 0; i < INPUT_PIPE_STAGES; i++) { + init_semaphore(inputs_arrived[i], detail::PRODUCER_BARRIER_ARRIVALS_v, 0); + init_semaphore(inputs_finished[i], detail::CONSUMER_BARRIER_ARRIVALS_v, 0); + } + for(int i = 0; i < OUTPUT_PIPE_STAGES; i++) { + init_semaphore(outputs_arrived[i], detail::CONSUMER_BARRIER_ARRIVALS_v, 0); + init_semaphore(outputs_finished[i], detail::PRODUCER_BARRIER_ARRIVALS_v, 0); + } + init_semaphore(finish_finished, detail::CONSUMER_BARRIER_ARRIVALS_v, 0); + } + + if constexpr (detail::CLUSTER_BLOCKS_v > 1) everyone::tma::cluster::sync(); + else everyone::sync(15); + + if(warpid() >= NUM_CONSUMER_WARPS) { + using producers = group; + producer_state p_state; + for(int task_iter = 0; true; task_iter++) { + int num_iters = 0; +#ifdef KITTENS_SM10X + common_setup_args unif{common, task_iter, num_iters, globals, *scratch_smem, tensor_alloc}; +#else + common_setup_args unif{common, task_iter, num_iters, globals, *scratch_smem}; +#endif + lcsft::common_setup(unif); + if(num_iters < 0) break; + int input_ring = 0; + int output_ring = 0; + int load_iter, store_iter = 0; + lcsft::producer::setup({p_state, unif}); + for(load_iter = 0; load_iter(semaphore_bitfield, input_ring)); + update_phasebit<1>(semaphore_bitfield, input_ring); + lcsft::producer::load({p_state, *input_smem[input_ring], inputs_arrived[input_ring], load_iter, unif}); + input_ring=ring_advance(input_ring); + } + wait(finish_finished, (task_iter%2)^1); + for(; load_iter(semaphore_bitfield, input_ring)); + update_phasebit<1>(semaphore_bitfield, input_ring); + lcsft::producer::load({p_state, *input_smem[input_ring], inputs_arrived[input_ring], load_iter, unif}); + input_ring=ring_advance(input_ring); + } + while(load_iter(semaphore_bitfield, output_ring))) { + update_phasebit<0>(semaphore_bitfield, output_ring); + lcsft::producer::store({p_state, *output_smem[output_ring], outputs_finished[output_ring], store_iter, unif}); + output_ring=ring_advance(output_ring); + store_iter++; + } + + if(load_iter(semaphore_bitfield, input_ring))) { + update_phasebit<1>(semaphore_bitfield, input_ring); + lcsft::producer::load({p_state, *input_smem[input_ring], inputs_arrived[input_ring], load_iter, unif}); + input_ring=ring_advance(input_ring); + load_iter++; + } + __nanosleep(10); + } + producers::sync(13); + } + } + else { + using consumers = group; + consumer_state c_state; + for(int task_iter = 0; true; task_iter++) { + int num_iters = 0; +#ifdef KITTENS_SM10X + common_setup_args unif{common, task_iter, num_iters, globals, *scratch_smem, tensor_alloc}; +#else + common_setup_args unif{common, task_iter, num_iters, globals, *scratch_smem}; +#endif + lcsft::common_setup(unif); + if(num_iters < 0) break; + int input_ring = 0; + int output_ring = 0; + lcsft::consumer::setup({c_state, unif}); +#ifdef CONSUMER_UNROLL + #pragma unroll CONSUMER_UNROLL_VALUE +#endif + for(int it = 0; it < num_iters; it++) { + wait(outputs_finished[output_ring], get_phasebit<1>(semaphore_bitfield, output_ring)); + update_phasebit<1>(semaphore_bitfield, output_ring); + wait(inputs_arrived[input_ring], get_phasebit<0>(semaphore_bitfield, input_ring)); + update_phasebit<0>(semaphore_bitfield, input_ring); + lcsft::consumer::compute({c_state, *input_smem[input_ring], *output_smem[output_ring], inputs_finished[input_ring], outputs_arrived[output_ring], it, unif}); + input_ring=ring_advance(input_ring); + output_ring=ring_advance(output_ring); + } + + #pragma unroll + for(int i = 0; i < OUTPUT_PIPE_STAGES; i++) { + wait(outputs_finished[i], get_phasebit<1>(semaphore_bitfield, i)); + } + + consumers::sync(14); + lcsft::consumer::finish({c_state, *finish_smem, finish_finished, unif}); + consumers::sync(14); + } + } + + if constexpr (detail::CLUSTER_BLOCKS_v > 1) everyone::tma::cluster::sync(); +#ifdef KITTENS_SM10X + else everyone::sync(15); +#endif +} + +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/lcsf/templates.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/lcsf/templates.cuh new file mode 100644 index 0000000..c4d7e87 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/lcsf/templates.cuh @@ -0,0 +1,141 @@ +#pragma once + +#include "../common/common.cuh" + +namespace kittens { +namespace prototype { +namespace lcsf { + + +template struct uniform_args { + using CKL = complete_kittens_layout; + typename CKL::common_state_t & common; + int & task_iter; + int & num_iters; + const typename CKL::globals_t & globals; + typename CKL::scratch_block_t & scratch; +#ifdef KITTENS_SM10X + kittens::tensor_allocator<512 / detail::NUM_BLOCKS_v, detail::CLUSTER_BLOCKS_v == 1 ? 1 : 2> & tt; +#endif + __device__ uniform_args( + typename CKL::common_state_t & _common, + int & _task_iter, + int & _num_iters, + const typename CKL::globals_t& _globals, + typename CKL::scratch_block_t& _scratch +#ifdef KITTENS_SM10X + , kittens::tensor_allocator<512 / detail::NUM_BLOCKS_v, detail::CLUSTER_BLOCKS_v == 1 ? 1 : 2> & _tt +#endif + ) : common(_common), + task_iter(_task_iter), + num_iters(_num_iters), + globals(_globals), + scratch(_scratch) +#ifdef KITTENS_SM10X + , tt(_tt) +#endif + {} + __device__ uniform_args(uniform_args &_args) : + common(_args.common), + task_iter(_args.task_iter), + num_iters(_args.num_iters), + globals(_args.globals), + scratch(_args.scratch) +#ifdef KITTENS_SM10X + , tt(_args.tt) +#endif + {} +}; + + +template using common_setup_args = uniform_args; + + +template struct producer_setup_args : uniform_args { + using CKL = complete_kittens_layout; + typename CKL::producer_state_t & state; + __device__ producer_setup_args( + typename CKL::producer_state_t& _state, + uniform_args &_args + ) : uniform_args(_args), state(_state) {} +}; + + +template struct producer_load_args : uniform_args { + using CKL = complete_kittens_layout; + typename CKL::producer_state_t & state; + typename CKL::input_block_t & input; + kittens::semaphore & inputs_arrived; + int iter; + __device__ producer_load_args( + typename CKL::producer_state_t& _state, + typename CKL::input_block_t& _input, + semaphore& _inputs_arrived, + int _iter, + uniform_args &_args + ) : uniform_args(_args), input(_input), state(_state), inputs_arrived(_inputs_arrived), iter(_iter) {} +}; + + +template struct producer_store_args : uniform_args { + using CKL = complete_kittens_layout; + typename CKL::producer_state_t & state; + typename CKL::output_block_t & output; + kittens::semaphore & outputs_finished; + int iter; + __device__ producer_store_args( + typename CKL::producer_state_t& _state, + typename CKL::output_block_t& _output, + semaphore& _outputs_finished, + int _iter, + uniform_args &_args + ) : uniform_args(_args), output(_output), state(_state), outputs_finished(_outputs_finished), iter(_iter) {} +}; + + +template struct consumer_setup_args : uniform_args { + using CKL = complete_kittens_layout; + typename CKL::consumer_state_t & state; + __device__ consumer_setup_args( + typename CKL::consumer_state_t& _state, + uniform_args &_args + ) : uniform_args(_args), state(_state) {} +}; + + +template struct consumer_compute_args : uniform_args { + using CKL = complete_kittens_layout; + typename CKL::consumer_state_t & state; + typename CKL::input_block_t & input; + typename CKL::output_block_t & output; + kittens::semaphore & inputs_finished; + kittens::semaphore & outputs_arrived; + int iter; + __device__ consumer_compute_args( + typename CKL::consumer_state_t& _state, + typename CKL::input_block_t& _input, + typename CKL::output_block_t& _output, + semaphore& _inputs_finished, + semaphore& _outputs_arrived, + int _iter, + uniform_args &_args + ) : uniform_args(_args), input(_input), output(_output), state(_state), inputs_finished(_inputs_finished), outputs_arrived(_outputs_arrived), iter(_iter) {} +}; + + +template struct consumer_finish_args : uniform_args { + using CKL = complete_kittens_layout; + typename CKL::consumer_state_t & state; + typename CKL::finish_block_t & finish; + kittens::semaphore & finish_finished; + __device__ consumer_finish_args( + typename CKL::consumer_state_t& _state, + typename CKL::finish_block_t& _finish, + semaphore& _finish_finished, + uniform_args &_args + ) : uniform_args(_args), finish(_finish), state(_state), finish_finished(_finish_finished) {} +}; + +} +} +} diff --git a/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/prototype.cuh b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/prototype.cuh new file mode 100644 index 0000000..cdc8340 --- /dev/null +++ b/deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/prototype.cuh @@ -0,0 +1,14 @@ + + + + + +#pragma once + +#include + +#include "common/common.cuh" +#include "lcf/lcf.cuh" +#include "lcsc/lcsc.cuh" +#include "lcsf/lcsf.cuh" +#include "interpreter/interpreter.cuh" diff --git a/deploy/nemotron_native/model.lock.json b/deploy/nemotron_native/model.lock.json new file mode 100644 index 0000000..4c90842 --- /dev/null +++ b/deploy/nemotron_native/model.lock.json @@ -0,0 +1,122 @@ +{ + "schema": "nemotron-native-model-lock.v1", + "model_identity": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa", + "inputs": { + "repository": "nvidia/easymagpietts_NEXT", + "revision": "3f12eb7fedcef8d432bd83d0a82525e8de6b3d39", + "token_environment_variable": "HF_TOKEN", + "files": [ + { + "id": "easymagpie_predictor_nemo", + "path": "2605_NemotronTTS_V0.2/v2/2605_EMTTS_SmallMamba_Step150k.nemo", + "sha256": "73f1591f9029624960dafd8e94f46f9cce7cc8141ba854f102d346fa2a8f4d01", + "bytes": 5067622400 + }, + { + "id": "codec_25fps_nemo", + "path": "25fps_spectral_codec_with_bandwidth_extension.nemo", + "sha256": "385dbccbb9eb94865c9745f28f2748845a9a567de7338375825fa2727095d243", + "bytes": 320471040 + }, + { + "id": "ipa_bpe_tokenizer", + "path": "2605_NemotronTTS_V0.2/bpe_ipa_tokenizer_2048_en_de_es_fr_hi_it_vi_zh_ko-KR_pt-BR_ar.json", + "sha256": "ba573df82176a963c56928e25166723e370e1af24f4aac70aa020c7bcf0ff0c5", + "bytes": 142076 + }, + { + "id": "english_sample_context_wav", + "path": "context_audios/english_sample.wav", + "sha256": "923f3265c8fc68940b88ce1e4bdc4860a579f447d90330b84bcb472d0e180466", + "bytes": 402286 + } + ] + }, + "converter": { + "repository": "https://github.com/styagi130/NeMo.git", + "commit": "bf562c17d255fc47ef41ce497fd4c2fb3b35f35d", + "script": "examples/tts/easymagpie_vllm_omni/scripts/convert_to_vllm.py", + "script_sha256": "6038d18a650f29a970b117e446a6db39908aa4470112c33717948ae7ec61ff62", + "image": "nvcr.io/nvidia/nemo@sha256:55f356677237cc7bcd318e505aa202022078fd4f0bb5ef527ed5d2ec8b3c5dd6", + "extra_pip_requirements": [ + "kaldialign==0.12.0" + ] + }, + "argv": [ + "python3", + "examples/tts/easymagpie_vllm_omni/scripts/convert_to_vllm.py", + "--nemo_file", + "/ckpt/2605_NemotronTTS_V0.2/v2/2605_EMTTS_SmallMamba_Step150k.nemo", + "--codec_model_path", + "/ckpt/25fps_spectral_codec_with_bandwidth_extension.nemo", + "--phoneme_tokenizer_path", + "/ckpt/2605_NemotronTTS_V0.2/bpe_ipa_tokenizer_2048_en_de_es_fr_hi_it_vi_zh_ko-KR_pt-BR_ar.json", + "--context_audio", + "/ckpt/context_audios/english_sample.wav", + "--speaker_name", + "english_sample", + "--context_audio_duration", + "5.0", + "--precompute_batch_size", + "1024", + "--dtype", + "bfloat16", + "--bundle_codec", + "--device", + "cuda", + "--outdir", + "/out" + ], + "outputs": [ + { + "path": "config.json", + "bytes": 1740, + "sha256": "5b0679ea0ec5d63b276ed40d17666d0b0fecbce4365595f078d323aa4278aaf2" + }, + { + "path": "model.safetensors", + "bytes": 2776868102, + "sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + { + "path": "model.safetensors.index.json", + "bytes": 71508, + "sha256": "a788e4e66aecf772059c68de035ca45d4ece163fbf9b63b2107c1461a5124eb5" + }, + { + "path": "tokenizer.json", + "bytes": 17077484, + "sha256": "623c34567aebb18582765289fbe23d901c62704d6518d71866e0e58db892b5b7" + }, + { + "path": "tokenizer_config.json", + "bytes": 177180, + "sha256": "50d5c40a1e06a86e8d0fc2d4c6a9bf73ed0ffca7bdad3c9b42cedca48baa25dc" + }, + { + "path": "special_tokens_map.json", + "bytes": 420, + "sha256": "e3a4f63da745f02317a45e00e6476c17fc66ac41faf14bb1b0be1f3211b0ca53" + }, + { + "path": "chat_template.jinja", + "bytes": 10504, + "sha256": "ab7813c3abdd9cb655905a410728b26c7884eca45ddfab8d9f931553485a7862" + }, + { + "path": "codec_native/config.json", + "bytes": 763, + "sha256": "3223ffaab2ae326acf9604e3c57c1477906d759242c209a242bdb2f9dc2a258a" + }, + { + "path": "codec_native/model.safetensors", + "bytes": 217238188, + "sha256": "7b855b3d8589a62ee9ce48462a1fbfe2f0fd44db55b4ecf5a36364a7735c9fc7" + }, + { + "path": "speaker_embeddings/english_sample.pt", + "bytes": 395034, + "sha256": "585ff1abb90e2200feedf9dbade84e5b64e0f4b2c7978482d1e2ef0005573401" + } + ] +} diff --git a/deploy/nemotron_native/profiles/playback-safe.json b/deploy/nemotron_native/profiles/playback-safe.json new file mode 100644 index 0000000..8a8bd68 --- /dev/null +++ b/deploy/nemotron_native/profiles/playback-safe.json @@ -0,0 +1,28 @@ +{ + "schema": "nemotron.playback-profile.v1", + "profiles": { + "1024": { + "physical_width": 1024, + "policy": { + "playback_reserve_guard": true, + "codec_deadline_priority": true, + "retire_batch": true + }, + "environment": { + "NEMOTRON_INITIAL_AUDIO_HOLD_MS": "40", + "NEMOTRON_MAX_ACTIVE_REQUESTS": null + } + }, + "2048": { + "physical_width": 2048, + "policy": { + "playback_reserve_guard": true, + "codec_deadline_priority": true + }, + "environment": { + "NEMOTRON_INITIAL_AUDIO_HOLD_MS": "300", + "NEMOTRON_MAX_ACTIVE_REQUESTS": "1024" + } + } + } +} diff --git a/deploy/nemotron_native/programs.json b/deploy/nemotron_native/programs.json new file mode 100644 index 0000000..8fe9645 --- /dev/null +++ b/deploy/nemotron_native/programs.json @@ -0,0 +1,1943 @@ +{ + "schema": "nemotron.native-program-catalog.v1", + "architecture": "sm_90a", + "startup": [ + 2, + 3, + 4, + 5 + ], + "steady": 8, + "programs": { + "1": { + "neural": { + "schema": "nemotron.native-readiness-plan.r22.v1", + "status": "PROPOSED_UNQUALIFIED", + "enabled": false, + "physical_width": 1, + "neural_cohort_partitioning": false, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "kernels": { + "prefill": "prefill", + "stage0": "stage0", + "codec": "codec" + }, + "execution": { + "mode": "serial", + "grid_policy": "overlap-only", + "stage0_grid": 96, + "full_stage0_grid": 96, + "codec_grid": 48 + }, + "generation": { + "cap": 511, + "startup_skip": true, + "chunk_frames": [ + 3, + 5, + 8 + ], + "contraction": false + }, + "state": { + "format": "FP32-private", + "recurrence_arithmetic": "FP32", + "canonical_binding": false, + "prefix_cache": false, + "initialization": "prefill-per-cohort" + }, + "hooks": [ + "predictor_active_mask", + "helper_graphs", + "controller_queues", + "admission_receipt" + ], + "id": "native-source-b1-recurrent-staging", + "source_build": 1 + }, + "serving": { + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "host": { + "tokens": "retained-engine-token-packing", + "serialization": "msgpack-c-1.2.2", + "pcm_conversion": "exact-avx2-with-scalar-fallback", + "backend_delivery_group": 64, + "batching_timer_ms": 0, + "deferred_audit": false, + "descriptor_ipc": true, + "torch_intra_threads": 4, + "torch_interop_threads": 4, + "tokio_threads": 8, + "backend_executor_threads": 1, + "cpu_policy": { + "server_physical_cores": 8, + "client_physical_cores": 4, + "include_smt_siblings": true, + "resolve_topology_at_launch": true + }, + "transport_options": { + "f5": true + }, + "fp32_prefix_snapshot": "immutable-fp32-v3" + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le", + "ownership": "immutable payload or retained slab lease until transport ACK", + "initial_singleton_readiness": true, + "readiness_bundle": "separate qualified physical-B1 bundle required", + "post_readiness_width": 1, + "incomplete_deadline_ms": 10000 + }, + "transport_compatibility": { + "identity_migration_enabled": false, + "execution_contract": "native_full_prefill_full_stage0_full_stage1_streaming_v1", + "output_contract": "pcm_f32_le_22050_mono_variable_streaming_v3", + "installed_transport_versions": { + "trt_version": "10.16.1.11", + "trtllm_version": "1.3.0rc25", + "trtllm_commit": "0b14c6bb0d2f2e9e6e8c76cf0667f0306ba68baa" + }, + "version_fields_describe_transport_image_not_neural_execution": true + }, + "schema": "nemotron.unified-serving.r22.v1", + "enabled": false, + "physical_width": 1, + "neural_cohort_partitioning": false, + "family": "readiness", + "id": "native-source-b1-recurrent-staging-streaming-v3" + }, + "libraries": { + "prefill": "programs/b1/artifacts/kernels/b1/prefill.so", + "stage0": "programs/b1/artifacts/kernels/b1/stage0.so", + "codec": "programs/b1/artifacts/kernels/b1/codec.so", + "stage0_inflight": "programs/b1/artifacts/kernels/b1/inflight.so" + } + }, + "2": { + "neural": { + "schema": "nemotron.small-width-plan.r22.v1", + "status": "BUILT_UNQUALIFIED", + "enabled": false, + "physical_width": 2, + "neural_cohort_partitioning": false, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "kernels": { + "prefill": "prefill", + "stage0": "stage0", + "codec": "codec" + }, + "execution": { + "mode": "serial", + "grid_policy": "static", + "stage0_grid": 132, + "codec_grid": 132 + }, + "resources": { + "occupancy_qualified": false, + "admission_checked_at_construction": true, + "compiled_arch": "sm_90a", + "producer_barrier": "actual-gridDim.x" + }, + "geometry": { + "physical_requests": 2, + "producer_tile_rows": 16, + "prefill_tile_rows": 64, + "sparse_expert_tile_rows": 16 + }, + "generation": { + "cap": 511, + "startup_skip": true, + "chunk_frames": [ + 3, + 5, + 8 + ], + "contraction": false + }, + "state": { + "format": "FP32-private", + "recurrence_arithmetic": "FP32", + "canonical_binding": false, + "prefix_cache": false, + "initialization": "prefill-per-cohort" + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le" + }, + "hooks": [ + "predictor_active_mask", + "helper_graphs", + "controller_queues", + "publication_receipt" + ], + "rng": { + "groups": 1, + "generated_rows_per_draw": 64, + "consumed_rows": 2, + "unused_rows": 62, + "context_draws": 1, + "all_physical_generations_advance": true, + "regrouping_equivalence": false + }, + "id": "native-source-b2", + "source_build": 1 + }, + "serving": { + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "host": { + "tokens": "retained-engine-token-packing", + "serialization": "msgpack-c-1.2.2", + "pcm_conversion": "exact-avx2-with-scalar-fallback", + "backend_delivery_group": 64, + "batching_timer_ms": 0, + "deferred_audit": false, + "descriptor_ipc": true, + "torch_intra_threads": 4, + "torch_interop_threads": 4, + "tokio_threads": 8, + "backend_executor_threads": 1, + "cpu_policy": { + "server_physical_cores": 8, + "client_physical_cores": 4, + "include_smt_siblings": true, + "resolve_topology_at_launch": true + }, + "transport_options": { + "f5": true + }, + "fp32_prefix_snapshot": "immutable-fp32-v3" + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le", + "ownership": "immutable payload or retained slab lease until transport ACK", + "initial_singleton_readiness": true, + "readiness_bundle": "separate qualified physical-B1 bundle required", + "post_readiness_width": 2, + "incomplete_deadline_ms": 10000 + }, + "transport_compatibility": { + "identity_migration_enabled": false, + "execution_contract": "native_full_prefill_full_stage0_full_stage1_streaming_v1", + "output_contract": "pcm_f32_le_22050_mono_variable_streaming_v3", + "installed_transport_versions": { + "trt_version": "10.16.1.11", + "trtllm_version": "1.3.0rc25", + "trtllm_commit": "0b14c6bb0d2f2e9e6e8c76cf0667f0306ba68baa" + }, + "version_fields_describe_transport_image_not_neural_execution": true + }, + "schema": "nemotron.unified-serving.r22.v1", + "enabled": false, + "physical_width": 2, + "neural_cohort_partitioning": false, + "family": "small", + "id": "native-source-b2-streaming-v3" + }, + "libraries": { + "stage0": "programs/b2/artifacts/kernels/stage0/native.so", + "prefill": "programs/b2/artifacts/kernels/prefill/native.so", + "codec": "programs/b2/artifacts/kernels/codec/native.so", + "stage0_inflight": "programs/b2/artifacts/kernels/stage0/inflight.so" + } + }, + "4": { + "neural": { + "schema": "nemotron.small-width-plan.r22.v1", + "status": "BUILT_UNQUALIFIED", + "enabled": false, + "physical_width": 4, + "neural_cohort_partitioning": false, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "kernels": { + "prefill": "prefill", + "stage0": "stage0", + "codec": "codec" + }, + "execution": { + "mode": "serial", + "grid_policy": "static", + "stage0_grid": 132, + "codec_grid": 132 + }, + "resources": { + "occupancy_qualified": false, + "admission_checked_at_construction": true, + "compiled_arch": "sm_90a", + "producer_barrier": "actual-gridDim.x" + }, + "geometry": { + "physical_requests": 4, + "producer_tile_rows": 16, + "prefill_tile_rows": 64, + "sparse_expert_tile_rows": 16 + }, + "generation": { + "cap": 511, + "startup_skip": true, + "chunk_frames": [ + 3, + 5, + 8 + ], + "contraction": false + }, + "state": { + "format": "FP32-private", + "recurrence_arithmetic": "FP32", + "canonical_binding": false, + "prefix_cache": false, + "initialization": "prefill-per-cohort" + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le" + }, + "hooks": [ + "predictor_active_mask", + "helper_graphs", + "controller_queues", + "publication_receipt" + ], + "rng": { + "groups": 1, + "generated_rows_per_draw": 64, + "consumed_rows": 4, + "unused_rows": 60, + "context_draws": 1, + "all_physical_generations_advance": true, + "regrouping_equivalence": false + }, + "id": "native-source-b4", + "source_build": 1 + }, + "serving": { + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "host": { + "tokens": "retained-engine-token-packing", + "serialization": "msgpack-c-1.2.2", + "pcm_conversion": "exact-avx2-with-scalar-fallback", + "backend_delivery_group": 64, + "batching_timer_ms": 0, + "deferred_audit": false, + "descriptor_ipc": true, + "torch_intra_threads": 4, + "torch_interop_threads": 4, + "tokio_threads": 8, + "backend_executor_threads": 1, + "cpu_policy": { + "server_physical_cores": 8, + "client_physical_cores": 4, + "include_smt_siblings": true, + "resolve_topology_at_launch": true + }, + "transport_options": { + "f5": true + }, + "fp32_prefix_snapshot": "immutable-fp32-v3" + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le", + "ownership": "immutable payload or retained slab lease until transport ACK", + "initial_singleton_readiness": true, + "readiness_bundle": "separate qualified physical-B1 bundle required", + "post_readiness_width": 4, + "incomplete_deadline_ms": 10000 + }, + "transport_compatibility": { + "identity_migration_enabled": false, + "execution_contract": "native_full_prefill_full_stage0_full_stage1_streaming_v1", + "output_contract": "pcm_f32_le_22050_mono_variable_streaming_v3", + "installed_transport_versions": { + "trt_version": "10.16.1.11", + "trtllm_version": "1.3.0rc25", + "trtllm_commit": "0b14c6bb0d2f2e9e6e8c76cf0667f0306ba68baa" + }, + "version_fields_describe_transport_image_not_neural_execution": true + }, + "schema": "nemotron.unified-serving.r22.v1", + "enabled": false, + "physical_width": 4, + "neural_cohort_partitioning": false, + "family": "small", + "id": "native-source-b4-streaming-v3" + }, + "libraries": { + "stage0": "programs/b4/artifacts/kernels/stage0/native.so", + "prefill": "programs/b4/artifacts/kernels/prefill/native.so", + "codec": "programs/b4/artifacts/kernels/codec/native.so", + "stage0_inflight": "programs/b4/artifacts/kernels/stage0/inflight.so" + } + }, + "8": { + "neural": { + "schema": "nemotron.small-width-plan.r22.v1", + "status": "BUILT_UNQUALIFIED", + "enabled": false, + "physical_width": 8, + "neural_cohort_partitioning": false, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "kernels": { + "prefill": "prefill", + "stage0": "stage0", + "codec": "codec" + }, + "execution": { + "mode": "serial", + "grid_policy": "static", + "stage0_grid": 132, + "codec_grid": 132 + }, + "resources": { + "occupancy_qualified": false, + "admission_checked_at_construction": true, + "compiled_arch": "sm_90a", + "producer_barrier": "actual-gridDim.x" + }, + "geometry": { + "physical_requests": 8, + "producer_tile_rows": 16, + "prefill_tile_rows": 64, + "sparse_expert_tile_rows": 16 + }, + "generation": { + "cap": 511, + "startup_skip": true, + "chunk_frames": [ + 3, + 5, + 8 + ], + "contraction": false + }, + "state": { + "format": "FP32-private", + "recurrence_arithmetic": "FP32", + "canonical_binding": false, + "prefix_cache": false, + "initialization": "prefill-per-cohort" + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le" + }, + "hooks": [ + "predictor_active_mask", + "helper_graphs", + "controller_queues", + "publication_receipt" + ], + "rng": { + "groups": 1, + "generated_rows_per_draw": 64, + "consumed_rows": 8, + "unused_rows": 56, + "context_draws": 1, + "all_physical_generations_advance": true, + "regrouping_equivalence": false + }, + "id": "native-source-b8", + "source_build": 1 + }, + "serving": { + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "host": { + "tokens": "retained-engine-token-packing", + "serialization": "msgpack-c-1.2.2", + "pcm_conversion": "exact-avx2-with-scalar-fallback", + "backend_delivery_group": 64, + "batching_timer_ms": 0, + "deferred_audit": false, + "descriptor_ipc": true, + "torch_intra_threads": 4, + "torch_interop_threads": 4, + "tokio_threads": 8, + "backend_executor_threads": 1, + "cpu_policy": { + "server_physical_cores": 8, + "client_physical_cores": 4, + "include_smt_siblings": true, + "resolve_topology_at_launch": true + }, + "transport_options": { + "f5": true + }, + "fp32_prefix_snapshot": "immutable-fp32-v3" + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le", + "ownership": "immutable payload or retained slab lease until transport ACK", + "initial_singleton_readiness": true, + "readiness_bundle": "separate qualified physical-B1 bundle required", + "post_readiness_width": 8, + "incomplete_deadline_ms": 10000 + }, + "transport_compatibility": { + "identity_migration_enabled": false, + "execution_contract": "native_full_prefill_full_stage0_full_stage1_streaming_v1", + "output_contract": "pcm_f32_le_22050_mono_variable_streaming_v3", + "installed_transport_versions": { + "trt_version": "10.16.1.11", + "trtllm_version": "1.3.0rc25", + "trtllm_commit": "0b14c6bb0d2f2e9e6e8c76cf0667f0306ba68baa" + }, + "version_fields_describe_transport_image_not_neural_execution": true + }, + "schema": "nemotron.unified-serving.r22.v1", + "enabled": false, + "physical_width": 8, + "neural_cohort_partitioning": false, + "family": "small", + "id": "native-source-b8-streaming-v3" + }, + "libraries": { + "stage0": "programs/b8/artifacts/kernels/stage0/native.so", + "prefill": "programs/b8/artifacts/kernels/prefill/native.so", + "codec": "programs/b8/artifacts/kernels/codec/native.so", + "stage0_inflight": "programs/b8/artifacts/kernels/stage0/inflight.so" + } + }, + "16": { + "neural": { + "schema": "nemotron.retained-narrow.r22.v1", + "enabled": false, + "physical_width": 16, + "neural_cohort_partitioning": false, + "execution": { + "mode": "overlap", + "grid_policy": "overlap-only", + "stage0_grid": 80, + "codec_grid": 52, + "full_stage0_grid": 132 + }, + "active_lanes": { + "mode": "disabled" + }, + "generation": { + "cap": 511, + "startup_skip": true, + "chunk_frames": [ + 3, + 5, + 8 + ] + }, + "state": { + "format": "FP32-private", + "recurrence": "FP32", + "prefix_cache": false + }, + "predictor": { + "full_width": "BF16", + "contracted": "BF16" + }, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "status": "PROPOSED_UNQUALIFIED", + "id": "native-source-b16", + "source_build": 1 + }, + "serving": { + "schema": "nemotron.unified-serving.r22.v1", + "enabled": false, + "physical_width": 16, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "host": { + "tokens": "retained-engine-token-packing", + "serialization": "msgpack-c-1.2.2", + "pcm_conversion": "exact-avx2-with-scalar-fallback", + "backend_delivery_group": 64, + "batching_timer_ms": 0, + "deferred_audit": false, + "descriptor_ipc": true, + "torch_intra_threads": 8, + "torch_interop_threads": 16, + "tokio_threads": 8, + "backend_executor_threads": 1, + "cpu_policy": { + "server_physical_cores": 8, + "client_physical_cores": 4, + "include_smt_siblings": true, + "resolve_topology_at_launch": true + }, + "transport_options": { + "f5": true + }, + "fp32_prefix_snapshot": "immutable-fp32-v3" + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le", + "ownership": "immutable payload or retained slab lease until transport ACK", + "initial_singleton_readiness": true, + "readiness_bundle": "separate qualified physical-B1 bundle required", + "post_readiness_width": 16, + "incomplete_deadline_ms": 10000 + }, + "transport_compatibility": { + "identity_migration_enabled": false, + "execution_contract": "native_full_prefill_full_stage0_full_stage1_streaming_v1", + "output_contract": "pcm_f32_le_22050_mono_variable_streaming_v3", + "installed_transport_versions": { + "trt_version": "10.16.1.11", + "trtllm_version": "1.3.0rc25", + "trtllm_commit": "0b14c6bb0d2f2e9e6e8c76cf0667f0306ba68baa" + }, + "version_fields_describe_transport_image_not_neural_execution": true + }, + "family": "narrow", + "id": "native-source-b16-streaming-v3" + }, + "libraries": { + "prefill": "programs/b16/artifacts/kernels/b16/prefill/native.so", + "stage0": "programs/b16/artifacts/kernels/b16/stage0/native.so", + "codec": "programs/b16/artifacts/kernels/b16/codec/libnative_stage1.so", + "stage0_inflight": "programs/b16/artifacts/kernels/b16/stage0/inflight.so" + } + }, + "32": { + "neural": { + "schema": "nemotron.retained-narrow.r22.v1", + "enabled": false, + "physical_width": 32, + "neural_cohort_partitioning": false, + "execution": { + "mode": "overlap", + "grid_policy": "overlap-only", + "stage0_grid": 80, + "codec_grid": 52, + "full_stage0_grid": 132 + }, + "active_lanes": { + "mode": "fixed-subset", + "target_batch": 16, + "after_chunk0": true, + "narrow_grid": 80, + "grid_policy": "follow-stage0" + }, + "generation": { + "cap": 511, + "startup_skip": true, + "chunk_frames": [ + 3, + 5, + 8 + ] + }, + "state": { + "format": "BF16-private", + "recurrence": "FP32", + "prefix_cache": true + }, + "predictor": { + "full_width": "FP8-FFN", + "contracted": "BF16" + }, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "status": "PROPOSED_UNQUALIFIED", + "id": "native-source-b32", + "source_build": 1 + }, + "serving": { + "schema": "nemotron.unified-serving.r22.v1", + "enabled": false, + "physical_width": 32, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "host": { + "tokens": "retained-engine-token-packing", + "serialization": "msgpack-c-1.2.2", + "pcm_conversion": "exact-avx2-with-scalar-fallback", + "backend_delivery_group": 64, + "batching_timer_ms": 0, + "deferred_audit": false, + "descriptor_ipc": true, + "torch_intra_threads": 8, + "torch_interop_threads": 16, + "tokio_threads": 8, + "backend_executor_threads": 1, + "cpu_policy": { + "server_physical_cores": 8, + "client_physical_cores": 4, + "include_smt_siblings": true, + "resolve_topology_at_launch": true + }, + "transport_options": { + "f5": true + } + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le", + "ownership": "immutable payload or retained slab lease until transport ACK", + "initial_singleton_readiness": true, + "readiness_bundle": "separate qualified physical-B1 bundle required", + "post_readiness_width": 32, + "incomplete_deadline_ms": 10000 + }, + "transport_compatibility": { + "identity_migration_enabled": false, + "execution_contract": "native_full_prefill_full_stage0_full_stage1_streaming_v1", + "output_contract": "pcm_f32_le_22050_mono_variable_streaming_v3", + "installed_transport_versions": { + "trt_version": "10.16.1.11", + "trtllm_version": "1.3.0rc25", + "trtllm_commit": "0b14c6bb0d2f2e9e6e8c76cf0667f0306ba68baa" + }, + "version_fields_describe_transport_image_not_neural_execution": true + }, + "family": "narrow", + "id": "native-source-b32-streaming-v3" + }, + "libraries": { + "prefill": "programs/b32/artifacts/kernels/b32/prefill/native.so", + "stage0": "programs/b32/artifacts/kernels/b32/stage0/native.so", + "codec": "programs/b32/artifacts/kernels/b32/codec/native.so", + "contraction_stage0": "programs/b32/artifacts/contraction/b32/stage0.so", + "contraction_transfer": "programs/b32/artifacts/contraction/b32/transfer.so", + "contraction_bf16": "programs/b32/artifacts/contraction/b32/bf16-native.so", + "stage0_inflight": "programs/b32/artifacts/kernels/b32/stage0/inflight.so" + } + }, + "64": { + "neural": { + "schema": "nemotron.retained-narrow.r22.v1", + "enabled": false, + "physical_width": 64, + "neural_cohort_partitioning": false, + "execution": { + "mode": "overlap", + "grid_policy": "overlap-only", + "stage0_grid": 80, + "codec_grid": 48, + "full_stage0_grid": 132 + }, + "active_lanes": { + "mode": "fixed-subset", + "target_batch": 16, + "after_chunk0": true, + "narrow_grid": 80, + "grid_policy": "follow-stage0" + }, + "generation": { + "cap": 511, + "startup_skip": true, + "chunk_frames": [ + 3, + 5, + 8 + ] + }, + "state": { + "format": "FP32-private", + "recurrence": "FP32", + "prefix_cache": false + }, + "predictor": { + "full_width": "BF16", + "contracted": "BF16" + }, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "status": "PROPOSED_UNQUALIFIED", + "id": "native-source-b64", + "source_build": 1 + }, + "serving": { + "schema": "nemotron.unified-serving.r22.v1", + "enabled": false, + "physical_width": 64, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "host": { + "tokens": "retained-engine-token-packing", + "serialization": "msgpack-c-1.2.2", + "pcm_conversion": "exact-avx2-with-scalar-fallback", + "backend_delivery_group": 64, + "batching_timer_ms": 0, + "deferred_audit": false, + "descriptor_ipc": true, + "torch_intra_threads": 8, + "torch_interop_threads": 16, + "tokio_threads": 8, + "backend_executor_threads": 1, + "cpu_policy": { + "server_physical_cores": 8, + "client_physical_cores": 4, + "include_smt_siblings": true, + "resolve_topology_at_launch": true + }, + "transport_options": { + "f5": true + }, + "fp32_prefix_snapshot": "immutable-fp32-v3" + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le", + "ownership": "immutable payload or retained slab lease until transport ACK", + "initial_singleton_readiness": true, + "readiness_bundle": "separate qualified physical-B1 bundle required", + "post_readiness_width": 64, + "incomplete_deadline_ms": 10000 + }, + "transport_compatibility": { + "identity_migration_enabled": false, + "execution_contract": "native_full_prefill_full_stage0_full_stage1_streaming_v1", + "output_contract": "pcm_f32_le_22050_mono_variable_streaming_v3", + "installed_transport_versions": { + "trt_version": "10.16.1.11", + "trtllm_version": "1.3.0rc25", + "trtllm_commit": "0b14c6bb0d2f2e9e6e8c76cf0667f0306ba68baa" + }, + "version_fields_describe_transport_image_not_neural_execution": true + }, + "family": "narrow", + "id": "native-source-b64-streaming-v3" + }, + "libraries": { + "prefill": "programs/b64/artifacts/kernels/b64/prefill/native.so", + "stage0": "programs/b64/artifacts/kernels/b64/stage0/native.so", + "codec": "programs/b64/artifacts/kernels/b64/codec/libnative_stage1.so", + "contraction_stage0": "programs/b64/artifacts/contraction/b64/stage0.so", + "contraction_transfer": "programs/b64/artifacts/contraction/b64/transfer.so", + "stage0_inflight": "programs/b64/artifacts/kernels/b64/stage0/inflight.so" + } + }, + "128": { + "neural": { + "schema": "nemotron.native-package-plan.r22.v1", + "status": "PROPOSED_UNQUALIFIED", + "enabled": false, + "physical_width": 128, + "neural_cohort_partitioning": false, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "kernels": { + "prefill": { + "artifact": "prefill", + "physical_width": 128, + "entrypoint": "mk2_native_prefill_run", + "workspace_sequences": 1 + }, + "stage0": { + "artifact": "stage0", + "physical_width": 128, + "entrypoint": "mk2_run_native_stage0", + "barrier_participants": 132, + "row_tiles": 2, + "row_tile": 64, + "k_tile": 128, + "pipeline_stages": 3 + }, + "codec": { + "artifact": "codec", + "physical_width": 128, + "entrypoint": "native_stage1_run", + "max_frames": 8, + "state_banks": 2 + } + }, + "generation": { + "cap": 511, + "startup_skip": true, + "chunk_frames": [ + 3, + 5, + 8 + ], + "contraction": false, + "rng": "independent64-lane-streams-with-retained-consumption" + }, + "state": { + "mode": 1, + "format": "canonical-FP32-private-BF16", + "recurrence_arithmetic": "FP32", + "scales_artifact": "state_scales", + "map_artifact": "state_map", + "map_kind": "identity", + "prefix_cache": "immutable-context-snapshot-with-qualified-key-and-lifetime" + }, + "host": { + "tokens": "pinned", + "serialization": "msgpack-c-1.2.2", + "pcm_conversion": "exact-avx2-with-scalar-fallback", + "backend_delivery_group": 64, + "batching_timer_ms": 0, + "deferred_audit": false, + "descriptor_ipc": true, + "torch_intra_threads": 8, + "torch_interop_threads": 16, + "tokio_threads": 8, + "backend_executor_threads": 1, + "cpu_policy": { + "server_physical_cores": 8, + "client_physical_cores": 4, + "include_smt_siblings": true, + "resolve_topology_at_launch": true + }, + "transport_options": { + "f5": true + } + }, + "execution": { + "mode": "serial", + "grid_policy": "overlap-only", + "stage0_grid": 132, + "full_stage0_grid": 132, + "codec_grid": 132 + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le", + "ownership": "immutable payload or retained slab lease until transport ACK", + "initial_singleton_readiness": true, + "readiness_bundle": "separate qualified physical-B1 bundle required", + "post_readiness_width": 128, + "incomplete_deadline_ms": 10000 + }, + "transport_compatibility": { + "identity_migration_enabled": false, + "execution_contract": "native_full_prefill_full_stage0_full_stage1_streaming_v1", + "output_contract": "pcm_f32_le_22050_mono_variable_streaming_v2", + "installed_transport_versions": { + "trt_version": "10.16.1.11", + "trtllm_version": "1.3.0rc25", + "trtllm_commit": "0b14c6bb0d2f2e9e6e8c76cf0667f0306ba68baa" + }, + "version_fields_describe_transport_image_not_neural_execution": true + }, + "id": "native-source-b128", + "source_build": 1 + }, + "serving": { + "schema": "nemotron.unified-serving.r22.v1", + "enabled": false, + "physical_width": 128, + "neural_cohort_partitioning": false, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "host": { + "tokens": "pinned", + "serialization": "msgpack-c-1.2.2", + "pcm_conversion": "exact-avx2-with-scalar-fallback", + "backend_delivery_group": 64, + "batching_timer_ms": 0, + "deferred_audit": false, + "descriptor_ipc": true, + "torch_intra_threads": 8, + "torch_interop_threads": 16, + "tokio_threads": 8, + "backend_executor_threads": 1, + "cpu_policy": { + "server_physical_cores": 8, + "client_physical_cores": 4, + "include_smt_siblings": true, + "resolve_topology_at_launch": true + }, + "transport_options": { + "f5": true + } + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le", + "ownership": "immutable payload or retained slab lease until transport ACK", + "initial_singleton_readiness": true, + "readiness_bundle": "separate qualified physical-B1 bundle required", + "post_readiness_width": 128, + "incomplete_deadline_ms": 10000 + }, + "transport_compatibility": { + "identity_migration_enabled": false, + "execution_contract": "native_full_prefill_full_stage0_full_stage1_streaming_v1", + "output_contract": "pcm_f32_le_22050_mono_variable_streaming_v3", + "installed_transport_versions": { + "trt_version": "10.16.1.11", + "trtllm_version": "1.3.0rc25", + "trtllm_commit": "0b14c6bb0d2f2e9e6e8c76cf0667f0306ba68baa" + }, + "version_fields_describe_transport_image_not_neural_execution": true + }, + "family": "wide", + "id": "native-source-b128-streaming-v3" + }, + "libraries": { + "prefill": "programs/b128/artifacts/kernels/b128/prefill.so", + "stage0": "programs/b128/artifacts/kernels/b128/stage0.so", + "codec": "programs/b128/artifacts/kernels/b128/codec.so", + "stage0_inflight": "programs/b128/artifacts/kernels/b128/inflight.so" + }, + "inflight_policy": { + "startup_phase_alignment": true, + "codec_first_packet_f2": false, + "startup_max_steps": 2, + "startup_budget_ms": 20.0, + "codec_merge_ready": false, + "partial_startup_single_step": false, + "serial_prefetch": true, + "codec_local_packing": true, + "codec_packed_transfer": true, + "defer_admission_reset": true + } + }, + "256": { + "neural": { + "schema": "nemotron.native-package-plan.r22.v1", + "status": "PROPOSED_UNQUALIFIED", + "enabled": false, + "physical_width": 256, + "neural_cohort_partitioning": false, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "kernels": { + "prefill": { + "artifact": "prefill", + "physical_width": 256, + "entrypoint": "mk2_native_prefill_run", + "workspace_sequences": 1 + }, + "stage0": { + "artifact": "stage0", + "physical_width": 256, + "entrypoint": "mk2_run_native_stage0", + "barrier_participants": 132, + "row_tiles": 4, + "row_tile": 64, + "k_tile": 128, + "pipeline_stages": 3 + }, + "codec": { + "artifact": "codec", + "physical_width": 256, + "entrypoint": "native_stage1_run", + "max_frames": 8, + "state_banks": 2 + } + }, + "generation": { + "cap": 511, + "startup_skip": true, + "chunk_frames": [ + 3, + 5, + 8 + ], + "contraction": false, + "rng": "independent64-lane-streams-with-retained-consumption" + }, + "state": { + "mode": 1, + "format": "canonical-FP32-private-BF16", + "recurrence_arithmetic": "FP32", + "scales_artifact": "state_scales", + "map_artifact": "state_map", + "map_kind": "identity", + "prefix_cache": "immutable-context-snapshot-with-qualified-key-and-lifetime" + }, + "host": { + "tokens": "pinned", + "serialization": "msgpack-c-1.2.2", + "pcm_conversion": "exact-avx2-with-scalar-fallback", + "backend_delivery_group": 64, + "batching_timer_ms": 0, + "deferred_audit": false, + "descriptor_ipc": true, + "torch_intra_threads": 8, + "torch_interop_threads": 16, + "tokio_threads": 8, + "backend_executor_threads": 1, + "cpu_policy": { + "server_physical_cores": 8, + "client_physical_cores": 4, + "include_smt_siblings": true, + "resolve_topology_at_launch": true + }, + "transport_options": { + "f5": true + } + }, + "execution": { + "mode": "serial", + "grid_policy": "overlap-only", + "stage0_grid": 132, + "full_stage0_grid": 132, + "codec_grid": 132 + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le", + "ownership": "immutable payload or retained slab lease until transport ACK", + "initial_singleton_readiness": true, + "readiness_bundle": "separate qualified physical-B1 bundle required", + "post_readiness_width": 256, + "incomplete_deadline_ms": 10000 + }, + "transport_compatibility": { + "identity_migration_enabled": false, + "execution_contract": "native_full_prefill_full_stage0_full_stage1_streaming_v1", + "output_contract": "pcm_f32_le_22050_mono_variable_streaming_v2", + "installed_transport_versions": { + "trt_version": "10.16.1.11", + "trtllm_version": "1.3.0rc25", + "trtllm_commit": "0b14c6bb0d2f2e9e6e8c76cf0667f0306ba68baa" + }, + "version_fields_describe_transport_image_not_neural_execution": true + }, + "id": "native-source-b256", + "source_build": 1 + }, + "serving": { + "schema": "nemotron.unified-serving.r22.v1", + "enabled": false, + "physical_width": 256, + "neural_cohort_partitioning": false, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "host": { + "tokens": "pinned", + "serialization": "msgpack-c-1.2.2", + "pcm_conversion": "exact-avx2-with-scalar-fallback", + "backend_delivery_group": 64, + "batching_timer_ms": 0, + "deferred_audit": false, + "descriptor_ipc": true, + "torch_intra_threads": 8, + "torch_interop_threads": 16, + "tokio_threads": 8, + "backend_executor_threads": 1, + "cpu_policy": { + "server_physical_cores": 8, + "client_physical_cores": 4, + "include_smt_siblings": true, + "resolve_topology_at_launch": true + }, + "transport_options": { + "f5": true + } + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le", + "ownership": "immutable payload or retained slab lease until transport ACK", + "initial_singleton_readiness": true, + "readiness_bundle": "separate qualified physical-B1 bundle required", + "post_readiness_width": 256, + "incomplete_deadline_ms": 10000 + }, + "transport_compatibility": { + "identity_migration_enabled": false, + "execution_contract": "native_full_prefill_full_stage0_full_stage1_streaming_v1", + "output_contract": "pcm_f32_le_22050_mono_variable_streaming_v3", + "installed_transport_versions": { + "trt_version": "10.16.1.11", + "trtllm_version": "1.3.0rc25", + "trtllm_commit": "0b14c6bb0d2f2e9e6e8c76cf0667f0306ba68baa" + }, + "version_fields_describe_transport_image_not_neural_execution": true + }, + "family": "wide", + "id": "native-source-b256-streaming-v3" + }, + "libraries": { + "prefill": "programs/b256/artifacts/kernels/b256/prefill.so", + "stage0": "programs/b256/artifacts/kernels/b256/stage0.so", + "codec": "programs/b256/artifacts/kernels/b256/codec.so", + "stage0_inflight": "programs/b256/artifacts/kernels/b256/inflight.so" + }, + "inflight_policy": { + "startup_phase_alignment": true, + "codec_first_packet_f2": false, + "startup_max_steps": 2, + "startup_budget_ms": 20.0, + "codec_merge_ready": false, + "partial_startup_single_step": false, + "serial_prefetch": true, + "codec_local_packing": true, + "codec_packed_transfer": true, + "defer_admission_reset": true + } + }, + "512": { + "neural": { + "schema": "nemotron.native-package-plan.r22.v1", + "status": "PROPOSED_UNQUALIFIED", + "enabled": false, + "physical_width": 512, + "neural_cohort_partitioning": false, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "kernels": { + "prefill": { + "artifact": "prefill", + "physical_width": 512, + "entrypoint": "mk2_native_prefill_run", + "workspace_sequences": 512, + "validation_kind": "retained-device-dedup" + }, + "stage0": { + "artifact": "stage0", + "physical_width": 512, + "entrypoint": "mk2_run_native_stage0", + "barrier_participants": 132, + "row_tiles": 8, + "row_tile": 64, + "k_tile": 128, + "pipeline_stages": 3 + }, + "codec": { + "artifact": "codec", + "physical_width": 512, + "entrypoint": "native_stage1_run", + "max_frames": 8, + "state_banks": 2 + } + }, + "generation": { + "cap": 511, + "startup_skip": true, + "chunk_frames": [ + 3, + 5, + 8 + ], + "contraction": false, + "rng": "independent64-lane-streams-with-retained-consumption" + }, + "state": { + "mode": 1, + "format": "canonical-FP32-private-BF16", + "recurrence_arithmetic": "FP32", + "scales_artifact": "state_scales", + "map_artifact": "state_map", + "map_kind": "identity", + "prefix_cache": "immutable-context-snapshot-with-qualified-key-and-lifetime" + }, + "host": { + "tokens": "pinned", + "serialization": "msgpack-c-1.2.2", + "pcm_conversion": "exact-avx2-with-scalar-fallback", + "backend_delivery_group": 64, + "batching_timer_ms": 0, + "deferred_audit": false, + "descriptor_ipc": true, + "torch_intra_threads": 8, + "torch_interop_threads": 16, + "tokio_threads": 8, + "backend_executor_threads": 1, + "cpu_policy": { + "server_physical_cores": 8, + "client_physical_cores": 4, + "include_smt_siblings": true, + "resolve_topology_at_launch": true + }, + "transport_options": { + "f5": true + } + }, + "execution": { + "mode": "serial", + "grid_policy": "overlap-only", + "stage0_grid": 132, + "full_stage0_grid": 132, + "codec_grid": 132 + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le", + "ownership": "immutable payload or retained slab lease until transport ACK", + "initial_singleton_readiness": true, + "readiness_bundle": "separate qualified physical-B1 bundle required", + "post_readiness_width": 512, + "incomplete_deadline_ms": 10000 + }, + "transport_compatibility": { + "identity_migration_enabled": false, + "execution_contract": "native_full_prefill_full_stage0_full_stage1_streaming_v1", + "output_contract": "pcm_f32_le_22050_mono_variable_streaming_v2", + "installed_transport_versions": { + "trt_version": "10.16.1.11", + "trtllm_version": "1.3.0rc25", + "trtllm_commit": "0b14c6bb0d2f2e9e6e8c76cf0667f0306ba68baa" + }, + "version_fields_describe_transport_image_not_neural_execution": true + }, + "id": "native-source-b512", + "source_build": 1 + }, + "serving": { + "schema": "nemotron.unified-serving.r22.v1", + "enabled": false, + "physical_width": 512, + "neural_cohort_partitioning": false, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "host": { + "tokens": "pinned", + "serialization": "msgpack-c-1.2.2", + "pcm_conversion": "exact-avx2-with-scalar-fallback", + "backend_delivery_group": 64, + "batching_timer_ms": 0, + "deferred_audit": false, + "descriptor_ipc": true, + "torch_intra_threads": 8, + "torch_interop_threads": 16, + "tokio_threads": 8, + "backend_executor_threads": 1, + "cpu_policy": { + "server_physical_cores": 8, + "client_physical_cores": 4, + "include_smt_siblings": true, + "resolve_topology_at_launch": true + }, + "transport_options": { + "f5": true + } + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le", + "ownership": "immutable payload or retained slab lease until transport ACK", + "initial_singleton_readiness": true, + "readiness_bundle": "separate qualified physical-B1 bundle required", + "post_readiness_width": 512, + "incomplete_deadline_ms": 10000 + }, + "transport_compatibility": { + "identity_migration_enabled": false, + "execution_contract": "native_full_prefill_full_stage0_full_stage1_streaming_v1", + "output_contract": "pcm_f32_le_22050_mono_variable_streaming_v3", + "installed_transport_versions": { + "trt_version": "10.16.1.11", + "trtllm_version": "1.3.0rc25", + "trtllm_commit": "0b14c6bb0d2f2e9e6e8c76cf0667f0306ba68baa" + }, + "version_fields_describe_transport_image_not_neural_execution": true + }, + "family": "wide", + "id": "native-source-b512-streaming-v3" + }, + "libraries": { + "prefill": "programs/b512/artifacts/kernels/b512/prefill.so", + "stage0": "programs/b512/artifacts/kernels/b512/stage0.so", + "codec": "programs/b512/artifacts/kernels/b512/codec.so", + "stage0_inflight": "programs/b512/artifacts/kernels/b512/inflight.so" + }, + "inflight_policy": { + "startup_phase_alignment": true, + "codec_first_packet_f2": false, + "startup_max_steps": 2, + "startup_budget_ms": 20.0, + "codec_merge_ready": false, + "partial_startup_single_step": false, + "serial_prefetch": true, + "codec_local_packing": true, + "codec_packed_transfer": true, + "defer_admission_reset": true + } + }, + "1024": { + "neural": { + "schema": "nemotron.native-package-plan.r22.v1", + "status": "PROPOSED_UNQUALIFIED", + "enabled": false, + "physical_width": 1024, + "neural_cohort_partitioning": false, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "kernels": { + "prefill": { + "artifact": "prefill", + "physical_width": 1024, + "entrypoint": "mk2_native_prefill_run", + "workspace_sequences": 1 + }, + "stage0": { + "artifact": "stage0", + "physical_width": 1024, + "entrypoint": "mk2_run_native_stage0", + "barrier_participants": 132, + "row_tiles": 16, + "row_tile": 64, + "k_tile": 128, + "pipeline_stages": 3 + }, + "codec": { + "artifact": "codec", + "physical_width": 1024, + "entrypoint": "native_stage1_run", + "max_frames": 8, + "state_banks": 2 + } + }, + "generation": { + "cap": 511, + "startup_skip": true, + "chunk_frames": [ + 3, + 5, + 8 + ], + "contraction": false, + "rng": "independent64-lane-streams-with-retained-consumption" + }, + "state": { + "mode": 1, + "format": "canonical-FP32-private-BF16", + "recurrence_arithmetic": "FP32", + "scales_artifact": "state_scales", + "map_artifact": "state_map", + "map_kind": "identity", + "prefix_cache": "immutable-context-snapshot-with-qualified-key-and-lifetime" + }, + "host": { + "tokens": "pinned", + "serialization": "msgpack-c-1.2.2", + "pcm_conversion": "exact-avx2-with-scalar-fallback", + "backend_delivery_group": 64, + "batching_timer_ms": 0, + "deferred_audit": false, + "descriptor_ipc": true, + "torch_intra_threads": 8, + "torch_interop_threads": 16, + "tokio_threads": 8, + "backend_executor_threads": 1, + "cpu_policy": { + "server_physical_cores": 8, + "client_physical_cores": 4, + "include_smt_siblings": true, + "resolve_topology_at_launch": true + }, + "transport_options": { + "f5": true + } + }, + "execution": { + "mode": "serial", + "grid_policy": "overlap-only", + "stage0_grid": 132, + "full_stage0_grid": 132, + "codec_grid": 132 + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le", + "ownership": "immutable payload or retained slab lease until transport ACK", + "initial_singleton_readiness": true, + "readiness_bundle": "separate qualified physical-B1 bundle required", + "post_readiness_width": 1024, + "incomplete_deadline_ms": 10000 + }, + "transport_compatibility": { + "identity_migration_enabled": false, + "execution_contract": "native_full_prefill_full_stage0_full_stage1_streaming_v1", + "output_contract": "pcm_f32_le_22050_mono_variable_streaming_v2", + "installed_transport_versions": { + "trt_version": "10.16.1.11", + "trtllm_version": "1.3.0rc25", + "trtllm_commit": "0b14c6bb0d2f2e9e6e8c76cf0667f0306ba68baa" + }, + "version_fields_describe_transport_image_not_neural_execution": true + }, + "id": "native-source-b1024", + "source_build": 1 + }, + "serving": { + "schema": "nemotron.unified-serving.r22.v1", + "enabled": false, + "physical_width": 1024, + "neural_cohort_partitioning": false, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "host": { + "tokens": "pinned", + "serialization": "msgpack-c-1.2.2", + "pcm_conversion": "exact-avx2-with-scalar-fallback", + "backend_delivery_group": 64, + "batching_timer_ms": 0, + "deferred_audit": false, + "descriptor_ipc": true, + "torch_intra_threads": 8, + "torch_interop_threads": 16, + "tokio_threads": 8, + "backend_executor_threads": 1, + "cpu_policy": { + "server_physical_cores": 8, + "client_physical_cores": 4, + "include_smt_siblings": true, + "resolve_topology_at_launch": true + }, + "transport_options": { + "f5": true + } + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le", + "ownership": "immutable payload or retained slab lease until transport ACK", + "initial_singleton_readiness": true, + "readiness_bundle": "separate qualified physical-B1 bundle required", + "post_readiness_width": 1024, + "incomplete_deadline_ms": 10000 + }, + "transport_compatibility": { + "identity_migration_enabled": false, + "execution_contract": "native_full_prefill_full_stage0_full_stage1_streaming_v1", + "output_contract": "pcm_f32_le_22050_mono_variable_streaming_v3", + "installed_transport_versions": { + "trt_version": "10.16.1.11", + "trtllm_version": "1.3.0rc25", + "trtllm_commit": "0b14c6bb0d2f2e9e6e8c76cf0667f0306ba68baa" + }, + "version_fields_describe_transport_image_not_neural_execution": true + }, + "family": "wide", + "id": "native-source-b1024-streaming-v3" + }, + "libraries": { + "prefill": "programs/b1024/artifacts/kernels/b1024/prefill.so", + "stage0": "programs/b1024/artifacts/kernels/b1024/stage0.so", + "codec": "programs/b1024/artifacts/kernels/b1024/codec.so", + "stage0_inflight": "programs/b1024/artifacts/kernels/b1024/inflight.so" + }, + "inflight_policy": { + "startup_phase_alignment": true, + "codec_first_packet_f2": false, + "startup_max_steps": 2, + "startup_budget_ms": 20.0, + "codec_merge_ready": false, + "partial_startup_single_step": false, + "serial_prefetch": true, + "codec_local_packing": true, + "codec_packed_transfer": true, + "defer_admission_reset": true + } + }, + "2048": { + "neural": { + "schema": "nemotron.native-package-plan.r22.v1", + "status": "PROPOSED_UNQUALIFIED", + "enabled": false, + "physical_width": 2048, + "neural_cohort_partitioning": false, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "kernels": { + "prefill": { + "artifact": "prefill", + "physical_width": 2048, + "entrypoint": "mk2_native_prefill_run", + "workspace_sequences": 1 + }, + "stage0": { + "artifact": "stage0", + "physical_width": 2048, + "entrypoint": "mk2_run_native_stage0", + "barrier_participants": 132, + "row_tiles": 32, + "row_tile": 64, + "k_tile": 128, + "pipeline_stages": 3 + }, + "codec": { + "artifact": "codec", + "physical_width": 2048, + "entrypoint": "native_stage1_run", + "max_frames": 8, + "state_banks": 2 + } + }, + "generation": { + "cap": 511, + "startup_skip": true, + "chunk_frames": [ + 3, + 5, + 8 + ], + "contraction": false, + "rng": "independent64-lane-streams-with-retained-consumption" + }, + "state": { + "mode": 1, + "format": "canonical-FP32-private-BF16", + "recurrence_arithmetic": "FP32", + "scales_artifact": "state_scales", + "map_artifact": "state_map", + "map_kind": "identity", + "prefix_cache": "immutable-context-snapshot-with-qualified-key-and-lifetime" + }, + "host": { + "tokens": "pinned", + "serialization": "msgpack-c-1.2.2", + "pcm_conversion": "exact-avx2-with-scalar-fallback", + "backend_delivery_group": 64, + "batching_timer_ms": 0, + "deferred_audit": false, + "descriptor_ipc": true, + "torch_intra_threads": 8, + "torch_interop_threads": 16, + "tokio_threads": 8, + "backend_executor_threads": 1, + "cpu_policy": { + "server_physical_cores": 8, + "client_physical_cores": 4, + "include_smt_siblings": true, + "resolve_topology_at_launch": true + }, + "transport_options": { + "f5": true + } + }, + "execution": { + "mode": "serial", + "grid_policy": "overlap-only", + "stage0_grid": 132, + "full_stage0_grid": 132, + "codec_grid": 132 + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le", + "ownership": "immutable payload or retained slab lease until transport ACK", + "initial_singleton_readiness": true, + "readiness_bundle": "separate qualified physical-B1 bundle required", + "post_readiness_width": 2048, + "incomplete_deadline_ms": 10000 + }, + "transport_compatibility": { + "identity_migration_enabled": false, + "execution_contract": "native_full_prefill_full_stage0_full_stage1_streaming_v1", + "output_contract": "pcm_f32_le_22050_mono_variable_streaming_v2", + "installed_transport_versions": { + "trt_version": "10.16.1.11", + "trtllm_version": "1.3.0rc25", + "trtllm_commit": "0b14c6bb0d2f2e9e6e8c76cf0667f0306ba68baa" + }, + "version_fields_describe_transport_image_not_neural_execution": true + }, + "id": "native-source-b2048", + "source_build": 1 + }, + "serving": { + "schema": "nemotron.unified-serving.r22.v1", + "enabled": false, + "physical_width": 2048, + "neural_cohort_partitioning": false, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "host": { + "tokens": "pinned", + "serialization": "msgpack-c-1.2.2", + "pcm_conversion": "exact-avx2-with-scalar-fallback", + "backend_delivery_group": 64, + "batching_timer_ms": 0, + "deferred_audit": false, + "descriptor_ipc": true, + "torch_intra_threads": 8, + "torch_interop_threads": 16, + "tokio_threads": 8, + "backend_executor_threads": 1, + "cpu_policy": { + "server_physical_cores": 8, + "client_physical_cores": 4, + "include_smt_siblings": true, + "resolve_topology_at_launch": true + }, + "transport_options": { + "f5": true + } + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le", + "ownership": "immutable payload or retained slab lease until transport ACK", + "initial_singleton_readiness": true, + "readiness_bundle": "separate qualified physical-B1 bundle required", + "post_readiness_width": 2048, + "incomplete_deadline_ms": 10000 + }, + "transport_compatibility": { + "identity_migration_enabled": false, + "execution_contract": "native_full_prefill_full_stage0_full_stage1_streaming_v1", + "output_contract": "pcm_f32_le_22050_mono_variable_streaming_v3", + "installed_transport_versions": { + "trt_version": "10.16.1.11", + "trtllm_version": "1.3.0rc25", + "trtllm_commit": "0b14c6bb0d2f2e9e6e8c76cf0667f0306ba68baa" + }, + "version_fields_describe_transport_image_not_neural_execution": true + }, + "family": "wide", + "id": "native-source-b2048-streaming-v3" + }, + "libraries": { + "prefill": "programs/b2048/artifacts/kernels/b2048/prefill.so", + "stage0": "programs/b2048/artifacts/kernels/b2048/stage0.so", + "codec": "programs/b2048/artifacts/kernels/b2048/codec.so", + "stage0_inflight": "programs/b2048/artifacts/kernels/b2048/inflight.so" + }, + "inflight_policy": { + "startup_phase_alignment": true, + "codec_first_packet_f2": false, + "startup_max_steps": 2, + "startup_budget_ms": 20.0, + "codec_merge_ready": false, + "partial_startup_single_step": false, + "serial_prefetch": true, + "codec_local_packing": true, + "codec_packed_transfer": true, + "defer_admission_reset": true + } + }, + "4096": { + "neural": { + "schema": "nemotron.native-package-plan.r22.v1", + "status": "PROPOSED_UNQUALIFIED", + "enabled": false, + "physical_width": 4096, + "neural_cohort_partitioning": false, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "kernels": { + "prefill": { + "artifact": "prefill", + "physical_width": 4096, + "entrypoint": "mk2_native_prefill_run", + "workspace_sequences": 1 + }, + "stage0": { + "artifact": "stage0", + "physical_width": 4096, + "entrypoint": "mk2_run_native_stage0", + "barrier_participants": 132, + "row_tiles": 64, + "row_tile": 64, + "k_tile": 128, + "pipeline_stages": 3 + }, + "codec": { + "artifact": "codec", + "physical_width": 4096, + "entrypoint": "native_stage1_run", + "max_frames": 8, + "state_banks": 2 + } + }, + "generation": { + "cap": 511, + "startup_skip": true, + "chunk_frames": [ + 3, + 5, + 8 + ], + "contraction": false, + "rng": "independent64-lane-streams-with-retained-consumption" + }, + "state": { + "mode": 1, + "format": "canonical-FP32-private-BF16", + "recurrence_arithmetic": "FP32", + "scales_artifact": "state_scales", + "map_artifact": "state_map", + "map_kind": "identity", + "prefix_cache": "immutable-context-snapshot-with-qualified-key-and-lifetime" + }, + "host": { + "tokens": "pinned", + "serialization": "msgpack-c-1.2.2", + "pcm_conversion": "exact-avx2-with-scalar-fallback", + "backend_delivery_group": 64, + "batching_timer_ms": 0, + "deferred_audit": false, + "descriptor_ipc": true, + "torch_intra_threads": 8, + "torch_interop_threads": 16, + "tokio_threads": 8, + "backend_executor_threads": 1, + "cpu_policy": { + "server_physical_cores": 8, + "client_physical_cores": 4, + "include_smt_siblings": true, + "resolve_topology_at_launch": true + }, + "transport_options": { + "f5": true + } + }, + "execution": { + "mode": "serial", + "grid_policy": "overlap-only", + "stage0_grid": 132, + "full_stage0_grid": 132, + "codec_grid": 132 + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le", + "ownership": "immutable payload or retained slab lease until transport ACK", + "initial_singleton_readiness": true, + "readiness_bundle": "separate qualified physical-B1 bundle required", + "post_readiness_width": 4096, + "incomplete_deadline_ms": 10000 + }, + "transport_compatibility": { + "identity_migration_enabled": false, + "execution_contract": "native_full_prefill_full_stage0_full_stage1_streaming_v1", + "output_contract": "pcm_f32_le_22050_mono_variable_streaming_v2", + "installed_transport_versions": { + "trt_version": "10.16.1.11", + "trtllm_version": "1.3.0rc25", + "trtllm_commit": "0b14c6bb0d2f2e9e6e8c76cf0667f0306ba68baa" + }, + "version_fields_describe_transport_image_not_neural_execution": true + }, + "id": "native-source-b4096", + "source_build": 1 + }, + "serving": { + "schema": "nemotron.unified-serving.r22.v1", + "enabled": false, + "physical_width": 4096, + "neural_cohort_partitioning": false, + "model": { + "checkpoint_sha256": "8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa" + }, + "host": { + "tokens": "pinned", + "serialization": "msgpack-c-1.2.2", + "pcm_conversion": "exact-avx2-with-scalar-fallback", + "backend_delivery_group": 64, + "batching_timer_ms": 0, + "deferred_audit": false, + "descriptor_ipc": true, + "torch_intra_threads": 8, + "torch_interop_threads": 16, + "tokio_threads": 8, + "backend_executor_threads": 1, + "cpu_policy": { + "server_physical_cores": 8, + "client_physical_cores": 4, + "include_smt_siblings": true, + "resolve_topology_at_launch": true + }, + "transport_options": { + "f5": true + } + }, + "publication": { + "continuation_proof_required": true, + "sample_rate": 22050, + "neural_output": "pcm_f32_le", + "http_output": "pcm_s16_le", + "ownership": "immutable payload or retained slab lease until transport ACK", + "initial_singleton_readiness": true, + "readiness_bundle": "separate qualified physical-B1 bundle required", + "post_readiness_width": 4096, + "incomplete_deadline_ms": 10000 + }, + "transport_compatibility": { + "identity_migration_enabled": false, + "execution_contract": "native_full_prefill_full_stage0_full_stage1_streaming_v1", + "output_contract": "pcm_f32_le_22050_mono_variable_streaming_v3", + "installed_transport_versions": { + "trt_version": "10.16.1.11", + "trtllm_version": "1.3.0rc25", + "trtllm_commit": "0b14c6bb0d2f2e9e6e8c76cf0667f0306ba68baa" + }, + "version_fields_describe_transport_image_not_neural_execution": true + }, + "family": "wide", + "id": "native-source-b4096-streaming-v3" + }, + "libraries": { + "prefill": "programs/b4096/artifacts/kernels/b4096/prefill.so", + "stage0": "programs/b4096/artifacts/kernels/b4096/stage0.so", + "codec": "programs/b4096/artifacts/kernels/b4096/codec.so", + "stage0_inflight": "programs/b4096/artifacts/kernels/b4096/inflight.so" + } + } + } +} diff --git a/deploy/nemotron_native/serve.py b/deploy/nemotron_native/serve.py new file mode 100644 index 0000000..e4052f2 --- /dev/null +++ b/deploy/nemotron_native/serve.py @@ -0,0 +1,233 @@ +from __future__ import annotations + +import argparse +import asyncio +import hashlib +import json +import signal +from pathlib import Path + + +def verify_model(root, seal, identity): + if seal["model_identity"] != identity: + raise RuntimeError("Deployment identity mismatch") + for name, row in seal["files"].items(): + p = (root / name).resolve(strict=True) + if not p.is_relative_to(root.resolve(strict=True)): + raise RuntimeError("Model path escapes root") + h = hashlib.sha256() + with p.open("rb") as f: + for block in iter(lambda: f.read(4 * 1024 * 1024), b""): + h.update(block) + if h.hexdigest() != row["sha256"] or p.stat().st_size != row["bytes"]: + raise RuntimeError("Model source mismatch: " + name) + + +def exercise_inflight_lifecycle(engine): + if any(slot is not None for slot in engine.slots): + raise RuntimeError("In-flight lifecycle checks require an empty owner") + text = "Hello from the native three megakernel runner." + request = engine.engine._requests([text] * engine.batch, None, "english_sample")[0] + + def collect(slot, cancel): + epoch = engine.epochs[slot] + 1 + identity = f"native-lifecycle-{slot}-{epoch}" + digest = hashlib.sha256() + chunks = [] + terminal = None + try: + engine.admit( + slot, epoch, text=text, token_ids=request.target_ids, seed=0, request_id=identity + ) + for _ in range(engine.steps + 16): + for event in engine.advance([slot]): + if ( + event["slot"] != slot + or event["epoch"] != epoch + or event["request_id"] != identity + ): + raise RuntimeError("Lifecycle execution changed request ownership") + if event["type"] == "chunk": + digest.update(event["pcm_f32_le"]) + chunks.append(event["samples"]) + elif event["type"] == "finish": + terminal = event + if cancel and chunks: + terminal = engine.cancel(slot, epoch) + break + if terminal is not None: + break + if terminal is None or not chunks: + raise RuntimeError("Lifecycle request did not complete or produce audio") + return { + "batch": engine.batch, + "mode": "request-local-inflight", + "slot": slot, + "epoch": epoch, + "cancel_after_first_chunk": cancel, + "pcm_sha256": digest.hexdigest(), + "chunk_samples": chunks, + "finish_reason": terminal["reason"], + "generations": terminal["receipt"]["generations"], + } + finally: + owner = engine.slots[slot] + if owner is not None and owner.epoch == epoch and not engine.poisoned: + if not owner.finished: + engine.cancel(slot, epoch) + engine.retire(slot, epoch) + + baseline = collect(0, False) + checks = [baseline, collect(0, True), collect(0, False), collect(engine.batch - 1, False)] + keys = ("pcm_sha256", "chunk_samples", "finish_reason", "generations") + for check in checks: + if not check["cancel_after_first_chunk"]: + if any(check[key] != baseline[key] for key in keys): + raise RuntimeError("In-flight exact replay failed after cancel or slot reuse") + check["exact_replay"] = True + else: + check["exact_replay"] = None + if any(slot is not None for slot in engine.slots): + raise RuntimeError("In-flight lifecycle checks did not release every slot") + return checks + + +async def serve(args): + from nemotron_native_backend.execution.program import ( + load_program, + load_serving_plan, + prepare_for_qualification, + ) + from nemotron_native_backend.execution.readiness import make_readiness_binding + + family = json.loads(args.plan.read_text())["family"] + neural = load_program(args.neural_plan, family) + plan = load_serving_plan(args.plan, neural, family) + identity = plan.contract["model"]["checkpoint_sha256"] + verify_model(args.model, json.loads(args.model_seal.read_text()), identity) + pins = json.loads(args.source_seal.read_text())["package_source_pins"] + readiness = make_readiness_binding( + args.readiness, + args.readiness_artifacts, + args.model, + source_pins=pins, + wide_plan=plan, + model_identity=identity, + qualification_only=True, + ) + args.receipts.mkdir(parents=True, exist_ok=True) + args.socket.parent.mkdir(parents=True, exist_ok=True) + worker = await prepare_for_qualification( + plan, + neural, + args.neural_plan, + args.artifacts, + args.model, + source_pins=pins, + model_identity=identity, + readiness=readiness, + socket_path=args.socket, + receipts=args.receipts, + qualification_only=True, + ) + try: + if args.exercise_lifecycle: + keys = ( + "batch", + "executed_steps", + "first_eos_generations", + "samples", + "finish_reasons", + "pcm_sha256", + "active_codes_sha256", + "active_phonemes_sha256", + "live_codes_sha256", + ) + report = {"status": "RUNNING", "checks": []} + loop = asyncio.get_running_loop() + + def collect(engine, cancel): + it = engine.stream( + ["Hello from the native three megakernel runner."] * engine.batch, + seed=0, + epoch=1, + ) + try: + for event in it: + if cancel and event["type"] == "chunk": + break + finally: + it.close() + + inflight = getattr(worker.backend, "inflight_engine", None) + for width, engine in worker.backend.engines.items(): + if inflight is not None and engine is inflight.engine: + checks = await loop.run_in_executor( + worker.backend.executor, exercise_inflight_lifecycle, inflight + ) + report["checks"].extend(checks) + worker.batcher._epochs = list(inflight.epochs) + continue + expected = {k: engine.last_receipt[k] for k in keys} + for cancel in (False, True, False): + await loop.run_in_executor( + worker.backend.executor, lambda e=engine, c=cancel: collect(e, c) + ) + if not cancel and {k: engine.last_receipt[k] for k in keys} != expected: + raise RuntimeError( + "Exact replay failed after construction/cancel atB" + str(width) + ) + report["checks"].append( + { + "batch": width, + "cancel_after_first_chunk": cancel, + "exact_replay": None if cancel else True, + } + ) + report["status"] = "PASS_NATIVE_FACTORY_LIFECYCLE" + (args.receipts / "LIFECYCLE.json").write_text(json.dumps(report, indent=2) + "\n") + stopping = asyncio.Event() + loop = asyncio.get_running_loop() + for sig in (signal.SIGINT, signal.SIGTERM): + loop.add_signal_handler(sig, stopping.set) + await worker.server.start() + receipt = dict( + status="QUALIFICATION_SERVER_LISTENING", + physical_widths=sorted(worker.backend.engines), + neural_plan_sha256=neural.manifest_sha256, + serving_plan_sha256=plan.manifest_sha256, + production_enabled=False, + ) + (args.receipts / "NATIVE_READY.json").write_text(json.dumps(receipt, indent=2) + "\n") + print(json.dumps(receipt), flush=True) + await stopping.wait() + finally: + await worker.server.shutdown() + await worker.backend.shutdown() + + +def main(): + p = argparse.ArgumentParser() + for key in ( + "neural-plan", + "plan", + "readiness", + "readiness-artifacts", + "artifacts", + "model", + "model-seal", + "source-seal", + "receipts", + "socket", + ): + p.add_argument("--" + key, type=Path, required=True) + p.add_argument("--qualification-only", action="store_true") + p.add_argument("--exercise-lifecycle", action="store_true") + args = p.parse_args() + if not args.qualification_only: + p.error("--qualification-only is required") + asyncio.run(serve(args)) + + +if __name__ == "__main__": + main() diff --git a/docker/Dockerfile.nemotron b/docker/Dockerfile.nemotron new file mode 100644 index 0000000..b375a09 --- /dev/null +++ b/docker/Dockerfile.nemotron @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +ARG NATIVE_RUNTIME=nvcr.io/nvidia/pytorch@sha256:222d8b18e671be5c3ef91cb41727a2572a0b23f59ded6c39f373a96946f6f2ba +ARG NORMALIZER_BUILDER=gcr.io/bazel-public/bazel:6.2.1@sha256:e25493c1c3bf74e0612a34cf6ae23c67243117aa1c42bc7730f46effd0a89006 +ARG TEXT_PREPARATION_BUILDER=rust:1.88.0-bookworm@sha256:af306cfa71d987911a781c37b59d7d67d934f49684058f96cf72079c3626bfe0 + +FROM ${TEXT_PREPARATION_BUILDER} AS text-preparation-builder +WORKDIR /build +COPY native/text_preparation/ ./ +ARG NATIVE_BUILD_JOBS=2 +RUN --mount=type=cache,target=/usr/local/cargo/registry,sharing=locked \ + cargo test --release --locked --workspace --jobs "${NATIVE_BUILD_JOBS}" \ + && cargo build --release --locked --jobs "${NATIVE_BUILD_JOBS}" \ + && install -D --mode=0555 target/release/libnemotron_text_preparation.so /out/libnemotron_text_preparation.so + +FROM ${NORMALIZER_BUILDER} AS normalizer-builder +USER root +RUN apt-get update && apt-get install --yes --no-install-recommends libboost-dev && rm -rf /var/lib/apt/lists/* +WORKDIR /build +COPY native/tts-normalizer/ ./ +ARG NATIVE_BUILD_JOBS=2 +RUN --mount=type=cache,target=/root/.cache/bazel,sharing=locked \ + bazel test --config=release --jobs="${NATIVE_BUILD_JOBS}" //:normalize_test //:speech_hints_test //:text_preprocessor_test //:text_preprocessor_c_test //:string_matcher_test \ + && bazel build --config=release --jobs="${NATIVE_BUILD_JOBS}" //:text_preprocessor_c \ + && install -D --mode=0555 bazel-bin/libtext_preprocessor_c.so /out/libnemotron_speech_tts_text_preprocessor.so + +FROM ${NATIVE_RUNTIME} AS cpp-builder +WORKDIR /source +COPY . . +ARG NATIVE_BUILD_JOBS=2 +RUN cmake -S . -B build -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_INSTALL_PREFIX=/opt/nemo-speech \ + -DNEMO_SPEECH_BUILD_ASR=OFF -DNEMO_SPEECH_BUILD_DIAR=OFF \ + -DNEMO_SPEECH_BUILD_TTS=ON -DNEMO_SPEECH_BUILD_HTTP=ON \ + -DNEMO_SPEECH_BUILD_MIC_CAPTURE=OFF -DNEMO_SPEECH_BUILD_TESTS=ON \ + -DGGML_CUDA=OFF -DGGML_NATIVE=OFF \ + && cmake --build build --parallel "${NATIVE_BUILD_JOBS}" \ + && ctest --test-dir build --output-on-failure \ + && cmake --install build + +FROM ${NATIVE_RUNTIME} AS kernel-builder +WORKDIR /source +COPY deploy/nemotron_native /source/deploy/nemotron_native +COPY scripts/nemotron_native /source/scripts/nemotron_native +COPY python/nemotron_native_backend /source/python/nemotron_native_backend +ARG NATIVE_WIDTHS=1,32,512 +ARG NATIVE_BUILD_JOBS=2 +ARG NATIVE_BATCHING_MODE=inflight +ARG NATIVE_INFLIGHT_PROFILE=retained +ENV LIBRARY_PATH=/usr/local/cuda/lib64/stubs +RUN python3 scripts/nemotron_native/build_kernels.py --output /kernel-build --widths "${NATIVE_WIDTHS}" --jobs "${NATIVE_BUILD_JOBS}" \ + && python3 scripts/nemotron_native/build_package.py --kernels /kernel-build --output /package --widths "${NATIVE_WIDTHS}" --batching-mode "${NATIVE_BATCHING_MODE}" --inflight-profile "${NATIVE_INFLIGHT_PROFILE}" + +FROM ${NATIVE_RUNTIME} AS runtime +USER root +ENV PATH=/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin \ + LD_LIBRARY_PATH=/usr/local/lib/python3.12/dist-packages/torch/lib:/usr/local/cuda/compat/lib:/usr/local/nvidia/lib:/usr/local/nvidia/lib64 \ + TRT_VERSION= TRTOSS_VERSION= POLYGRAPHY_VERSION= +COPY scripts/nemotron_native/runtime_dependencies.py /opt/nemotron-native/runtime_dependencies.py +RUN python3 /opt/nemotron-native/runtime_dependencies.py --remove-tensorrt --output /opt/nemotron-native/dependencies.json \ + && python3 -m pip install --no-cache-dir msgpack==1.2.2 numpy==2.3.5 safetensors==0.8.0 tokenizers==0.23.1 +COPY --from=cpp-builder /opt/nemo-speech /opt/nemo-speech +COPY --from=kernel-builder /package /package +COPY --from=normalizer-builder /out/libnemotron_speech_tts_text_preprocessor.so /opt/nemotron-speech/lib/processors/libnemotron_speech_tts_text_preprocessor.so +COPY --from=text-preparation-builder /out/libnemotron_text_preparation.so /opt/nemotron-speech/lib/processors/libnemotron_text_preparation.so +COPY python /opt/nemotron-speech/python +COPY deploy/nemotron_native/serve.py /opt/nemotron-speech/deploy/nemotron_native/serve.py +COPY deploy/nemotron_native/LICENSE /opt/nemotron-speech/deploy/nemotron_native/LICENSE +COPY scripts/nemotron_native/launch.py /opt/nemotron-speech/scripts/nemotron_native/launch.py +COPY LICENSE NOTICE THIRD_PARTY_NOTICES.md /opt/nemo-speech/share/licenses/nemo-speech/ +COPY deploy/nemotron_native/LICENSE /opt/nemo-speech/share/licenses/nemo-speech/third_party/nemotron-native/LICENSE +COPY native/tts-normalizer/LICENSE.legacy native/tts-normalizer/LICENSE.riva native/tts-normalizer/THIRD_PARTY_NOTICES.md \ + /opt/nemo-speech/share/licenses/nemo-speech/third_party/nemotron-normalizer/ +COPY native/text_preparation/LICENSE.imported native/text_preparation/NOTICE \ + /opt/nemo-speech/share/licenses/nemo-speech/third_party/nemotron-text-preparation/ +COPY deploy/nemotron_native/kernels/third_party/thunderkittens/LICENSE \ + /opt/nemo-speech/share/licenses/nemo-speech/third_party/thunderkittens/LICENSE +COPY deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/LICENSE \ + /opt/nemo-speech/share/licenses/nemo-speech/third_party/thunderkittens-v4-prefill/LICENSE +ENV PATH=/opt/nemo-speech/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin \ + PYTHONPATH=/opt/nemotron-speech/python \ + PYTHONDONTWRITEBYTECODE=1 PYTHONNOUSERSITE=1 PYTHONUNBUFFERED=1 TOKENIZERS_PARALLELISM=false \ + NVIDIA_DRIVER_CAPABILITIES=compute,utility NVIDIA_TF32_OVERRIDE=0 CUDA_MODULE_LOADING=LAZY \ + NEMOTRON_TN_STRING_MATCHER=1 \ + NEMOTRON_NATIVE_PREPARATION_LIBRARY=/opt/nemotron-speech/lib/processors/libnemotron_text_preparation.so \ + NEMOTRON_FRONTEND_PREPARATION=1 \ + NEMOTRON_BRIDGE_FORWARD_F32=1 \ + NEMOTRON_BRIDGE_FAIR_BATCH=64 \ + NEMOTRON_HTTP_TCP_NODELAY=1 +ENTRYPOINT ["python3", "/opt/nemotron-speech/scripts/nemotron_native/launch.py"] diff --git a/docs/nemotron-integration-results.md b/docs/nemotron-integration-results.md new file mode 100644 index 0000000..212212f --- /dev/null +++ b/docs/nemotron-integration-results.md @@ -0,0 +1,591 @@ +# Nemotron native backend integration results + +This report records local integration checks and performance samples from +October 15, 2026. The C++ frontend successfully serves the retained native +Nemotron backend, including concurrent streaming and cancellation. Initial +B512 host overhead and its subsequent repairs are retained below. The latest +matched direct-worker comparison measures 205.04 ms in-flight TTFA p95; it does +not establish the 200-ms target, production readiness, or fresh acoustic quality. + +Use the [build and run guide](nemotron-megakernel.md) for reproducible setup and +the [measurement instructions](nemotron-megakernel.md#validate-and-measure) for +client commands and metric definitions. Model weights, tokenizer assets, +normalization grammars, and generated validation receipts are supplied +separately from the repository. + +## Architecture and provenance + +The request path is C++ HTTP frontend → C++ synthesizer → Python bridge → +retained Python native controller → CUDA megakernels. The legacy Rust frontend +is not part of this request path. The bridge runs outside the sealed native +package, so frontend integration does not alter that package's source inventory. + +The C++ frontend multiplexes requests over one Unix socket using bounded +per-request audio queues. The bridge validates the tokenizer against the model +seal, applies the retained native English normalizer when configured, tokenizes +without added special tokens, and appends EOS `131073`. It submits the resulting +text and IDs through the controller's existing MessagePack transport. Native +PCM slabs are validated and copied before acknowledgment; the default bridge +path converts them to PCM16 with the legacy asymmetric scaling, rounding, and +saturation rules. The native controller retains GPU state, scheduling, kernel +dispatch, codec execution, and slot lifecycle ownership. + +| Component | Revision or identity | +|---|---| +| C++ repository base | `083861c241061562b8ec29eb06d7f6cdc3af8775` | +| Retained native source | `2cc27a4e6b06e5e6bf33020ce4faa1fc00bfd25b` | +| Model checkpoint SHA-256 | `8bb6bfeb9cd4335f014a70bf87059af58945399be454a4c4e353a983d09eeffa` | +| Tokenizer SHA-256 | `623c34567aebb18582765289fbe23d901c62704d6518d71866e0e58db892b5b7` | +| Ten-text benchmark corpus SHA-256 | `8bae1d20bfd5d17bee363afdbd6d25e0167d4bcc5c11ea6663d75c3b8d9450a2` | + +A fresh byte comparison verified **579 imported files unchanged** against the +retained source revision: 468 kernel files including the recipe inventory, 65 +Python runtime files, 38 native normalizer files, and eight deployment/build +files. The two additional copies of the retained root license also match their +source. The inventory counts each destination once and preserves the vendor +licenses within the imported trees. Recompilation does not imply byte-identical +shared libraries. + +These components retain their original terms; see the +[license scope](../THIRD_PARTY_NOTICES.md#imported-nemotron-native-deployment-closure). +Local qualification does not resolve redistribution approval. + +The tested GPU was an H100 80 GB HBM3 with 132 SMs, compute capability 9.0, +without MIG, using `sm_90a` kernels. The v3 runtime reported CUDA 13.2, +PyTorch `2.12.0a0+5aff3928d8.nv26.05`, tokenizers 0.23.1, MessagePack 1.2.2, +and NumPy 2.3.5. TensorRT and TensorRT-LLM imports were unavailable in that +runtime. The [Docker recipe](../docker/Dockerfile.nemotron) pins its base images +and builds the frontend, normalizer, kernels, and native package from source. + +## Qualification completed + +- The native worker reached qualification readiness at widths 1, 32, and 512. + Its readiness receipt explicitly reports `production_enabled: false`. +- B1 returned two successful HTTP PCM streams with identical bytes and sample + counts on repeated input. This is repeatability within the new stack, not a + cross-frontend acoustic comparison. +- At B32 and B512, native lifecycle checks passed EOS completion, cancellation + after the first chunk, exact replay after slot reuse, and replay in the last + physical slot. The replay PCM hashes and chunk sample counts matched within + each width. +- Eighteen bridge/launcher unit tests passed, covering multiplexing, token + preparation, copy-before-ACK ownership, malformed output, slow-consumer + failure, cancellation, tokenizer identity, PCM conversion, and bounded child + shutdown. HTTP tests additionally exercise streaming before completion, + WAV metadata, disconnect cancellation, and backend failure handling. +- The v6 container build passed all 13 registered CTest checks. +- The actual tokenizer produced identical IDs for 38 corpus and edge cases + when changing from `encode` to singleton `encode_batch`. All 35 cases within + the token budget also produced byte-identical MessagePack submit payloads. + +The tests and lifecycle checks do not provide new MOS, intelligibility, speaker +similarity, ASR error-rate, or other acoustic scores. No such scores are claimed. +Cross-frontend exact audio parity also remains separate from the timing samples +below: lane ordering and effective fixed-cohort seeds differed between runs. + +## v3 performance samples + +Each v3 mode below had **one server boot and three measured trials**, with text +normalization enabled and the retained ten-text corpus. Fixed-cohort trials +submitted one synchronized wave of 512 requests; in-flight trials submitted +2,048 requests at concurrency 512 with refills. Each client used two warmup +rounds. The two load patterns and client timing definitions differ and should +not be compared as equivalent latency experiments. + +Values below are the equal mean of the three trial statistics, not percentiles +of a pooled request set. Peak VRAM is the maximum observed device memory usage. + +| B512 sample | Successful requests | TTFA p50 | TTFA p95 | TTFA p99 | RTFx | Peak VRAM | +|---|---:|---:|---:|---:|---:|---:| +| Legacy same-host fixed-cohort control | 1,536 / 1,536 | 180.00 ms | 191.70 ms | 196.63 ms | 532.30 | 28,615 MiB | +| C++/bridge v3 fixed-cohort | 1,536 / 1,536 | 276.33 ms | 298.18 ms | 306.37 ms | 481.18 | 28,587 MiB | +| C++/bridge v3 in-flight | 6,144 / 6,144 | 146.70 ms | 349.25 ms | 374.34 ms | 1,267.48 | 24,673 MiB | + +The fixed-cohort p95 difference is **106.48 ms** in this sample. The legacy +control used the same host, native package, normalizer library, corpus, and CPU +role allocations, with a frozen legacy frontend binary. All ten normalized +strings, their token IDs, and per-cohort multiplicities matched. The legacy +frontend performed an additional B1 readiness synthesis, and effective seeds +and lane ordering differed; generated durations therefore differed too. These +controls help locate host overhead but do not establish an audio-matched +throughput comparison. No same-host legacy in-flight control was collected. + +The v3 in-flight client recorded five playback underruns in its first trial and +none in the other two. The largest reported audio-buffer deficit was 18.64 ms. +Zero request failures should not be interpreted as zero playback interruption. + +Earlier diagnostic receipts localized most of the startup difference to request +admission: the legacy worker received a B512 cohort over a median 116.35 ms, +versus 186.79 ms for one bridge diagnostic trial. Backend startup after the final +submission was about 2–3 ms in both arms; backend first-chunk time was about +64–67 ms for the legacy control and 70.46 ms for that diagnostic. These numbers +describe the diagnostic run, not a decomposition of the three-trial v3 table. + +An isolated CPU preparation check on four physical cores with SMT and 16 +preparation threads measured 140.03 ms for 512 requests with `encode`, and +134.18 ms with singleton `encode_batch`. It used the real tokenizer and grammar, +with `TOKENIZERS_PARALLELISM=false`. Reducing the preparation pool or retaining +the GIL during normalizer cleanup did not improve total preparation time. +Moving the bridge onto frontend cores also worsened the HTTP diagnostic. +No further host optimization is promoted by these checks. + +## v4 optional float32 forwarding screen + +The v4 screen enabled `NEMOTRON_BRIDGE_FORWARD_F32=1`. This optional path forwards +native float32 samples to C++ for PCM16 conversion; the public HTTP format +remains PCM16. The default bridge still performs PCM16 conversion in Python. +Each mode again used one server boot and three measured trials. + +| B512 v4 sample | Successful requests | TTFA p50 | TTFA p95 | TTFA p99 | RTFx | Peak VRAM | +|---|---:|---:|---:|---:|---:|---:| +| Fixed-cohort | 1,536 / 1,536 | 267.34 ms | 283.33 ms | 289.81 ms | 636.16 | 28,583 MiB | +| In-flight | 6,144 / 6,144 | 149.00 ms | 323.73 ms | 359.19 ms | 1,245.80 | 24,673 MiB | + +The in-flight client reported no playback deficit responses in these three +trials. Fixed-cohort RTFx ranged from 470.34 to 795.44 across trials; the mean +alone does not establish a throughput gain. Fixed-cohort p95 TTFA remained +91.62 ms above the same-host legacy control. This screen does not promote the +optional transport to the default or establish acoustic equivalence. + +## Final representative samples + +The v6 image includes CLI seed propagation and bounded preparation cancellation. +All six arms completed with text normalization enabled and optional +`NEMOTRON_BRIDGE_FORWARD_F32=1`. The bridge default remains PCM16 forwarding. +The final tests used the image's freshly built CUDA libraries; all 15 packaged +shared-library hashes match the tested in-flight package. The fixed-cohort +package was assembled from the same build using the fixed serving policy. + +Each arm used one fresh boot and three measured trials. Fixed-cohort B1/B32 +used seven synchronized waves per trial, and B512 used one. In-flight trials +contained 32/128/2048 requests at concurrency 1/32/512. Warmups were excluded. +Numbers are equal means of trial statistics; latency and ITL are milliseconds. + +| B | Mode | RTFx | req/s | TTFA p50 | p95 | p99 | ITL avg | ITL p95 | Peak MiB | Requests | Failures | Deficit responses | +|---:|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| 1 | Fixed cohort | 16.73 | 2.14 | 20.25 | 20.95 | 20.95 | 4.40 | 4.59 | 5,191 | 21 | 0 | — | +| 1 | In-flight | 17.50 | 2.51 | 19.88 | 20.86 | 20.99 | 4.39 | 4.54 | 5,199 | 96 | 0 | 0 | +| 32 | Fixed cohort | 259.35 | 34.61 | 32.18 | 48.82 | 49.54 | 3.92 | 4.73 | 10,335 | 672 | 0 | — | +| 32 | In-flight | 301.40 | 42.30 | 25.31 | 43.32 | 44.51 | 7.47 | 12.59 | 6,835 | 384 | 0 | 0 | +| 512 | Fixed cohort | 511.06 | 67.62 | 265.75 | 279.29 | 285.44 | 19.25 | 23.31 | 28,587 | 1,536 | 0 | — | +| 512 | In-flight | 1238.79 | 170.87 | 149.06 | 336.34 | 361.50 | 28.67 | 37.00 | 24,673 | 6,144 | 0 | 0 | + +All **8,853 measured requests completed**. The in-flight client recorded zero +playback-deficit responses across its 6,624 measured requests. The historical +client does not report that metric; its dashes are not evidence of zero deficits. +The repository portable client additionally completed 64 B32 requests in each +mode, with no failures or measured playback deficits. Those smoke tests have +different timing definitions and are excluded from this table. + +B512 fixed-cohort p95 remains **87.59 ms above** the legacy control from the +earlier same-host diagnostic comparison. This integration therefore does not +yet reproduce the legacy frontend's B512 latency. The diagnostic admission +gap is the main identified remaining host-path cost. The control and final +samples were not collected as matched boot pairs; no confidence interval or +performance-parity claim is made. No fresh acoustic campaign was performed. + +The host used one H100 SXM 80 GB (132 SMs), driver 595.91.07, and an AMD EPYC +7413 CPU. The native worker and Python bridge shared physical cores 0–3 plus +SMT siblings 24–27; C++ HTTP used 4–7/28–31; clients used 8–11/32–35. +No GPU clock or power settings were changed. GPU telemetry and source/binary +identities accompany the raw observations. + +Image index: `sha256:af3ce40c4f1ee973630da290fcf4f6d1ec96fb4b9f1c2d6ab0de7b03a5fb1474`. +Image config: `sha256:93e9f1ba8ce1e7ef094f13554bcf8147203f825dcfb391096630605ea1b19daf`. +The image contains no TensorRT or TensorRT-LLM installation used by this backend. +Its build passed 13 CTests, including the 18-case bridge/launcher suite. +Scalar/AVX2 PCM conversion also passed independent boundary/random-value and +ASan/UBSan checks. Remote CI, gRPC builds, other GPU geometries, and wider +batch configurations have not been validated by this campaign. + +## Evidence inventory + +The following names identify receipts relative to the external validation +directory. They are evidence identifiers, not bundled downloads or links to +machine-specific paths. + +| Evidence | Receipt | +|---|---| +| Imported-file byte comparison | `frozen-import-inventory.json` | +| Qualified runtime imports | `runtime-v3.json`, `runtime-v6.json` | +| Native readiness and lifecycle | `b1-qualification/worker/NATIVE_READY.json`, `b32-qualification/worker/LIFECYCLE.json`, `b512-qualification/worker/LIFECYCLE.json` | +| B1 HTTP repeatability | `b1-execution.json` | +| Final registration and summary | `HTTP_REGISTRATION-final.json`, `final-summary.json`, `final-summary.md` | +| Final measurements and portable smoke tests | `http-final/b{1,32,512}-{fixed_cohort,inflight}/` | +| Image/tested library identity | `image-tested-package-v6.json`, `kernels-v6/BUILD.json` | +| v3 experiment registration | `HTTP_REGISTRATION-r3.json` | +| v3 fixed-cohort measurements | `http-r3/b512-fixed_cohort/historical.json` | +| v3 in-flight measurements | `http-r3/b512-inflight/inflight-{0,1,2}.json` | +| v4 optional-transport screen | `http-r4/b512-fixed_cohort/historical.json`, `http-r4/b512-inflight/inflight-{0,1,2}.json` | +| Legacy control and provenance | `http-legacy-control/COMPARISON.json`, `http-legacy-control/historical.json` | +| Admission diagnostics | `http-diag2/b512-fixed_cohort/bridge-trace.jsonl` and worker cohort receipts | +| CPU preparation and normalization controls | `tokenizer-qualification/encode-comparison.json`, `thread-sweep.json`, `normalizer-profile.json` | + +Every timing arm retains its launch command, image identity, worker receipts, +client records, GPU telemetry, and shutdown result. Repeating a performance +claim requires additional independent boots and matched measurement conditions; +the current one-boot samples have no confidence intervals. + +## Subsequent B512 host-path repair + +The following campaign supersedes the v6 performance samples for the repaired +host path; the earlier observations above remain unchanged. With the same native +CUDA package, moving normalization and tokenization into the C++ frontend's +native preparation path reduced B512 HTTP TTFA p95 from **274.60 to 224.30 ms** +for fixed cohorts and from **335.23 to 229.95 ms** for in-flight requests. This +repairs most of the port's measured host overhead, but does not establish parity +with the legacy frontend or a change in neural execution speed. + +### Repair and executable scope + +The repaired path is C++ HTTP → native normalization/tokenization → validated +prepared-request transport → Python bridge → unchanged native controller and +CUDA. Preparation uses the retained English normalizer and a native tokenizer +binding, with the same trimming, punctuation, tokenizer assets, token budget, +and EOS rules. The bridge negotiates prepared-input support and validates token +IDs before submission. Its event loop no longer performs each request's text +preparation. This backend-specific path does not replace preparation for other +models. + +The tested bundle also enables bounded bridge event-loop yielding, TCP_NODELAY +on this frontend's HTTP connections, and float32 forwarding with exact PCM16 +conversion in C++. Both comparison arms already used float32 forwarding. +Consequently, the table attributes the improvement to the complete tested host +bundle; it does not assign independent HTTP gains to these mechanisms. Worker +affinity changes and process-pool preparation were investigated but are not +enabled by the selected configuration. + +There are **no CUDA kernel, native controller, state-format, generation, or +chunk-schedule changes** in this comparison. Both arms use the exact retained v6 +native package, built for `sm_90a` at widths 1, 32, and 512. Its kernel source +manifest SHA-256 is +`3976dd371b717d3a5c3b7bb3c2ec1e8dbdada2ffec6f4b9f5cee7f542f3ed222`. +The previously verified 15 packaged shared libraries are reused, rather than +substituted with independently rebuilt binaries. The fixed-cohort and in-flight +serving policies remain distinct, with the same policy package used by both +arms of each comparison. + +| Tested image | Image index SHA-256 | +|---|---| +| v6 baseline | `af3ce40c4f1ee973630da290fcf4f6d1ec96fb4b9f1c2d6ab0de7b03a5fb1474` | +| v7 repaired host | `dea07919df447b9ca57ab15a390d5b592c95a3287d35fe8301554b4f8da6433d` | + +The v7 image rebuilt the C++ frontend and native preparation component while +copying `/package` from v6. The C++ build passed all 13 CTests. The Python suite +ran 44 tests with three fixture-dependent skips, and a separate real-fixture +qualification ran nine tests successfully. The latter checks normalization and +token-ID agreement, native allocation ownership, and concurrent preparation. +The normalizer/tokenizer comparison includes the corpus and edge cases with +normalization both enabled and disabled. No fresh acoustic scoring was performed. + +The complete source-build recipe remains `docker/Dockerfile.nemotron`, including +CUDA compilation. This host-only campaign deliberately canceled its redundant +CUDA rebuild and reused the qualified v6 package; it does not certify a newly +completed clean build of every component at v7. Remote CI, gRPC, other GPU +architectures, and widths beyond the recorded cases remain outside this result. + +### Matched HTTP comparison + +Each B512 arm and mode used **two fresh server boots and three measured trials +per boot**, ordered baseline/repaired/repaired/baseline. Fixed-cohort trials used +one synchronized 512-request wave; in-flight trials used 2,048 requests at +concurrency 512 with refills. Two warmup rounds were excluded. The ten-text +corpus, normalization, voice, generation cap, 2/3/4/5 → 8 chunk schedule, and +publication behavior were unchanged. The clients' existing timers were +preserved, including admission delay; no interval was subtracted. + +CPU placement remained four physical cores plus SMT for the worker/bridge +(0–3/24–27), another four for C++ HTTP and preparation (4–7/28–31), and four for +the client (8–11/32–35). This is a **4+4 server-core comparison**. It does not +include the separate shared-eight-core frontend-placement control. + +Statistics are equal means of per-trial values across both boots, not pooled +request percentiles. Latency and ITL are in milliseconds; peak memory is the +maximum telemetry after readiness, including warmups. + +| B | Host | Mode | RTFx | req/s | TTFA p50 | p95 | p99 | ITL avg | ITL p95 | Peak MiB | Requests | Failures | Deficit responses | +|---:|---|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| 512 | v6 baseline | Fixed cohort | 471.21 | 62.43 | 259.61 | 274.60 | 281.01 | 19.64 | 23.36 | 28,585 | 3,072 | 0 | — | +| 512 | v7 repaired | Fixed cohort | 543.67 | 72.74 | 203.77 | 224.30 | 231.48 | 19.51 | 23.39 | 28,621 | 3,072 | 0 | — | +| 512 | v6 baseline | In-flight | 1252.60 | 172.77 | 146.15 | 335.23 | 364.24 | 28.38 | 36.95 | 24,673 | 12,288 | 0 | 0 | +| 512 | v7 repaired | In-flight | 1263.37 | 174.26 | 147.26 | 229.95 | 249.26 | 28.23 | 36.61 | 24,673 | 12,288 | 0 | 0 | + +All **30,720 measured B512 requests completed**. Both in-flight arms recorded +zero playback-deficit responses. The historical fixed-cohort client does not +measure playback deficits; the dashes do not imply zero. + +| Mode | TTFA p95 change | Repaired / baseline | 95% matched interval for ratio | +|---|---:|---:|---:| +| Fixed cohort | −50.30 ms (−18.32%) | 0.8168 | [0.7576, 0.8621] | +| In-flight | −105.28 ms (−31.41%) | 0.6859 | [0.6453, 0.7216] | + +Intervals use 20,000 paired hierarchical bootstrap draws, resampling matched +boot pairs and then trial positions within each pair. Two boot pairs on one GPU +provide limited uncertainty evidence, not independent hardware replication. +In-flight p50 changed little: this repair primarily reduced startup/tail +preparation delay in that workload. + +Fixed-cohort RTFx must not be interpreted as a kernel speedup. Asynchronous +preparation changes lane order and therefore batch-seeded generated trajectories. +The baseline's six measured cohorts all reached 511 producer steps; repaired +cohorts ran for 483/511/511 and 435/511/275 steps. Length-limit terminations were +8/3 for baseline/repaired, respectively, with the remaining requests ending at +EOS. All 12,288 requests in each B512 in-flight arm ended at EOS. These +differences are retained in the result rather than removed from throughput. + +### Legacy residual and B32 check + +A fresh same-host legacy fixed-cohort control used one boot and three trials, +with the same four-plus-four server-core assignment and retained native package. +Its TTFA p95 was **204.19 ms**, versus the repaired port's **224.30 ms**: a +remaining **20.10 ms, or 9.84%**, difference. The legacy control is not a matched +two-boot arm and has different lane/seed trajectories, so this comparison has +no matched confidence interval and does not demonstrate performance parity. +The earlier 191.70-ms legacy result remains a historical observation above. + +The repaired B32 regression check used one boot and three trials per mode. +Fixed cohorts retained seven waves per trial; in-flight trials contained 128 +requests. It is a smoke/regression sample, not a new matched improvement claim. + +| B | Mode | RTFx | req/s | TTFA p50 | p95 | p99 | ITL avg | ITL p95 | Peak MiB | Requests | Failures | Deficit responses | +|---:|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| 32 | Fixed cohort | 233.72 | 31.28 | 32.68 | 45.16 | 46.35 | 3.90 | 4.73 | 10,335 | 672 | 0 | — | +| 32 | In-flight | 298.11 | 41.84 | 26.08 | 39.98 | 41.34 | 7.51 | 13.05 | 6,835 | 384 | 0 | 0 | + +There are no new CER, WER, speaker-similarity, or other acoustic figures for +these arms. Exact text/token preparation tests and unchanged neural binaries +do not replace acoustic evaluation. Additional CPU-placement controls should +be reported separately from this frozen 4+4 comparison. + +### Repair evidence inventory + +The following receipts are relative to the external host-repair evidence +directory. They supplement the earlier validation inventory. + +| Evidence | Receipt | +|---|---| +| Frozen campaign and independent summary | `FINAL_REGISTRATION.json`, `FINAL_COMPARISON.json`, `FINAL_COMPARISON.md`, `summarize_final.py` | +| Per-boot clients, telemetry, commands, termination records | `http-final/` | +| Fresh legacy control | `legacy-fixed-refresh/REGISTRATION.json`, `legacy-fixed-refresh/RECEIPT.json`, `legacy-fixed-refresh/historical.json` | +| Host-only build scope and image | `HOST_BUILD_SCOPE.json`, `Dockerfile.host-only`, `build-v7-host-only.log` | +| C++ and Python test receipts | `build-v7.log`, `tests-frontend-preparation-all.log`, `native-v7-qualification.log` | +| Native preparation source/fixture qualification | `preparation-process/native-qualification-formatted.log` | +| Retained negative process-pool experiment | `preparation-process/result.json` | + +## Subsequent shared-eight-core placement + +A separate placement experiment allowed the C++ frontend and its native +preparation workers to use **all eight existing server cores**, 0–7/24–31, +while keeping the worker and bridge on 0–3/24–27. The total service allocation +remains eight physical cores plus SMT. This relaxes the earlier rigid split; +it does not add CPU cores or change GPU execution. The v7 image, native package, +client commands, normalization, and generation contract were unchanged. + +After an exploratory placement screen, two fresh boots per mode each ran three +measured trials with the same warmups and request counts as the comparison +above. These measurements were collected **after** that comparison, without +interleaving another baseline. Differences against the earlier arms are +therefore descriptive chronological comparisons, not matched randomized +estimates of the placement's isolated effect. + +| B | Mode | RTFx | req/s | TTFA p50 | p95 | p99 | ITL avg | ITL p95 | Peak MiB | Requests | Failures | Deficit responses | +|---:|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| 512 | Fixed cohort | 621.76 | 83.45 | 160.36 | 176.80 | 184.06 | 19.70 | 24.52 | 28,613 | 3,072 | 0 | — | +| 512 | In-flight | 1310.30 | 180.73 | 140.95 | 224.45 | 239.73 | 27.03 | 36.06 | 24,673 | 12,288 | 0 | 0 | + +Fixed-cohort boot means were **173.77 and 179.84 ms** TTFA p95; in-flight boot +means were **217.30 and 231.61 ms**. Within-arm hierarchical bootstrap intervals +from 20,000 draws were [170.21, 182.93] ms and [216.16, 235.40] ms, respectively. +These describe the limited two-boot sample on one GPU and do not account for +placement selection or chronological effects. + +The fixed-cohort mean is 47.49 ms below the earlier repaired 4+4 sample and +27.39 ms below the fresh 204.19-ms legacy control. The in-flight mean is 5.49 ms +below the repaired 4+4 sample; its boot-to-boot variation prevents attributing +that small difference confidently to placement. Relative to the earlier v6 +baseline, the observed complete-stack p95 differences are −97.80 ms fixed and +−110.78 ms in-flight. No paired confidence interval is claimed for these +chronological contrasts. + +All **15,360 measured requests completed**. Fixed cohorts recorded 3,071 EOS +terminations and one length limit; cohort steps were 436/387/332 and +511/363/307. All 12,288 in-flight requests ended at EOS with zero measured +playback-deficit responses. As above, fixed-cohort RTFx varies with generated +trajectories and cap tails, and its client does not record playback deficits. + +No fresh acoustic scores or measured cross-frontend PCM identity are claimed. +The in-flight population matches normalized texts, sample counts, and +termination reasons between the original baseline and repaired arms. Measured +request receipts do not contain PCM hashes; hashes in HTTP preparation receipts +belong to pre-serving qualification and cannot establish exact audio identity +for the timed requests. + +This separate sample is recorded by +`HTTP_REGISTRATION-shared-cpu-final.json`, `shared-cpu-final/`, +`SHARED_CPU_COMPARISON.json`, `SHARED_CPU_COMPARISON.md`, and +`summarize_shared_cpu.py` in the external host-repair evidence directory. The +frozen 4+4 comparison above remains the matched evidence for the code repair. + +## Fresh legacy in-flight control + +A subsequent legacy Rust frontend control completed one fresh boot with three +unchanged `tts_perf_client` trials: 2,048 measured requests at concurrency 512 +per trial, after two warmup rounds. Its TTFA p95 was **212.93 ms**, versus +**224.45 ms** for the selected C++ host path above. The remaining point-estimate +difference is **11.53 ms, or 5.41%**. The selected fixed-cohort point estimate +remains 176.80 ms versus the fresh legacy control's 204.19 ms. + +| Frontend | Mode | Boots | RTFx | req/s | TTFA p50 | p95 | p99 | ITL avg | ITL p95 | Peak MiB | Requests | Failures | Deficit responses | +|---|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| Legacy Rust | Fixed cohort | 1 | 490.04 | 64.74 | 191.48 | 204.19 | 209.66 | 18.88 | 22.13 | 28,615 | 1,536 | 0 | — | +| Selected C++ | Fixed cohort | 2 | 621.76 | 83.45 | 160.36 | 176.80 | 184.06 | 19.70 | 24.52 | 28,613 | 3,072 | 0 | — | +| Legacy Rust | In-flight | 1 | 1308.91 | 180.54 | 141.35 | 212.93 | 224.95 | 27.14 | 35.26 | 24,673 | 6,144 | 0 | 0 | +| Selected C++ | In-flight | 2 | 1310.30 | 180.73 | 140.95 | 224.45 | 239.73 | 27.03 | 36.06 | 24,673 | 12,288 | 0 | 0 | + +Both implementations used the same H100, native package, model, normalization +assets, corpus, and client timing definitions. Both had eight physical service +cores plus SMT. Legacy kept frontend and worker on disjoint four-core sets; +selected C++ allowed frontend preparation to share all eight cores while its +worker/bridge retained four. Legacy also performed one excluded readiness +synthesis before the registered client warmups. These are chronological samples +with **one legacy boot versus two C++ boots**, not interleaved matched pairs. +There is no resolved parity or superiority claim, and no matched confidence +interval is assigned to the legacy comparison. + +All 6,144 measured legacy in-flight requests ended at EOS with no failures or +playback-deficit responses. Each selected C++ in-flight boot matched that +control's complete normalized-text/sample-count/termination population and +**44,544.72 seconds of generated audio** across 6,144 measured requests. This +comparison explicitly excludes legacy readiness and each client's warmups. +It verifies population and duration agreement, not PCM identity or acoustic +quality. Fixed-cohort lane-order/seed trajectories still differ, so their RTFx +difference cannot be attributed solely to frontend speed. + +The additional receipts are `legacy-inflight-refresh/REGISTRATION.json`, +`legacy-inflight-refresh/RECEIPT.json`, the three client results under +`legacy-inflight-refresh/`, and the independently reconstructed +`LEGACY_COMPARISON.json`, `LEGACY_COMPARISON.md`, and `summarize_legacy.py` in +the external host-repair evidence directory. + +### Launcher and small-width checks + +The optional `--native-cpus` and `--frontend-cpus` controls were subsequently +added to the Nemotron launcher. Image `nemo-speech-cpp:nemotron-v7.1` changes +only that launcher relative to v7; CUDA, the C++ frontend, and native text +preparation binaries are identical. The launch options reproduce the selected +CPU masks before child threads initialize and record their effective affinity. +No placement is imposed when the options are absent. + +Live checks using these options measured B1 fixed-cohort p95 **21.11 ms** +(21 measured requests) and B32 in-flight p95 **33.40 ms** (384 measured requests), +with three trials per fresh boot and normalization enabled. All requests +completed, and B32 recorded zero playback deficits. These are regression +samples, not matched cross-width performance claims. All existing child-thread +affinities matched the requested role masks. Commands, results, and launch +receipts are retained in `launcher-regression/`. + +Validation includes 13 CMake tests, six Rust tests, and 51 Python controller, +bridge, and launcher tests (three optional model-fixture tests skipped in that +generic invocation). The separate nine-test model/native-library invocation +passed with those fixtures available. Formatting and repository hooks passed. +An initial test run inherited the image's F32 forwarding default while one +test expected PCM16; the test now explicitly selects its PCM16 control, while +the separate F32 test retains coverage of the selected path. Runtime behavior +was unchanged by that test correction. Logs are retained in `build-v7.log`, +`tests-final.log`, `tests-final-r2.log`, `native-v7-qualification.log`, and +`precommit-final.log`. + +## Subsequent direct-worker transport + +The opt-in `--direct-worker` path removes the Python compatibility bridge from +request submission and audio delivery. C++ prepares and validates text, sends +the native controller's existing MessagePack protocol, and copies shared-memory +audio into owned PCM16 before acknowledging slab reuse. The native controller +still owns admission, request slots, GPU work and retirement. No CUDA kernel, +arithmetic, RNG, state format, generation policy or chunk schedule changed. +The default remains the compatibility bridge; other models are unaffected. + +This section reports the selected direct **base** implementation. Experimental +descriptor batching, publication-before-ACK and other development candidates +are excluded. Both arms retain normalization and the shared-eight-core placement: +frontend/preparation on 0–7/24–31 and native worker on 0–3/24–27. The baseline's +bridge shares the worker's mask. The total service allocation is eight physical +cores plus SMT, with the client on separate cores 8–11/32–35. + +### Matched B512 in-flight result + +The comparison used two fresh boots per arm in baseline/direct/direct/baseline +order, three measured trials per boot, 2,048 requests per trial at concurrency +512, and two excluded warmup rounds. Both arms used the exact retained external +`tts_perf_client` binary, SHA-256 +`d7713f0bbd418c95395959ebd7497bc3c1110c5313da41f0b0dacf019730abbf`. +The client stopwatch, arrivals, preparation and admission accounting were +unchanged. Values are equal means of six per-trial statistics, not pooled +request percentiles; latency and ITL are milliseconds. + +| B | Transport | RTFx | req/s | TTFA p50 | p95 | p99 | ITL avg | ITL p95 | Peak MiB | Requests | Failures | Deficit responses | +|---:|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| 512 | Compatibility bridge | 1332.86 | 183.84 | 138.40 | 218.53 | 235.37 | 26.60 | 35.79 | 24,673 | 12,288 | 0 | 0 | +| 512 | Direct worker | 1353.61 | 186.70 | 135.59 | 205.04 | 217.24 | 26.06 | 34.94 | 24,673 | 12,288 | 0 | 0 | + +Direct transport reduces the primary statistic by **13.49 ms (6.17%)**. The +95% matched bootstrap interval is **[−20.36, −4.22] ms**, with a direct/baseline +ratio interval of **[0.9068, 0.9801]**. The calculation uses 20,000 draws of +matched boot blocks and paired trial positions. Only two blocks on one GPU +were measured, so this is limited single-device evidence. + +Baseline boot means were 211.50 and 225.57 ms; direct boot means were 199.59 +and 210.50 ms. The direct mean's within-arm interval is [191.20, 214.58] ms. +Consequently, **TTFA p95 ≤200 ms is not established**. Development screens and +diagnostics do not enter this comparison. The older 212.93-ms legacy Rust +sample was collected separately and is not a paired arm of this experiment. + +### Separate supporting checks + +The direct path also completed the following single-boot, three-trial samples. +They are regression/smoke measurements, not matched improvement estimates. + +| B | Mode | RTFx | req/s | TTFA p50 | p95 | p99 | Requests | Failures | Deficit responses | +|---:|---|---:|---:|---:|---:|---:|---:|---:|---:| +| 1 | In-flight | 17.54 | 2.51 | 18.32 | 19.25 | 19.33 | 96 | 0 | 0 | +| 512 | Fixed cohort | 485.26 | 63.87 | 154.16 | 171.09 | 177.34 | 1,536 | 0 | — | + +Fixed-cohort lane ordering and batch-seeded output lengths differ across +separate runs; the RTFx above is not evidence of a kernel or transport throughput +gain. Its historical client does not measure playback deficits. + +B32 received a separate same-seed HTTP qualification, not a fresh timing claim. +Each arm served 64 requests, including four client cancellations and four slow +readers. All 60 uncancelled requests completed; the four cancellations were +intentional. All 64 corresponding received PCM hashes and sample counts matched +between bridge and direct paths, including the cancelled prefixes. Native +receipts checked physical B32, request seeds, slot generations, cancellations +and completed sample totals. These checks establish this fixture's exact +output/lifecycle agreement, not fresh CER, WER or speaker-similarity scores. + +A clean CPU-only CMake build passed all **14 CTests**, including 121 native +transport checks and the 59-test bridge/launcher suite (three optional fixture +tests skipped). It does not rebuild or requalify CUDA kernels. Live execution +reused the same qualified native package; remote CI and additional widths or +GPU architectures remain outside this result. + +### Direct-worker evidence and reproduction + +The following receipts are relative to the external in-flight-r2 evidence +directory. The tested direct frontend library SHA-256 is +`5577751830e46b0defdcd903eeae13a86207d308d41b59558dcd1f3f5c7dd8de`; +the B512 in-flight serving manifest SHA-256 is +`f7d92e89925ddca8ee8f419700456e2700e6d82b629a67364c6776f76a8caee5`. + +| Evidence | Receipt | +|---|---| +| Registration, identities and matched statistics | `FROZEN_CLIENT_COMPARISON.json`, `DIRECT_COMPARISON.json` | +| Matched client observations, telemetry and launch commands | `screens/{baseline_frozen_client,direct_frozen_client}-boot{1,2}/` | +| B1 and fixed-cohort supporting samples | `screens/direct_b1-boot1/`, `screens/direct_fixed512-boot1/` | +| B32 exact PCM and lifecycle comparison | `screens/{qual_baseline,qual_direct}-boot1/pcm-qualification/RESULT.json` | +| Clean CPU build and full CTest results | `ci-build/RECEIPT.json`, `ci-build/ctest.log` | + +Follow [direct-worker B512 reproduction](nemotron-megakernel.md#direct-worker-b512-comparison) +with the sealed package, model, normalization assets and frozen client. Remove +only `--direct-worker` for the bridge control, keep all other settings matched, +and use fresh server boots and new receipt directories for each arm. diff --git a/docs/nemotron-megakernel.md b/docs/nemotron-megakernel.md new file mode 100644 index 0000000..13931d3 --- /dev/null +++ b/docs/nemotron-megakernel.md @@ -0,0 +1,444 @@ +# Nemotron megakernel TTS backend + +The `nemotron` backend connects the NeMo-Speech.cpp C++ frontend to the retained +Nemotron native controller and CUDA megakernels. It supports the pinned +EasyMagpie checkpoint, `english_sample` voice, English text, and mono 22050 Hz +PCM16 output. It is a different model from the default MagpieTTS/NanoCodec GGUF +backend. + +The Python controller is retained (Python 3.11 or newer; the pinned image uses +Python 3.12). Python owns native package loading, GPU state, in-flight +scheduling, and the compatibility bridge. The C++ frontend owns public HTTP +handling and, in the Nemotron image, calls the native normalization/tokenization +library before multiplexing prepared requests over one Unix socket. The +bridge-side preparation path remains available for diagnostics. This +integration is not a pure C++/CUDA rewrite of the controller. The neural +megakernels do not execute through GGML, vLLM, TensorRT, or TensorRT-LLM. + +The vendored controller, kernels, and native normalizer retain their source +notices. Their licensing and redistribution review is pending; a successful +local build does not authorize publishing those sources or model assets. The +imported closure is not relicensed by the repository's Apache-2.0 license; see +the [retained license scope](../THIRD_PARTY_NOTICES.md#imported-nemotron-native-deployment-closure). +The container installs the project and imported component notices under +`/opt/nemo-speech/share/licenses/nemo-speech/`. The retained normalizer README +describes its historical Rust integration and build paths. Use this document +for the current frontend, Docker recipe, and Python bridge setup. + +## Build + +The container recipe builds the frontend, native normalizer, selected CUDA +kernels, and controller package from source, using pinned builder/runtime +images. Model and grammar assets are mounted separately at runtime: + +```bash +docker build -f docker/Dockerfile.nemotron \ + --build-arg NATIVE_WIDTHS=1,32,512 \ + --build-arg NATIVE_BATCHING_MODE=inflight \ + --build-arg NATIVE_INFLIGHT_PROFILE=retained \ + -t nemo-speech:nemotron . +``` + +The batching mode is fixed when the native package is assembled. To build a +separate fixed-cohort image, use `--build-arg NATIVE_BATCHING_MODE=fixed_cohort` +and a separate tag such as `nemo-speech:nemotron-fixed`. Start that image with +`--mode fixed_cohort`; start the default image with `--mode inflight`. +The launcher rejects a mode that differs from the package manifest. Fixed-cohort +execution needs a full wave of the selected width, so a lone request at B32 or +B512 waits for the rest of its cohort. + +Initialize the repository's submodules before configuring CMake. A CPU frontend +build can serve the GPU controller; the controller and native kernels have +their own CUDA runtime requirements. + +```bash +cmake -S . -B build-native -DCMAKE_BUILD_TYPE=Release \ + -DNEMO_SPEECH_BUILD_TTS=ON -DNEMO_SPEECH_BUILD_HTTP=ON \ + -DNEMO_SPEECH_BUILD_ASR=OFF -DNEMO_SPEECH_BUILD_DIAR=OFF \ + -DNEMO_SPEECH_BUILD_NMT=OFF -DNEMO_SPEECH_BUILD_S2S=OFF \ + -DNEMO_SPEECH_BUILD_GRPC=OFF -DNEMO_SPEECH_BUILD_MIC_CAPTURE=OFF \ + -DGGML_CUDA=OFF +cmake --build build-native --target nemo_speech_cli -j 4 +``` + +Use the retained CUDA/PyTorch runtime for native compilation and execution. +The kernel recipes target Hopper `sm_90a` and the qualified H100 configuration +has 132 SMs. An H100 80 GB PCIe card can have fewer SMs; its name and memory +capacity alone do not establish compatibility with these launch recipes. +Changing the GPU geometry, architecture, or runtime requires separate +qualification. The bridge additionally uses the `tokenizers` +package and NumPy. Python support dependencies are pinned in +[`scripts/nemotron_native/requirements.txt`](../scripts/nemotron_native/requirements.txt); +PyTorch comes from the qualified CUDA image and is not replaced by that file. +For a local virtual environment inside that CUDA/PyTorch environment, retain +the qualified system PyTorch installation: + +```bash +python3 -m venv --system-site-packages .venv-nemotron +. .venv-nemotron/bin/activate +python3 -m pip install -r scripts/nemotron_native/requirements.txt +``` + +Validate the runtime's actual imports and shared libraries: + +```bash +PYTHONPATH="$PWD/python${PYTHONPATH:+:$PYTHONPATH}" \ + python3 scripts/nemotron_native/runtime_dependencies.py \ + --include-backend --output /tmp/nemotron-runtime-dependencies.json +python3 scripts/nemotron_native/build_kernels.py --validate +python3 scripts/nemotron_native/build_kernels.py \ + --widths 1,32,512 --output /tmp/nemotron-kernels --jobs 2 +python3 scripts/nemotron_native/build_package.py \ + --kernels /tmp/nemotron-kernels --output /tmp/nemotron-package \ + --widths 1,32,512 --batching-mode inflight --inflight-profile retained +``` + +For a local fixed-cohort package, reuse the compiled kernels but choose a fresh +output directory and `--batching-mode fixed_cohort`. Select that directory with +the launcher's `--package` and pass `--mode fixed_cohort`. + +Build outputs and package directories must be fresh. Compilation produces an +unqualified package; it does not establish audio correctness or performance. +Keep the model directory and its seal outside the repository. The expected +checkpoint and file hashes are recorded in +[`deploy/nemotron_native/model.lock.json`](../deploy/nemotron_native/model.lock.json). +Do not substitute the default `nemo-speech pull magpie` assets. + +## Run + +The Nemotron Docker image enables native frontend text preparation. It uses +the retained normalization ABI and tokenizer 0.21.4 directly from the C++ +request path, then forwards validated text/token IDs through the bridge. +Preparation remains inside each HTTP request's measured lifetime. There is no +text-result cache or normalization bypass. Other model backends do not load +this library or use its protocol. + +For an equivalent source-built deployment, build and select the preparation +library before launching: + +```bash +cargo build --manifest-path native/text_preparation/Cargo.toml --release --locked +export NEMOTRON_NATIVE_PREPARATION_LIBRARY="$PWD/native/text_preparation/target/release/libnemotron_text_preparation.so" +export NEMOTRON_FRONTEND_PREPARATION=1 +export NEMOTRON_BRIDGE_FORWARD_F32=1 +export NEMOTRON_BRIDGE_FAIR_BATCH=64 +export NEMOTRON_HTTP_TCP_NODELAY=1 +``` + +The Docker builder pins Rust 1.88.0 and the dependency lockfile. Setting +`NEMOTRON_FRONTEND_PREPARATION=0` retains bridge-side preparation for diagnostic +comparisons. `NEMOTRON_BRIDGE_FAIR_BATCH` yields between bounded groups of ready +messages without introducing a batching timer. `NEMOTRON_HTTP_TCP_NODELAY` +affects only the Nemotron HTTP backend. Raw-F32 forwarding preserves the +existing PCM16 conversion and moves that conversion into the C++ reader. + +Start the controller and bridge in the qualified Python/CUDA environment: + +```bash +python3 scripts/nemotron_native/launch.py \ + --package /tmp/nemotron-package \ + --model /models/nemotron --model-seal /models/model-seal.json \ + --width 32 --mode inflight \ + --listen /tmp/nemotron-b32.sock \ + --receipts /tmp/nemotron-b32-receipts \ + --http-host 127.0.0.1 --http-port 8080 --http-threads 40 \ + --frontend-binary build-native/bin/nemo-speech +``` + +The launcher verifies the selected package's width and mode, starts its native +worker, waits for readiness, starts the bridge, then starts and supervises the +C++ HTTP frontend when `--http-port` is present. The receipts directory must be +new or empty, and existing socket paths are never replaced. Keep the launcher +running while serving. It sets `TOKENIZERS_PARALLELISM=false`; concurrent requests +use the bounded preparation pool instead of nested tokenizer worker pools. + +`--direct-worker` connects the C++ HTTP frontend directly to the native worker, +bypassing the Python compatibility bridge. It requires `--http-port` and the +native text-preparation library. The launcher seals tokenizer and transport +identity before startup and writes the private `DIRECT_WORKER.json` configuration +only for the frontend. Audio is copied from the worker's shared slabs into owned +PCM16 before acknowledging reuse. Native inference, request seeds, text +normalization and the public HTTP contract remain unchanged. Omitting the flag +retains the compatibility bridge. In direct mode, `--native-cpus` applies only +to the worker. + +CPU placement affects end-to-end latency even when the CUDA binaries are +identical. The optional `--native-cpus` controls the worker and bridge; +`--frontend-cpus` controls the C++ frontend. Both use Linux `taskset` before +starting child threads, must be subsets of the launcher's allowed CPUs, and +are recorded in `LAUNCH.json` and `EXIT.json`. Unset options inherit affinity. +Inspect `lscpu -e=CPU,CORE,SOCKET,NODE,ONLINE` before selecting CPU IDs. + +On the measured EPYC 7413 host, the following options let preparation use all +eight allocated physical server cores and their SMT siblings, while retaining +four of those cores for native control and the bridge: + +```text +--native-cpus 0-3,24-27 --frontend-cpus 0-7,24-31 +``` + +The container or parent process must allow `0-7,24-31`. The measured client used +separate cores `8-11,32-35`. These IDs describe that host, not a portable CPU +topology. The placement shares existing server cores; it does not add cores or +remove preparation/admission time from TTFA. Keep B512 HTTP worker capacity at +least 512 (`--http-threads 520` in the measurements), because fixed cohorts must +admit a complete physical wave. + +The Docker image uses this launcher as its entrypoint, with the package at +`/package` and `nemo-speech` on `PATH`. Pass `--http-port`, `--http-host`, and +`--http-threads` to start the complete service in one container. Mount model, +model-seal, normalizer assets, configuration, and receipt paths explicitly. + +For text normalization, add `--normalization-config /config/normalization.json`. +The configuration uses the retained normalizer ABI and identifies its native +library and grammar directories. Its `pool_size`, `append_punctuation`, input +limits, and normalization settings are preserved. For comparisons with the +TN-enabled historical runs, use the same grammar hashes and configuration, +including pool size 16 and `NEMOTRON_TN_STRING_MATCHER=1`. Omitting the +configuration disables normalization and preserves trimmed input text. + +This complete Docker example starts B32 in-flight serving with normalization. +Set the three asset paths to absolute paths on your machine. Model assets must +match the model seal; the grammar root must contain the matching `en-US` +directory, including `tokenize_and_classify.far` and `verbalize.far`. These +assets are user-supplied and are not downloaded by the image. + +```bash +NEMOTRON_MODEL_DIR=/absolute/path/to/converted-model +NEMOTRON_MODEL_SEAL=/absolute/path/to/model-seal.json +NEMOTRON_GRAMMAR_DIR=/absolute/path/to/normalizer-assets +NEMOTRON_RUN_DIR="$(mktemp -d)" +mkdir -p "$NEMOTRON_RUN_DIR/config" "$NEMOTRON_RUN_DIR/receipts" +cat > "$NEMOTRON_RUN_DIR/config/normalization.json" <<'JSON' +{ + "enabled": true, + "native_library_root": "/opt/nemotron-speech/lib/processors", + "native_processor": "/opt/nemotron-speech/lib/processors/libnemotron_speech_tts_text_preprocessor.so", + "normalizer_root": "/normalizer", + "normalizer_path": "en-US", + "language_code": "en-US", + "direct_grammar": true, + "append_punctuation": true, + "pool_size": 16, + "max_input_bytes": 65536, + "max_normalized_bytes": 262144, + "max_segments": 256, + "bounded_preparation": true +} +JSON +docker run --rm --gpus device=0 --shm-size=1g --stop-timeout=120 \ + -p 127.0.0.1:8080:8080 \ + --mount "type=bind,src=$NEMOTRON_MODEL_DIR,dst=/model,readonly" \ + --mount "type=bind,src=$NEMOTRON_MODEL_SEAL,dst=/model-seal.json,readonly" \ + --mount "type=bind,src=$NEMOTRON_GRAMMAR_DIR,dst=/normalizer,readonly" \ + --mount "type=bind,src=$NEMOTRON_RUN_DIR/config,dst=/config,readonly" \ + --mount "type=bind,src=$NEMOTRON_RUN_DIR/receipts,dst=/receipts" \ + nemo-speech:nemotron \ + --package /package --model /model --model-seal /model-seal.json \ + --width 32 --mode inflight --listen /tmp/nemotron-b32.sock \ + --receipts /receipts --normalization-config /config/normalization.json \ + --http-host 0.0.0.0 --http-port 8080 --http-threads 40 +``` + +Run the readiness and synthesis requests below from another terminal. Receipts +remain in `$NEMOTRON_RUN_DIR/receipts` after the container stops. Use a fresh run +directory for another launch. For serving without normalization, omit the +grammar/config mounts and `--normalization-config`. The container's native +worker still requires the qualified model and GPU configuration. + +The Nemotron image enables `NEMOTRON_BRIDGE_FORWARD_F32=1`, moving +float32-to-PCM16 conversion to the C++ frontend; public HTTP audio remains +PCM16. Override it with `-e NEMOTRON_BRIDGE_FORWARD_F32=0` to reproduce the +bridge-side conversion control. Outside the image, an unset variable retains +bridge-side conversion. Record the selected values with benchmark receipts. +The [results report](nemotron-integration-results.md) preserves the earlier +candidate measurements and the subsequent complete host-path comparison. + +Alternatively omit `--http-port` from the launcher and start the frontend in +another terminal: + +```bash +build-native/bin/nemo-speech serve \ + --tts-backend nemotron --native-socket /tmp/nemotron-b32.sock \ + --host 127.0.0.1 --port 8080 --http-threads 40 --no-warmup +curl --fail http://127.0.0.1:8080/ready +curl --fail http://127.0.0.1:8080/v1/audio/speech \ + -H 'Content-Type: application/json' \ + -d '{"model":"nemotron-tts","input":"Hello from Nemotron.","voice":"english_sample","response_format":"wav"}' \ + --output /tmp/nemotron.wav +``` + +`response_format: "pcm"` streams S16LE audio with `x-audio-sample-rate: 22050`. +`wav` buffers the completed audio before returning its WAV header. Worker errors +after PCM response headers have been sent abort the HTTP stream; clients must +treat truncated chunked responses as failures. A clean HTTP EOF and nonempty, +sample-aligned audio are required for success. Client disconnects abort the +corresponding backend request. + +The HTTP worker pool must accommodate the intended request concurrency plus +control requests. Suggested initial settings for B1/B32/B512 are 8/40/520 HTTP +workers. Changing HTTP concurrency does not change the package's physical +width. Keep TTS preemption disabled for throughput measurements. + +For diagnostic runs only, set `NEMOTRON_BRIDGE_DIAGNOSTICS` to a new JSONL +path. The bridge records request receipt, preparation, native submission and +first-audio delivery timestamps from one monotonic clock. This does not replace +client TTFA or GPU event timing. Leave this unset during performance collection. + +## Validate and measure + +The mock tests require no model or GPU: + +```bash +python3 -B tests/nemotron_http_test.py --binary build-native/bin/nemo-speech +python3 -B tests/nemotron_normalization_test.py +python3 -B -m unittest discover -s scripts/nemotron_native/tests -v +``` + +The HTTP suite covers first audio before completion, overlapping requests, WAV +metadata, worker failure, recovery, disconnect cancellation, and malformed PCM. +The normalizer test uses the real library when +`NEMOTRON_NORMALIZATION_CONFIG=/config/normalization.json` is set; otherwise its +native ABI test is reported as skipped. + +For GPU qualification, compare the old and new frontend at the same physical +width, package, text, normalized text/token IDs, context, seed, sampling +parameters, precision, and chunk schedule. Compare exact PCM hashes and sample +counts, EOS/length-cap termination, and replay after cancellation and slot +reuse. Include overlapping short/long requests and refills. Fixed-cohort and +in-flight sampling have different RNG contracts, so exact PCM equality is not +an appropriate gate between those two modes. + +The portable benchmark is included in this repository and requires only the +Python standard library. It sends requests directly to the supplied endpoint +and does not start, stop, or reconfigure a server: + +```bash +python3 scripts/nemotron_benchmark.py \ + --url http://127.0.0.1:8080/v1/audio/speech \ + --mode closed-loop --width 32 --concurrency 32 --num-requests 128 \ + --warmup-waves 2 --trials 3 --corpus-tsv /data/corpus.tsv \ + --seed 0 --output /results/portable-b32.json +``` + +`--mode fixed-cohort` gates each wave before connecting and waits for every +request in that wave to terminate before starting the next. `closed-loop` +maintains the configured concurrency through worker-thread refill. `--width` +records the declared physical backend width; it cannot verify that width from +HTTP. The output path must be new. All raw warmup and measured request records, +failures, observed samples, PCM hashes, and read arrivals are retained. + +The default request stopwatch includes TCP/TLS connection establishment. +`--exclude-connect` excludes connection time from request latencies while +keeping it in trial wall time. Trial wall time includes final response drain. +TTFA is arrival of the first complete PCM16 sample. ITL is reconstructed from +the configured 2/3/4/5-frame startup and 8-frame steady boundaries, with 1764 +samples per frame; incomplete terminal chunks contribute audio duration but +not ITL. Read coalescing can produce zero reconstructed intervals. These are +client observations, not server codec timestamps. + +`--timeout` is the per-request budget including connection establishment, +checked at socket-operation and read boundaries. It still includes connection +time when `--exclude-connect` changes the latency stopwatch. Requests already +submitted are drained before the final report is written, including failures +and partial PCM from interrupted streams. An interrupted run retains submitted +requests and records how many planned requests were not submitted. + +Reports contain per-trial nearest-rank percentiles and the equal mean of trial +statistics, including p50/p95/p99 TTFA and completion latency. `rtfx` counts all +received sample-derived audio, including partial failed responses; +`completed_rtfx` counts successful responses only. Failed-request timings and +playback deficits remain visible. GPU memory is not measured by this client; +collect telemetry separately. Its timing and scheduling are distinct from the +historical clients even when connection timing is excluded. + +For comparison against historical results, retain the original external client +and its configuration. The historical clients are not included in this image; +the command below assumes a separately supplied binary at the displayed path. The following `tts_perf_client` B512 example measures +four refill cohorts after two warmup rounds: + +```bash +/opt/nemotron-speech/bin/tts_perf_client \ + --base-url http://127.0.0.1:8080 --transport speech-http \ + --concurrency 512 --num-requests 2048 --warmup-rounds 2 \ + --model easymagpie --voice english_sample --corpus-tsv /data/corpus.tsv \ + --legacy-speech-request --disable-http-pooling \ + --request-timeout-seconds 300 --record-audio-arrivals \ + --codec-frame-samples 1764 --startup-chunk-frames 2,3,4,5 \ + --steady-chunk-frames 8 --json-output /results/inflight-b512.json +``` + +### Direct-worker B512 comparison + +After completing the build, package and normalization setup above, start the +opt-in direct transport with a B512 in-flight package. This example uses the +measured host's eight physical server cores and their SMT siblings; adapt the +CPU IDs to the actual topology. It does not change GPU execution or HTTP timing. + +```bash +export NEMOTRON_NATIVE_PREPARATION_LIBRARY="$PWD/native/text_preparation/target/release/libnemotron_text_preparation.so" +export NEMOTRON_FRONTEND_PREPARATION=1 +export NEMOTRON_BRIDGE_FORWARD_F32=1 +export NEMOTRON_BRIDGE_FAIR_BATCH=64 +export NEMOTRON_HTTP_TCP_NODELAY=1 +export NEMOTRON_TN_STRING_MATCHER=1 +taskset -c 0-7,24-31 python3 scripts/nemotron_native/launch.py \ + --package /tmp/nemotron-package \ + --model /models/nemotron --model-seal /models/model-seal.json \ + --normalization-config /config/normalization.json \ + --width 512 --mode inflight --direct-worker \ + --listen /tmp/nemotron-b512-direct.sock \ + --receipts /results/b512-direct-boot1 \ + --http-host 127.0.0.1 --http-port 8080 --http-threads 520 \ + --native-cpus 0-3,24-27 --frontend-cpus 0-7,24-31 \ + --frontend-binary build-native/bin/nemo-speech +``` + +In a separate terminal, run three unchanged historical-client trials. The +normalizer pool remains 16. The bridge-specific environment settings above are +ignored in direct mode and retained for the matched bridge control. + +```bash +sha256sum /opt/nemotron-speech/bin/tts_perf_client +for trial in 0 1 2; do + taskset -c 8-11,32-35 /opt/nemotron-speech/bin/tts_perf_client \ + --base-url http://127.0.0.1:8080 --transport speech-http \ + --concurrency 512 --num-requests 2048 --warmup-rounds 2 \ + --model easymagpie --voice english_sample --corpus-tsv /data/corpus.tsv \ + --legacy-speech-request --disable-http-pooling \ + --request-timeout-seconds 300 --record-audio-arrivals \ + --codec-frame-samples 1764 --startup-chunk-frames 2,3,4,5 \ + --steady-chunk-frames 8 \ + --json-output "/results/b512-direct-boot1/inflight-${trial}.json" +done +``` + +The reported comparison used client SHA-256 +`d7713f0bbd418c95395959ebd7497bc3c1110c5313da41f0b0dacf019730abbf`. +Stop each launcher cleanly before starting the next arm. For the compatibility +control, omit only `--direct-worker`, choose a new socket and empty receipt +directory, and keep package, assets, CPU masks and client arguments identical. +Collect two fresh boots per arm in bridge/direct/direct/bridge order. Preserve +every trial, then average its TTFA p95 values rather than pooling requests. +The [recorded result](nemotron-integration-results.md#subsequent-direct-worker-transport) +is a single-device comparison and does not establish a ≤200-ms bound. + +For a small width sweep, boot separate B1/B32/B512 controllers and use 32/128/2048 +measured requests respectively. Freeze client binaries and the corpus; the +historical ten-text corpus hash is +`8bae1d20bfd5d17bee363afdbd6d25e0167d4bcc5c11ea6663d75c3b8d9450a2`. +Retain failures, slow requests, audio arrival traces, playback deficits, VRAM +telemetry, and launch/package/model/normalization receipts. Run compared arms +sequentially on the same idle GPU with matched CPU allocations. + +Historical `measure_public_api.py` burst measurements and `tts_perf_client` +closed-loop measurements have different timers and arrival patterns. Report +them separately. The built-in `nemo-speech bench tts` measures synthesis +callbacks without the HTTP transport and must also be labelled separately. +Neither successful streaming nor a performance run establishes acoustic +quality; preserve WAVs and run the same intelligibility and acoustic checks on +both implementations. + +The [integration results](nemotron-integration-results.md) record the completed +checks, measured samples, and remaining limitations separately from these +reproduction instructions. diff --git a/include/nemo_speech/tts.h b/include/nemo_speech/tts.h index 9cb2547..5a56454 100644 --- a/include/nemo_speech/tts.h +++ b/include/nemo_speech/tts.h @@ -71,6 +71,9 @@ typedef struct nemo_speech_tts_model_config { const char* tokenizer_model_dir; // Optional written-to-spoken TN grammar root. Empty keeps text unchanged. const char* text_normalizer_model_dir; + const char* model_backend; + const char* native_socket; + int32_t native_timeout_seconds; } nemo_speech_tts_model_config; typedef struct nemo_speech_tts_runtime_config { diff --git a/native/text_preparation/Cargo.lock b/native/text_preparation/Cargo.lock new file mode 100644 index 0000000..3d5a7b9 --- /dev/null +++ b/native/text_preparation/Cargo.lock @@ -0,0 +1,771 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "getrandom", + "once_cell", + "serde", + "version_check", + "zerocopy", +] + +[[package]] +name = "aho-corasick" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" +dependencies = [ + "memchr", +] + +[[package]] +name = "base64" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" + +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" + +[[package]] +name = "castaway" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a" +dependencies = [ + "rustversion", +] + +[[package]] +name = "cc" +version = "1.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "509591b7bcd67f4ef775afad7662703b4935daaa6ec0e5605cfb1090b32a2b6d" +dependencies = [ + "find-msvc-tools", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "compact_str" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dfdd1c2274d9aa354115b09dc9a901d6c5576818cdf70d14cae2bdb47df00ab" +dependencies = [ + "castaway", + "cfg-if", + "itoa", + "rustversion", + "ryu", + "serde", + "static_assertions", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" + +[[package]] +name = "darling" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +dependencies = [ + "darling_core", + "darling_macro", +] + +[[package]] +name = "darling_core" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn 2.0.119", +] + +[[package]] +name = "darling_macro" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +dependencies = [ + "darling_core", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "dary_heap" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe" +dependencies = [ + "serde", +] + +[[package]] +name = "derive_builder" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947" +dependencies = [ + "derive_builder_macro", +] + +[[package]] +name = "derive_builder_core" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "derive_builder_macro" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c" +dependencies = [ + "derive_builder_core", + "syn 2.0.119", +] + +[[package]] +name = "either" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e5e8f6c15a24b9a3ee5efec809ccd006d3b30e8b3bb63c39af737c7f87daa1d" + +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" + +[[package]] +name = "find-msvc-tools" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d45db016d36b838f563236e9193d0ee6ce38f3f68b6c94e914b4929c96bbb890" + +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "wasip2", +] + +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "libc" +version = "0.2.189" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "macro_rules_attribute" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b3ae8f6d608c795738406608304d30a2dfbdc8e58e44f7ba43236da5208ded3c" +dependencies = [ + "macro_rules_attribute-proc_macro", + "pastey", +] + +[[package]] +name = "macro_rules_attribute-proc_macro" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc04a4c58212d57930a24bf47d3fa87485264a3a054e9c10e042eb373573ad3c" + +[[package]] +name = "media-schema" +version = "0.1.0" +dependencies = [ + "serde", + "serde_bytes", + "thiserror", +] + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "monostate" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67" +dependencies = [ + "monostate-impl", + "serde", + "serde_core", +] + +[[package]] +name = "monostate-impl" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "nemotron-text-preparation" +version = "0.1.0" +dependencies = [ + "media-schema", + "processor-abi", + "regex", + "serde", + "serde_json", + "tokenizers", +] + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "onig" +version = "6.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2" +dependencies = [ + "bitflags", + "libc", + "once_cell", + "onig_sys", +] + +[[package]] +name = "onig_sys" +version = "69.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e68317604e77e53b85896388e1a803c1d21b74c899ec9e5e1112db90735edd7" +dependencies = [ + "cc", + "pkg-config", +] + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "pastey" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4" + +[[package]] +name = "pkg-config" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6b464fbc74e149a392436b17d523f769e057cb6877f6a5c4618bc6f11800548" + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "processor-abi" +version = "0.1.0" +dependencies = [ + "libloading", + "media-schema", + "thiserror", +] + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "rand" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" +dependencies = [ + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom", +] + +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-cond" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f" +dependencies = [ + "either", + "itertools", + "rayon", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_bytes" +version = "0.11.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5d440709e79d88e51ac01c4b72fc6cb7314017bb7da9eeff678aa94c10e3ea8" +dependencies = [ + "serde", + "serde_core", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "serde_json" +version = "1.0.151" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "spm_precompiled" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" +dependencies = [ + "base64", + "nom", + "serde", + "unicode-segmentation", +] + +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "tokenizers" +version = "0.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a620b996116a59e184c2fa2dfd8251ea34a36d0a514758c6f966386bd2e03476" +dependencies = [ + "ahash", + "aho-corasick", + "compact_str", + "dary_heap", + "derive_builder", + "esaxx-rs", + "getrandom", + "itertools", + "log", + "macro_rules_attribute", + "monostate", + "onig", + "paste", + "rand", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unicode-normalization-alignments" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de" +dependencies = [ + "smallvec", +] + +[[package]] +name = "unicode-segmentation" +version = "1.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" + +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "zerocopy" +version = "0.8.56" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "556764e583adb45a9f8d413c2a147fa7e8d821e48e12b14fd560b607998b75eb" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.56" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ab42fc20575779bd240faa45f94a74256f755c0fa9e89f0ede20d91d0cdfc1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/native/text_preparation/Cargo.toml b/native/text_preparation/Cargo.toml new file mode 100644 index 0000000..5c6dcc5 --- /dev/null +++ b/native/text_preparation/Cargo.toml @@ -0,0 +1,37 @@ +[package] +name = "nemotron-text-preparation" +version = "0.1.0" +edition = "2021" +license = "Apache-2.0" + +[lib] +crate-type = ["cdylib", "rlib"] + +[workspace] +members = ["crates/processor-abi", "crates/media-schema"] +resolver = "2" + +[workspace.package] +edition = "2021" +license = "Apache-2.0" +version = "0.1.0" + +[workspace.dependencies] +libloading = "=0.8.9" +serde = { version = "1", features = ["derive"] } +serde_bytes = "0.11" +thiserror = "2" + +[workspace.lints.rust] +unsafe_op_in_unsafe_fn = "deny" + +[workspace.lints.clippy] +all = "warn" + +[dependencies] +processor-abi = { path = "crates/processor-abi" } +media-schema = { path = "crates/media-schema" } +serde = { workspace = true } +serde_json = "1" +regex = "1" +tokenizers = { version = "=0.21.4", default-features = false, features = ["onig"] } diff --git a/native/text_preparation/LICENSE.imported b/native/text_preparation/LICENSE.imported new file mode 100644 index 0000000..009736b --- /dev/null +++ b/native/text_preparation/LICENSE.imported @@ -0,0 +1,7 @@ +Copyright (c) 2020-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. diff --git a/native/text_preparation/NOTICE b/native/text_preparation/NOTICE new file mode 100644 index 0000000..de72a43 --- /dev/null +++ b/native/text_preparation/NOTICE @@ -0,0 +1 @@ +processor-abi and media-schema are adapted from NVIDIA Nemotron Speech commit 2cc27a4e6b06e5e6bf33020ce4faa1fc00bfd25b. Original copyright and license are retained in LICENSE.imported. Only explanatory comments were removed. The native text normalizer remains a separately loaded library under its existing license. diff --git a/native/text_preparation/crates/media-schema/Cargo.toml b/native/text_preparation/crates/media-schema/Cargo.toml new file mode 100644 index 0000000..498e947 --- /dev/null +++ b/native/text_preparation/crates/media-schema/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "media-schema" +edition.workspace = true +license.workspace = true +version.workspace = true + +[dependencies] +serde.workspace = true +serde_bytes.workspace = true +thiserror.workspace = true + +[lints] +workspace = true diff --git a/native/text_preparation/crates/media-schema/src/lib.rs b/native/text_preparation/crates/media-schema/src/lib.rs new file mode 100644 index 0000000..f187e9c --- /dev/null +++ b/native/text_preparation/crates/media-schema/src/lib.rs @@ -0,0 +1,335 @@ +// Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +use std::collections::HashSet; + +use serde::{Deserialize, Serialize}; +use thiserror::Error; + +pub const SCHEMA_VERSION: u32 = 1; +pub const STREAM_ERROR_CODE_RECORD: &str = "stream_error_code"; +pub const STREAM_ERROR_MESSAGE_RECORD: &str = "stream_error_message"; + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct RecordEnvelope { + pub schema_version: u32, + pub request_id: String, + pub session_id: Option, + pub sequence: Option, + pub is_final: bool, + pub records: Vec, +} + +impl RecordEnvelope { + pub fn new(request_id: impl Into, records: Vec) -> Self { + Self { + schema_version: SCHEMA_VERSION, + request_id: request_id.into(), + session_id: None, + sequence: None, + is_final: true, + records, + } + } + + pub fn record(&self, name: &str) -> Option<&MediaRecord> { + self.records + .iter() + .find(|record| record.name == name) + .map(|record| &record.value) + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct NamedRecord { + pub name: String, + pub value: MediaRecord, +} + +impl NamedRecord { + pub fn new(name: impl Into, value: MediaRecord) -> Self { + Self { + name: name.into(), + value, + } + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "kind", content = "value", rename_all = "snake_case")] +pub enum MediaRecord { + Utf8(String), + Bytes(#[serde(with = "serde_bytes")] Vec), + Audio(AudioBuffer), + TokenIds(Vec), + Tensor(Tensor), + Integer(i64), + Float(f64), + Boolean(bool), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum AudioEncoding { + PcmS16Le, +} + +impl AudioEncoding { + fn bytes_per_sample(self) -> usize { + match self { + Self::PcmS16Le => 2, + } + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct AudioBuffer { + pub encoding: AudioEncoding, + pub sample_rate_hz: u32, + pub channels: u16, + #[serde(with = "serde_bytes")] + pub data: Vec, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ElementType { + U8, + I32, + I64, + F32, +} + +impl ElementType { + fn size_bytes(self) -> usize { + match self { + Self::U8 => 1, + Self::I32 | Self::F32 => 4, + Self::I64 => 8, + } + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct Tensor { + pub element_type: ElementType, + pub shape: Vec, + #[serde(with = "serde_bytes")] + pub data: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ValidationLimits { + pub max_request_id_bytes: usize, + pub max_session_id_bytes: usize, + pub max_records: usize, + pub max_record_name_bytes: usize, + pub max_text_bytes: usize, + pub max_binary_bytes: usize, + pub max_audio_bytes: usize, + pub max_tokens: usize, + pub max_tensor_rank: usize, + pub max_tensor_bytes: usize, +} + +impl Default for ValidationLimits { + fn default() -> Self { + Self { + max_request_id_bytes: 256, + max_session_id_bytes: 256, + max_records: 256, + max_record_name_bytes: 128, + max_text_bytes: 256 * 1024, + max_binary_bytes: 16 * 1024 * 1024, + max_audio_bytes: 64 * 1024 * 1024, + max_tokens: 1_048_576, + max_tensor_rank: 8, + max_tensor_bytes: 256 * 1024 * 1024, + } + } +} + +#[derive(Debug, Error, Clone, PartialEq, Eq)] +pub enum ValidationError { + #[error("unsupported schema version {actual}; expected {expected}")] + UnsupportedSchema { actual: u32, expected: u32 }, + #[error("{field} must not be empty")] + Empty { field: &'static str }, + #[error("{field} exceeds the configured limit of {limit}")] + LimitExceeded { field: &'static str, limit: usize }, + #[error("record name is duplicated: {0}")] + DuplicateRecord(String), + #[error("audio configuration is invalid: {0}")] + InvalidAudio(&'static str), + #[error("tensor shape or byte length is invalid")] + InvalidTensor, + #[error("floating-point record is not finite")] + NonFiniteFloat, +} + +pub fn validate_envelope( + envelope: &RecordEnvelope, + limits: &ValidationLimits, +) -> Result<(), ValidationError> { + if envelope.schema_version != SCHEMA_VERSION { + return Err(ValidationError::UnsupportedSchema { + actual: envelope.schema_version, + expected: SCHEMA_VERSION, + }); + } + validate_identifier( + "request_id", + &envelope.request_id, + limits.max_request_id_bytes, + )?; + if let Some(session_id) = &envelope.session_id { + validate_identifier("session_id", session_id, limits.max_session_id_bytes)?; + } + if envelope.records.len() > limits.max_records { + return Err(ValidationError::LimitExceeded { + field: "records", + limit: limits.max_records, + }); + } + + let mut names = HashSet::with_capacity(envelope.records.len()); + for record in &envelope.records { + validate_identifier("record name", &record.name, limits.max_record_name_bytes)?; + if !names.insert(&record.name) { + return Err(ValidationError::DuplicateRecord(record.name.clone())); + } + validate_record(&record.value, limits)?; + } + Ok(()) +} + +pub fn validate_record( + record: &MediaRecord, + limits: &ValidationLimits, +) -> Result<(), ValidationError> { + match record { + MediaRecord::Utf8(text) => { + if text.is_empty() { + return Err(ValidationError::Empty { field: "text" }); + } + check_limit("text", text.len(), limits.max_text_bytes) + } + MediaRecord::Bytes(data) => { + check_limit("binary record", data.len(), limits.max_binary_bytes) + } + MediaRecord::Audio(audio) => validate_audio(audio, limits), + MediaRecord::TokenIds(tokens) => check_limit("token IDs", tokens.len(), limits.max_tokens), + MediaRecord::Tensor(tensor) => validate_tensor(tensor, limits), + MediaRecord::Float(value) if !value.is_finite() => Err(ValidationError::NonFiniteFloat), + MediaRecord::Integer(_) | MediaRecord::Float(_) | MediaRecord::Boolean(_) => Ok(()), + } +} + +fn validate_identifier( + field: &'static str, + value: &str, + limit: usize, +) -> Result<(), ValidationError> { + if value.is_empty() { + return Err(ValidationError::Empty { field }); + } + check_limit(field, value.len(), limit) +} + +fn check_limit(field: &'static str, actual: usize, limit: usize) -> Result<(), ValidationError> { + if actual > limit { + Err(ValidationError::LimitExceeded { field, limit }) + } else { + Ok(()) + } +} + +fn validate_audio(audio: &AudioBuffer, limits: &ValidationLimits) -> Result<(), ValidationError> { + check_limit("audio", audio.data.len(), limits.max_audio_bytes)?; + if audio.sample_rate_hz == 0 { + return Err(ValidationError::InvalidAudio( + "sample rate must be positive", + )); + } + if audio.channels == 0 || audio.channels > 32 { + return Err(ValidationError::InvalidAudio("channel count must be 1-32")); + } + let frame_bytes = audio.encoding.bytes_per_sample() * usize::from(audio.channels); + if !audio.data.len().is_multiple_of(frame_bytes) { + return Err(ValidationError::InvalidAudio( + "payload must contain complete sample frames", + )); + } + Ok(()) +} + +fn validate_tensor(tensor: &Tensor, limits: &ValidationLimits) -> Result<(), ValidationError> { + check_limit("tensor rank", tensor.shape.len(), limits.max_tensor_rank)?; + check_limit("tensor", tensor.data.len(), limits.max_tensor_bytes)?; + let elements = tensor + .shape + .iter() + .try_fold(1usize, |size, dimension| size.checked_mul(*dimension)) + .ok_or(ValidationError::InvalidTensor)?; + let expected = elements + .checked_mul(tensor.element_type.size_bytes()) + .ok_or(ValidationError::InvalidTensor)?; + if expected != tensor.data.len() { + return Err(ValidationError::InvalidTensor); + } + if tensor.element_type == ElementType::F32 { + for sample in tensor.data.chunks_exact(4) { + let value = f32::from_le_bytes(sample.try_into().expect("four-byte chunk")); + if !value.is_finite() { + return Err(ValidationError::NonFiniteFloat); + } + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rejects_duplicate_record_names() { + let envelope = RecordEnvelope::new( + "request-1", + vec![ + NamedRecord::new("text", MediaRecord::Utf8("one".into())), + NamedRecord::new("text", MediaRecord::Utf8("two".into())), + ], + ); + assert_eq!( + validate_envelope(&envelope, &ValidationLimits::default()), + Err(ValidationError::DuplicateRecord("text".into())) + ); + } + + #[test] + fn rejects_malformed_audio_frames() { + let audio = MediaRecord::Audio(AudioBuffer { + encoding: AudioEncoding::PcmS16Le, + sample_rate_hz: 16_000, + channels: 1, + data: vec![0], + }); + assert!(matches!( + validate_record(&audio, &ValidationLimits::default()), + Err(ValidationError::InvalidAudio(_)) + )); + } + + #[test] + fn rejects_tensor_size_mismatch() { + let tensor = MediaRecord::Tensor(Tensor { + element_type: ElementType::I64, + shape: vec![2, 2], + data: vec![0; 24], + }); + assert_eq!( + validate_record(&tensor, &ValidationLimits::default()), + Err(ValidationError::InvalidTensor) + ); + } +} diff --git a/native/text_preparation/crates/processor-abi/Cargo.toml b/native/text_preparation/crates/processor-abi/Cargo.toml new file mode 100644 index 0000000..5c931f0 --- /dev/null +++ b/native/text_preparation/crates/processor-abi/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "processor-abi" +edition.workspace = true +license.workspace = true +version.workspace = true + +[dependencies] +libloading.workspace = true +media-schema = { path = "../media-schema" } +thiserror.workspace = true + +[lints] +workspace = true diff --git a/native/text_preparation/crates/processor-abi/src/lib.rs b/native/text_preparation/crates/processor-abi/src/lib.rs new file mode 100644 index 0000000..bea496f --- /dev/null +++ b/native/text_preparation/crates/processor-abi/src/lib.rs @@ -0,0 +1,678 @@ +// Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +pub mod raw; + +use std::{ + mem::size_of, + path::Path, + slice, + sync::{Arc, Condvar, Mutex}, + time::{Duration, Instant}, +}; + +use libloading::Library; +use media_schema::{ + AudioBuffer, AudioEncoding, ElementType, MediaRecord, NamedRecord, RecordEnvelope, Tensor, +}; +use thiserror::Error; + +use raw::{ + ByteSpanV1, EnvelopeOutputV1, EnvelopeViewV1, ProcessorApiV1, ProcessorCapabilitiesV1, + ProcessorHandleV1, RecordViewV1, StatusV1, +}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProcessorCapabilities { + pub processor_id: String, + pub batch_input: bool, + pub streaming_input: bool, + pub thread_safe: bool, + pub stateful: bool, +} + +#[derive(Debug, Error)] +pub enum ProcessorError { + #[error("unable to load native processor: {0}")] + Load(String), + #[error("processor ABI version or layout is incompatible")] + IncompatibleAbi, + #[error("native processor returned status {code}: {message}")] + Native { code: i32, message: String }, + #[error("native processor returned malformed output: {0}")] + MalformedOutput(String), + #[error("native processor received malformed input: {0}")] + MalformedInput(String), +} + +struct NativeHandle(ProcessorHandleV1); + +unsafe impl Send for NativeHandle {} + +struct ConcurrentNativeHandle(NativeHandle); + +unsafe impl Sync for ConcurrentNativeHandle {} + +enum NativeHandleAccess { + Serialized(Mutex), + Concurrent(ConcurrentNativeHandle), +} + +pub struct NativeProcessor { + _library: Library, + api: &'static ProcessorApiV1, + handle: NativeHandleAccess, + capabilities: ProcessorCapabilities, +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct ProcessorTimings { + pub queue_wait: Duration, + pub active: Duration, +} + +pub struct NativeProcessorPool { + processors: Vec>, + available: Mutex>, + ready: Condvar, +} + +impl NativeProcessorPool { + pub unsafe fn load( + path: &Path, + config: &RecordEnvelope, + instances: usize, + ) -> Result, ProcessorError> { + if instances == 0 { + return Err(ProcessorError::Load( + "processor pool must contain at least one instance".into(), + )); + } + let first = unsafe { NativeProcessor::load(path, config) }?; + let actual_instances = if first.capabilities().thread_safe { + 1 + } else { + instances + }; + let mut processors = Vec::with_capacity(actual_instances); + processors.push(first); + for _ in 1..actual_instances { + processors.push(unsafe { NativeProcessor::load(path, config) }?); + } + let available = (0..actual_instances).rev().collect(); + Ok(Arc::new(Self { + processors, + available: Mutex::new(available), + ready: Condvar::new(), + })) + } + + pub fn capabilities(&self) -> &ProcessorCapabilities { + self.processors[0].capabilities() + } + + pub fn size(&self) -> usize { + self.processors.len() + } + + pub fn process(&self, request: &RecordEnvelope) -> Result { + self.process_with_timings(request).map(|(output, _)| output) + } + + pub fn process_with_timings( + &self, + request: &RecordEnvelope, + ) -> Result<(RecordEnvelope, ProcessorTimings), ProcessorError> { + if self.capabilities().thread_safe { + let active_started = Instant::now(); + let output = self.processors[0].process(request)?; + return Ok(( + output, + ProcessorTimings { + queue_wait: Duration::ZERO, + active: active_started.elapsed(), + }, + )); + } + let queue_started = Instant::now(); + let index = { + let mut available = self.available.lock().expect("processor pool lock poisoned"); + while available.is_empty() { + available = self + .ready + .wait(available) + .expect("processor pool lock poisoned"); + } + available.pop().expect("availability checked") + }; + let queue_wait = queue_started.elapsed(); + let _checkout = ProcessorCheckout { pool: self, index }; + let active_started = Instant::now(); + let output = self.processors[index].process(request)?; + Ok(( + output, + ProcessorTimings { + queue_wait, + active: active_started.elapsed(), + }, + )) + } +} + +struct ProcessorCheckout<'a> { + pool: &'a NativeProcessorPool, + index: usize, +} + +impl Drop for ProcessorCheckout<'_> { + fn drop(&mut self) { + self.pool + .available + .lock() + .expect("processor pool lock poisoned") + .push(self.index); + self.pool.ready.notify_one(); + } +} + +impl NativeProcessor { + pub unsafe fn load(path: &Path, config: &RecordEnvelope) -> Result, ProcessorError> { + let library = unsafe { Library::new(path) } + .map_err(|error| ProcessorError::Load(error.to_string()))?; + let get_api: libloading::Symbol<'_, raw::GetProcessorApi> = unsafe { + library + .get(raw::API_SYMBOL) + .map_err(|error| ProcessorError::Load(error.to_string()))? + }; + let api_ptr = unsafe { get_api(raw::PROCESSOR_ABI_VERSION) }; + if api_ptr.is_null() { + return Err(ProcessorError::IncompatibleAbi); + } + let api = unsafe { &*api_ptr }; + if api.abi_version != raw::PROCESSOR_ABI_VERSION + || api.struct_size < size_of::() + { + return Err(ProcessorError::IncompatibleAbi); + } + + let capabilities = read_capabilities(api)?; + let borrowed = BorrowedEnvelope::new(config)?; + let mut handle = std::ptr::null_mut(); + let status = unsafe { (api.create)(&borrowed.view, &mut handle) }; + check_status(api, status)?; + if handle.is_null() { + return Err(ProcessorError::MalformedOutput( + "create returned a null handle".into(), + )); + } + let api = unsafe { &*(api as *const ProcessorApiV1) }; + let handle = if capabilities.thread_safe { + NativeHandleAccess::Concurrent(ConcurrentNativeHandle(NativeHandle(handle))) + } else { + NativeHandleAccess::Serialized(Mutex::new(NativeHandle(handle))) + }; + Ok(Arc::new(Self { + _library: library, + api, + handle, + capabilities, + })) + } + + pub fn capabilities(&self) -> &ProcessorCapabilities { + &self.capabilities + } + + pub fn process(&self, request: &RecordEnvelope) -> Result { + let borrowed = BorrowedEnvelope::new(request)?; + match &self.handle { + NativeHandleAccess::Serialized(handle) => { + let handle = handle.lock().expect("native processor lock poisoned"); + self.process_handle(handle.0, &borrowed) + } + NativeHandleAccess::Concurrent(handle) => self.process_handle((handle.0).0, &borrowed), + } + } + + fn process_handle( + &self, + handle: ProcessorHandleV1, + borrowed: &BorrowedEnvelope, + ) -> Result { + let mut output = raw::empty_output(); + let status = unsafe { (self.api.process)(handle, &borrowed.view, &mut output) }; + check_status(self.api, status)?; + let converted = copy_output(&output); + unsafe { (self.api.release_output)(handle, &mut output) }; + converted + } +} + +impl Drop for NativeProcessor { + fn drop(&mut self) { + match &mut self.handle { + NativeHandleAccess::Serialized(handle) => { + if let Ok(handle) = handle.get_mut() { + destroy_handle(self.api, handle); + } + } + NativeHandleAccess::Concurrent(handle) => destroy_handle(self.api, &mut handle.0), + } + } +} + +fn destroy_handle(api: &ProcessorApiV1, handle: &mut NativeHandle) { + if !handle.0.is_null() { + unsafe { (api.destroy)(handle.0) }; + handle.0 = std::ptr::null_mut(); + } +} + +struct BorrowedRecord { + name: Vec, + data: Vec, + shape: Vec, + view: RecordViewV1, +} + +impl BorrowedRecord { + fn new(record: &NamedRecord) -> Result { + let name = record.name.as_bytes().to_vec(); + let (kind, element_type, data, shape, sample_rate_hz, channels, audio_encoding) = + encode_record(&record.value); + let mut value = Self { + name, + data, + shape, + view: RecordViewV1 { + struct_size: size_of::(), + name: ByteSpanV1 { + data: std::ptr::null(), + len: 0, + }, + kind, + element_type, + data: ByteSpanV1 { + data: std::ptr::null(), + len: 0, + }, + shape: std::ptr::null(), + rank: 0, + sample_rate_hz, + channels, + audio_encoding, + sequence: 0, + has_sequence: 0, + is_final: 0, + reserved: [0; 6], + }, + }; + value.view.name = span(&value.name); + value.view.data = span(&value.data); + value.view.shape = value.shape.as_ptr(); + value.view.rank = value.shape.len(); + Ok(value) + } +} + +struct BorrowedEnvelope { + request_id: Vec, + session_id: Vec, + records: Vec, + views: Vec, + view: EnvelopeViewV1, +} + +impl BorrowedEnvelope { + fn new(envelope: &RecordEnvelope) -> Result { + let request_id = envelope.request_id.as_bytes().to_vec(); + let session_id = envelope + .session_id + .as_deref() + .unwrap_or_default() + .as_bytes() + .to_vec(); + let records = envelope + .records + .iter() + .map(BorrowedRecord::new) + .collect::, _>>()?; + let views = records.iter().map(|record| record.view).collect::>(); + let mut value = Self { + request_id, + session_id, + records, + views, + view: EnvelopeViewV1 { + struct_size: size_of::(), + schema_version: envelope.schema_version, + request_id: ByteSpanV1 { + data: std::ptr::null(), + len: 0, + }, + session_id: ByteSpanV1 { + data: std::ptr::null(), + len: 0, + }, + records: std::ptr::null(), + record_count: 0, + sequence: envelope.sequence.unwrap_or_default(), + has_sequence: u8::from(envelope.sequence.is_some()), + is_final: u8::from(envelope.is_final), + reserved: [0; 6], + }, + }; + value.view.request_id = span(&value.request_id); + value.view.session_id = span(&value.session_id); + value.view.records = value.views.as_ptr(); + value.view.record_count = value.views.len(); + debug_assert_eq!(value.records.len(), value.views.len()); + Ok(value) + } +} + +fn encode_record(record: &MediaRecord) -> (u32, u32, Vec, Vec, u32, u16, u16) { + match record { + MediaRecord::Utf8(value) => ( + raw::RECORD_UTF8, + raw::ELEMENT_U8, + value.as_bytes().to_vec(), + Vec::new(), + 0, + 0, + raw::AUDIO_NONE as u16, + ), + MediaRecord::Bytes(value) => ( + raw::RECORD_BYTES, + raw::ELEMENT_U8, + value.clone(), + Vec::new(), + 0, + 0, + raw::AUDIO_NONE as u16, + ), + MediaRecord::Audio(value) => ( + raw::RECORD_AUDIO, + raw::ELEMENT_U8, + value.data.clone(), + Vec::new(), + value.sample_rate_hz, + value.channels, + match value.encoding { + AudioEncoding::PcmS16Le => raw::AUDIO_PCM_S16_LE as u16, + }, + ), + MediaRecord::TokenIds(value) => ( + raw::RECORD_TOKEN_IDS, + raw::ELEMENT_I32, + value.iter().flat_map(|item| item.to_le_bytes()).collect(), + vec![value.len()], + 0, + 0, + raw::AUDIO_NONE as u16, + ), + MediaRecord::Tensor(value) => ( + raw::RECORD_TENSOR, + match value.element_type { + ElementType::U8 => raw::ELEMENT_U8, + ElementType::I32 => raw::ELEMENT_I32, + ElementType::I64 => raw::ELEMENT_I64, + ElementType::F32 => raw::ELEMENT_F32, + }, + value.data.clone(), + value.shape.clone(), + 0, + 0, + raw::AUDIO_NONE as u16, + ), + MediaRecord::Integer(value) => ( + raw::RECORD_INTEGER, + raw::ELEMENT_I64, + value.to_le_bytes().to_vec(), + Vec::new(), + 0, + 0, + raw::AUDIO_NONE as u16, + ), + MediaRecord::Float(value) => ( + raw::RECORD_FLOAT, + raw::ELEMENT_I64, + value.to_le_bytes().to_vec(), + Vec::new(), + 0, + 0, + raw::AUDIO_NONE as u16, + ), + MediaRecord::Boolean(value) => ( + raw::RECORD_BOOLEAN, + raw::ELEMENT_U8, + vec![u8::from(*value)], + Vec::new(), + 0, + 0, + raw::AUDIO_NONE as u16, + ), + } +} + +fn read_capabilities(api: &ProcessorApiV1) -> Result { + let mut raw_capabilities = ProcessorCapabilitiesV1 { + struct_size: size_of::(), + processor_id: ByteSpanV1 { + data: std::ptr::null(), + len: 0, + }, + batch_input: 0, + streaming_input: 0, + thread_safe: 0, + stateful: 0, + reserved: [0; 4], + }; + let status = unsafe { (api.capabilities)(&mut raw_capabilities) }; + check_status(api, status)?; + Ok(ProcessorCapabilities { + processor_id: copy_utf8(raw_capabilities.processor_id, "processor ID")?, + batch_input: raw_capabilities.batch_input != 0, + streaming_input: raw_capabilities.streaming_input != 0, + thread_safe: raw_capabilities.thread_safe != 0, + stateful: raw_capabilities.stateful != 0, + }) +} + +fn check_status(api: &ProcessorApiV1, mut status: StatusV1) -> Result<(), ProcessorError> { + let result = if status.code == 0 { + Ok(()) + } else { + Err(ProcessorError::Native { + code: status.code, + message: copy_utf8(status.message, "status message") + .unwrap_or_else(|_| "invalid native error message".into()), + }) + }; + unsafe { (api.release_status)(&mut status) }; + result +} + +fn copy_output(output: &EnvelopeOutputV1) -> Result { + if output.struct_size < size_of::() { + return Err(ProcessorError::MalformedOutput( + "output structure is too small".into(), + )); + } + let request_id = copy_utf8(output.request_id, "request ID")?; + let session_id = if output.session_id.len == 0 { + None + } else { + Some(copy_utf8(output.session_id, "session ID")?) + }; + if output.record_count > 0 && output.records.is_null() { + return Err(ProcessorError::MalformedOutput( + "record array is null".into(), + )); + } + let views = if output.record_count == 0 { + &[] + } else { + unsafe { slice::from_raw_parts(output.records, output.record_count) } + }; + let records = views + .iter() + .map(copy_record) + .collect::, _>>()?; + Ok(RecordEnvelope { + schema_version: output.schema_version, + request_id, + session_id, + sequence: (output.has_sequence != 0).then_some(output.sequence), + is_final: output.is_final != 0, + records, + }) +} + +fn copy_record(view: &RecordViewV1) -> Result { + if view.struct_size < size_of::() { + return Err(ProcessorError::MalformedOutput( + "record structure is too small".into(), + )); + } + let name = copy_utf8(view.name, "record name")?; + let data = copy_bytes(view.data, "record data")?; + let value = match view.kind { + raw::RECORD_UTF8 => MediaRecord::Utf8( + String::from_utf8(data) + .map_err(|_| ProcessorError::MalformedOutput("invalid UTF-8 record".into()))?, + ), + raw::RECORD_BYTES => MediaRecord::Bytes(data), + raw::RECORD_AUDIO => MediaRecord::Audio(AudioBuffer { + encoding: match u32::from(view.audio_encoding) { + raw::AUDIO_PCM_S16_LE => AudioEncoding::PcmS16Le, + _ => { + return Err(ProcessorError::MalformedOutput( + "unknown audio encoding".into(), + )) + } + }, + sample_rate_hz: view.sample_rate_hz, + channels: view.channels, + data, + }), + raw::RECORD_TOKEN_IDS => { + if data.len() % 4 != 0 { + return Err(ProcessorError::MalformedOutput( + "token byte length is invalid".into(), + )); + } + MediaRecord::TokenIds( + data.chunks_exact(4) + .map(|bytes| u32::from_le_bytes(bytes.try_into().expect("four-byte chunk"))) + .collect(), + ) + } + raw::RECORD_TENSOR => { + if view.rank > 0 && view.shape.is_null() { + return Err(ProcessorError::MalformedOutput( + "tensor shape is null".into(), + )); + } + let shape = if view.rank == 0 { + Vec::new() + } else { + unsafe { slice::from_raw_parts(view.shape, view.rank) }.to_vec() + }; + MediaRecord::Tensor(Tensor { + element_type: match view.element_type { + raw::ELEMENT_U8 => ElementType::U8, + raw::ELEMENT_I32 => ElementType::I32, + raw::ELEMENT_I64 => ElementType::I64, + raw::ELEMENT_F32 => ElementType::F32, + _ => { + return Err(ProcessorError::MalformedOutput( + "unknown tensor element type".into(), + )) + } + }, + shape, + data, + }) + } + raw::RECORD_INTEGER if data.len() == 8 => MediaRecord::Integer(i64::from_le_bytes( + data.as_slice().try_into().expect("eight-byte value"), + )), + raw::RECORD_FLOAT if data.len() == 8 => MediaRecord::Float(f64::from_le_bytes( + data.as_slice().try_into().expect("eight-byte value"), + )), + raw::RECORD_BOOLEAN if data.len() == 1 => MediaRecord::Boolean(data[0] != 0), + _ => { + return Err(ProcessorError::MalformedOutput(format!( + "unknown or malformed record kind {}", + view.kind + ))) + } + }; + Ok(NamedRecord { name, value }) +} + +fn span(bytes: &[u8]) -> ByteSpanV1 { + ByteSpanV1 { + data: bytes.as_ptr(), + len: bytes.len(), + } +} + +fn copy_utf8(value: ByteSpanV1, field: &str) -> Result { + String::from_utf8(copy_bytes(value, field)?) + .map_err(|_| ProcessorError::MalformedOutput(format!("{field} is not valid UTF-8"))) +} + +fn copy_bytes(value: ByteSpanV1, field: &str) -> Result, ProcessorError> { + if value.len == 0 { + return Ok(Vec::new()); + } + if value.data.is_null() { + return Err(ProcessorError::MalformedOutput(format!( + "{field} pointer is null" + ))); + } + Ok(unsafe { slice::from_raw_parts(value.data, value.len) }.to_vec()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn zero_length_abi_arrays_accept_null_pointers() { + let mut output = raw::empty_output(); + output.schema_version = media_schema::SCHEMA_VERSION; + let copied = copy_output(&output).unwrap(); + assert!(copied.records.is_empty()); + + let name = b"tensor"; + let tensor = RecordViewV1 { + struct_size: size_of::(), + name: ByteSpanV1 { + data: name.as_ptr(), + len: name.len(), + }, + kind: raw::RECORD_TENSOR, + element_type: raw::ELEMENT_U8, + data: ByteSpanV1 { + data: std::ptr::null(), + len: 0, + }, + shape: std::ptr::null(), + rank: 0, + sample_rate_hz: 0, + channels: 0, + audio_encoding: raw::AUDIO_NONE as u16, + sequence: 0, + has_sequence: 0, + is_final: 0, + reserved: [0; 6], + }; + let copied = copy_record(&tensor).unwrap(); + assert!(matches!( + copied.value, + MediaRecord::Tensor(Tensor { shape, data, .. }) if shape.is_empty() && data.is_empty() + )); + } +} diff --git a/native/text_preparation/crates/processor-abi/src/raw.rs b/native/text_preparation/crates/processor-abi/src/raw.rs new file mode 100644 index 0000000..4819ce9 --- /dev/null +++ b/native/text_preparation/crates/processor-abi/src/raw.rs @@ -0,0 +1,156 @@ +// Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +use std::{ + ffi::{c_char, c_void}, + mem::size_of, +}; + +pub const PROCESSOR_ABI_VERSION: u32 = 1; + +pub const STATUS_OK: i32 = 0; +pub const STATUS_INTERNAL: i32 = 1; +pub const STATUS_INVALID_ARGUMENT: i32 = 2; + +pub const RECORD_UTF8: u32 = 1; +pub const RECORD_BYTES: u32 = 2; +pub const RECORD_AUDIO: u32 = 3; +pub const RECORD_TOKEN_IDS: u32 = 4; +pub const RECORD_TENSOR: u32 = 5; +pub const RECORD_INTEGER: u32 = 6; +pub const RECORD_FLOAT: u32 = 7; +pub const RECORD_BOOLEAN: u32 = 8; + +pub const ELEMENT_NONE: u32 = 0; +pub const ELEMENT_U8: u32 = 1; +pub const ELEMENT_I32: u32 = 2; +pub const ELEMENT_I64: u32 = 3; +pub const ELEMENT_F32: u32 = 4; + +pub const AUDIO_NONE: u32 = 0; +pub const AUDIO_PCM_S16_LE: u32 = 1; +pub const AUDIO_PCM_F32_LE: u32 = 2; + +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct ByteSpanV1 { + pub data: *const u8, + pub len: usize, +} + +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct StatusV1 { + pub struct_size: usize, + pub code: i32, + pub message: ByteSpanV1, + pub context: *mut c_void, +} + +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct RecordViewV1 { + pub struct_size: usize, + pub name: ByteSpanV1, + pub kind: u32, + pub element_type: u32, + pub data: ByteSpanV1, + pub shape: *const usize, + pub rank: usize, + pub sample_rate_hz: u32, + pub channels: u16, + pub audio_encoding: u16, + pub sequence: u64, + pub has_sequence: u8, + pub is_final: u8, + pub reserved: [u8; 6], +} + +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct EnvelopeViewV1 { + pub struct_size: usize, + pub schema_version: u32, + pub request_id: ByteSpanV1, + pub session_id: ByteSpanV1, + pub records: *const RecordViewV1, + pub record_count: usize, + pub sequence: u64, + pub has_sequence: u8, + pub is_final: u8, + pub reserved: [u8; 6], +} + +#[repr(C)] +#[derive(Debug)] +pub struct EnvelopeOutputV1 { + pub struct_size: usize, + pub schema_version: u32, + pub request_id: ByteSpanV1, + pub session_id: ByteSpanV1, + pub records: *const RecordViewV1, + pub record_count: usize, + pub sequence: u64, + pub has_sequence: u8, + pub is_final: u8, + pub reserved: [u8; 6], + pub context: *mut c_void, +} + +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct ProcessorCapabilitiesV1 { + pub struct_size: usize, + pub processor_id: ByteSpanV1, + pub batch_input: u8, + pub streaming_input: u8, + pub thread_safe: u8, + pub stateful: u8, + pub reserved: [u8; 4], +} + +pub type ProcessorHandleV1 = *mut c_void; + +#[repr(C)] +pub struct ProcessorApiV1 { + pub struct_size: usize, + pub abi_version: u32, + pub capabilities: unsafe extern "C" fn(*mut ProcessorCapabilitiesV1) -> StatusV1, + pub create: unsafe extern "C" fn(*const EnvelopeViewV1, *mut ProcessorHandleV1) -> StatusV1, + pub destroy: unsafe extern "C" fn(ProcessorHandleV1), + pub process: unsafe extern "C" fn( + ProcessorHandleV1, + *const EnvelopeViewV1, + *mut EnvelopeOutputV1, + ) -> StatusV1, + pub release_output: unsafe extern "C" fn(ProcessorHandleV1, *mut EnvelopeOutputV1), + pub release_status: unsafe extern "C" fn(*mut StatusV1), +} + +pub type GetProcessorApi = unsafe extern "C" fn(u32) -> *const ProcessorApiV1; + +pub const API_SYMBOL: &[u8] = b"nemotron_speech_processor_get_api\0"; + +pub fn empty_output() -> EnvelopeOutputV1 { + EnvelopeOutputV1 { + struct_size: size_of::(), + schema_version: 0, + request_id: ByteSpanV1 { + data: std::ptr::null(), + len: 0, + }, + session_id: ByteSpanV1 { + data: std::ptr::null(), + len: 0, + }, + records: std::ptr::null(), + record_count: 0, + sequence: 0, + has_sequence: 0, + is_final: 0, + reserved: [0; 6], + context: std::ptr::null_mut(), + } +} + +#[allow(dead_code)] +const _: *const c_char = std::ptr::null(); diff --git a/native/text_preparation/src/lib.rs b/native/text_preparation/src/lib.rs new file mode 100644 index 0000000..0d6dbde --- /dev/null +++ b/native/text_preparation/src/lib.rs @@ -0,0 +1,464 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +use media_schema::{MediaRecord, NamedRecord, RecordEnvelope}; +use processor_abi::{NativeProcessorPool, ProcessorError}; +use regex::Regex; +use serde::Deserialize; +use std::{ + ffi::{c_char, c_void, CStr, CString}, + panic::{catch_unwind, AssertUnwindSafe}, + path::{Path, PathBuf}, + ptr, slice, + sync::Arc, +}; +use tokenizers::Tokenizer; + +type Outcome = Result; +fn invalid(message: impl Into) -> (i32, String) { + (1, message.into()) +} +fn internal(message: impl Into) -> (i32, String) { + (2, message.into()) +} +fn yes() -> bool { + true +} +fn language() -> String { + "en-US".into() +} +fn config_name() -> String { + "sparrowhawk_configuration.ascii_proto".into() +} +fn pool() -> usize { + 2 +} +fn input_limit() -> usize { + 65536 +} +fn output_limit() -> usize { + 262144 +} +fn segments_limit() -> usize { + 256 +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct Config { + #[serde(default = "yes")] + enabled: bool, + native_library_root: PathBuf, + native_processor: PathBuf, + normalizer_root: PathBuf, + #[serde(default)] + normalizer_path: Option, + #[serde(default = "language")] + language_code: String, + #[serde(default = "config_name")] + normalizer_config: String, + #[serde(default)] + direct_grammar: bool, + #[serde(default)] + abbreviations_path: Option, + #[serde(default = "yes")] + append_punctuation: bool, + #[serde(default = "pool")] + pool_size: usize, + #[serde(default)] + bounded_preparation: bool, + #[serde(default = "input_limit")] + max_input_bytes: usize, + #[serde(default = "output_limit")] + max_normalized_bytes: usize, + #[serde(default = "segments_limit")] + max_segments: usize, +} + +struct Preparation { + tokenizer: Tokenizer, + normalizer: Option>, + enabled: bool, + max_input: usize, + max_output: usize, + max_segments: usize, + punctuation: Regex, +} + +fn canonical(path: &Path, directory: bool) -> Outcome { + let path = path + .canonicalize() + .map_err(|e| invalid(format!("{}: {e}", path.display())))?; + if if directory { + !path.is_dir() + } else { + !path.is_file() + } { + return Err(invalid("normalizer path has the wrong file type")); + } + Ok(path) +} +fn contained(root: &Path, path: &Path, directory: bool) -> Outcome { + let path = canonical(&root.join(path), directory)?; + if !path.starts_with(root) { + return Err(invalid("normalizer path escapes its root")); + } + Ok(path) +} +fn utf8_path(path: PathBuf) -> Outcome { + path.into_os_string() + .into_string() + .map_err(|_| invalid("normalizer path must be UTF-8")) +} +fn processor_error(error: ProcessorError) -> (i32, String) { + let code = if matches!( + error, + ProcessorError::Native { + code: 100..=102, + .. + } + ) { + 1 + } else { + 2 + }; + (code, error.to_string()) +} +fn whitespace(c: char) -> bool { + c.is_whitespace() || ('\u{1c}'..='\u{1f}').contains(&c) +} + +impl Preparation { + fn new(config: Option, tokenizer_path: &str) -> Outcome { + let tokenizer = Tokenizer::from_file(tokenizer_path) + .map_err(|e| invalid(format!("tokenizer failed: {e}")))?; + let mut result = Self { + tokenizer, + normalizer: None, + enabled: false, + max_input: input_limit(), + max_output: output_limit(), + max_segments: segments_limit(), + punctuation: Regex::new(r" +([,.?!।])").map_err(|e| internal(e.to_string()))?, + }; + if let Some(c) = config { + let _ = c.bounded_preparation; + if c.language_code != "en-US" { + return Err(invalid("normalization language_code must be en-US")); + } + for n in [ + c.max_input_bytes, + c.max_normalized_bytes, + c.max_segments, + c.pool_size, + ] { + if n == 0 || n > i64::MAX as usize { + return Err(invalid( + "normalization limits must fit positive ABI integers", + )); + } + } + let library_root = canonical(&c.native_library_root, true)?; + let asset_root = canonical(&c.normalizer_root, true)?; + if library_root.starts_with(&asset_root) || asset_root.starts_with(&library_root) { + return Err(invalid( + "normalizer libraries and assets must have separate roots", + )); + } + let library = canonical(&c.native_processor, false)?; + if !library.starts_with(&library_root) { + return Err(invalid("native_processor escapes native_library_root")); + } + let language_path = contained( + &asset_root, + c.normalizer_path.as_deref().unwrap_or(Path::new("en-US")), + true, + )?; + let name = Path::new(&c.normalizer_config); + if name.components().count() != 1 + || name.file_name().and_then(|n| n.to_str()) != Some(c.normalizer_config.as_str()) + { + return Err(invalid( + "normalizer_config must name a file in normalizer_path", + )); + } + let required: Vec<&str> = if c.direct_grammar { + vec!["tokenize_and_classify.far", "verbalize.far"] + } else { + vec![&c.normalizer_config] + }; + for name in required { + contained(&language_path, Path::new(name), false)?; + } + if c.direct_grammar { + for name in ["pre_process.far", "post_process.far", "speech_class.far"] { + let p = language_path.join(name); + if p.exists() || p.is_symlink() { + contained(&language_path, Path::new(name), false)?; + } + } + } + let mut records = vec![ + NamedRecord::new("language_code", MediaRecord::Utf8("en-US".into())), + NamedRecord::new( + "normalizer_path", + MediaRecord::Utf8(utf8_path(language_path.clone())?), + ), + NamedRecord::new("normalizer_config", MediaRecord::Utf8(c.normalizer_config)), + NamedRecord::new("direct_grammar", MediaRecord::Boolean(c.direct_grammar)), + NamedRecord::new( + "append_punctuation", + MediaRecord::Boolean(c.append_punctuation), + ), + ]; + for (key, value) in [ + ("max_input_bytes", c.max_input_bytes), + ("max_normalized_bytes", c.max_normalized_bytes), + ("max_segments", c.max_segments), + ] { + records.push(NamedRecord::new(key, MediaRecord::Integer(value as i64))); + } + if let Some(path) = c.abbreviations_path { + let path = canonical(&language_path.join(path), false)?; + if !path.starts_with(&asset_root) { + return Err(invalid("abbreviations_path escapes normalizer_root")); + } + records.push(NamedRecord::new( + "abbreviations_path", + MediaRecord::Utf8(utf8_path(path)?), + )); + } + let normalizer = unsafe { + NativeProcessorPool::load( + &library, + &RecordEnvelope::new("nemotron-cpp-normalization", records), + c.pool_size, + ) + } + .map_err(processor_error)?; + if normalizer.capabilities().processor_id != "nemotron_speech.tts.text_preprocessor" + || normalizer.capabilities().stateful + { + return Err(invalid( + "native processor is not the stateless TTS normalizer", + )); + } + result.normalizer = Some(normalizer); + result.enabled = c.enabled; + result.max_input = c.max_input_bytes; + result.max_output = c.max_normalized_bytes; + result.max_segments = c.max_segments; + } + Ok(result) + } + fn prepare(&self, input: &str, apply: bool) -> Outcome { + let text = input.trim_matches(whitespace); + if text.is_empty() || text.len() > self.max_input { + return Err(invalid("text is empty or exceeds max_input_bytes")); + } + let normalized = if self.enabled && apply { + let request = RecordEnvelope::new( + "nemotron-cpp-normalization", + vec![ + NamedRecord::new("text", MediaRecord::Utf8(text.into())), + NamedRecord::new("language_code", MediaRecord::Utf8("en-US".into())), + ], + ); + let output = self + .normalizer + .as_ref() + .ok_or_else(|| internal("normalizer unavailable"))? + .process(&request) + .map_err(processor_error)?; + if output.schema_version != request.schema_version + || output.request_id != request.request_id + || output.session_id != request.session_id + || output.sequence != request.sequence + || output.is_final != request.is_final + { + return Err(internal("native normalizer changed envelope identity")); + } + if output.records.is_empty() || output.records.len() > self.max_segments { + return Err(internal("native normalizer returned invalid segment count")); + } + let mut bytes = output.records.len() - 1; + let mut segments = Vec::with_capacity(output.records.len()); + for (index, record) in output.records.into_iter().enumerate() { + if record.name != format!("normalized_text.{index}") { + return Err(internal( + "native normalizer returned invalid segment order/type", + )); + } + let MediaRecord::Utf8(segment) = record.value else { + return Err(internal("native normalizer text is not UTF-8")); + }; + bytes = bytes + .checked_add(segment.len()) + .ok_or_else(|| internal("normalizer length overflow"))?; + if segment.trim_matches(whitespace).is_empty() || bytes > self.max_output { + return Err(internal( + "native normalizer returned empty or oversized text", + )); + } + segments.push(segment); + } + self.punctuation + .replace_all(&segments.join(" "), "$1") + .into_owned() + } else { + text.to_owned() + }; + let encoding = self + .tokenizer + .encode(normalized.as_str(), false) + .map_err(|e| internal(format!("tokenizer failed: {e}")))?; + let mut tokens = encoding.get_ids().to_vec(); + tokens.push(131073); + if tokens.len() > 512 { + return Err(invalid( + "native TTS supports at most 512 text tokens including EOS", + )); + } + if tokens.iter().any(|t| *t >= 131075) { + return Err(invalid("tokenizer returned an out-of-vocabulary token")); + } + serde_json::to_string(&serde_json::json!({"text":normalized,"tokens":tokens})) + .map_err(|e| internal(e.to_string())) + } +} + +unsafe fn cstring<'a>(value: *const c_char) -> Outcome<&'a str> { + if value.is_null() { + return Err(invalid("null string argument")); + } + CStr::from_ptr(value) + .to_str() + .map_err(|_| invalid("argument must be UTF-8")) +} +unsafe fn execute(error: *mut *mut c_char, f: impl FnOnce() -> Outcome<()>) -> i32 { + if error.is_null() { + return 1; + } + *error = ptr::null_mut(); + match catch_unwind(AssertUnwindSafe(f)) { + Ok(Ok(())) => 0, + result => { + let (code, message) = match result { + Ok(Err(e)) => e, + _ => internal("native text preparation panicked"), + }; + *error = CString::new(message.replace('\0', "\\0")) + .unwrap() + .into_raw(); + code + } + } +} + +#[no_mangle] +pub unsafe extern "C" fn nemotron_text_create( + config_json: *const c_char, + tokenizer_path: *const c_char, + handle: *mut *mut c_void, + error: *mut *mut c_char, +) -> i32 { + execute(error, || { + if handle.is_null() { + return Err(invalid("null handle output")); + } + *handle = ptr::null_mut(); + let config: Option = if config_json.is_null() { + None + } else { + serde_json::from_str(cstring(config_json)?) + .map_err(|e| invalid(format!("normalization config: {e}")))? + }; + *handle = Box::into_raw(Box::new(Preparation::new( + config, + cstring(tokenizer_path)?, + )?)) + .cast(); + Ok(()) + }) +} + +#[no_mangle] +pub unsafe extern "C" fn nemotron_text_prepare( + handle: *mut c_void, + text: *const u8, + text_len: usize, + apply_normalize: u8, + output: *mut *mut c_char, + error: *mut *mut c_char, +) -> i32 { + execute(error, || { + if output.is_null() { + return Err(invalid("null output pointer")); + } + *output = ptr::null_mut(); + if handle.is_null() || text.is_null() || text_len > 1_000_000 || apply_normalize > 1 { + return Err(invalid("invalid preparation arguments")); + } + let text = std::str::from_utf8(slice::from_raw_parts(text, text_len)) + .map_err(|_| invalid("text must be UTF-8"))?; + let payload = (&*handle.cast::()).prepare(text, apply_normalize != 0)?; + *output = CString::new(payload) + .map_err(|_| internal("invalid JSON encoding"))? + .into_raw(); + Ok(()) + }) +} + +#[no_mangle] +pub unsafe extern "C" fn nemotron_text_free(value: *mut c_char) { + if !value.is_null() { + drop(CString::from_raw(value)); + } +} + +#[no_mangle] +pub unsafe extern "C" fn nemotron_text_destroy(handle: *mut c_void) { + if !handle.is_null() { + let _ = catch_unwind(AssertUnwindSafe(|| { + drop(Box::from_raw(handle.cast::())) + })); + } +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn trimming_matches_python_controls() { + for value in ['\u{1c}', '\u{1d}', '\u{1e}', '\u{1f}', '\u{85}', '\u{a0}'] { + assert_eq!( + format!("{value}text{value}").trim_matches(whitespace), + "text" + ); + } + assert_eq!( + "\u{200b}text\u{200b}".trim_matches(whitespace), + "\u{200b}text\u{200b}" + ); + } + #[test] + fn invalid_ffi_arguments_report_errors() { + unsafe { + let mut error = ptr::null_mut(); + let mut output = ptr::null_mut(); + assert_eq!( + nemotron_text_prepare( + ptr::null_mut(), + b"x".as_ptr(), + 1, + 1, + &mut output, + &mut error + ), + 1 + ); + assert!(output.is_null()); + assert!(!error.is_null()); + nemotron_text_free(error); + } + } +} diff --git a/native/text_preparation/tests/qualification.py b/native/text_preparation/tests/qualification.py new file mode 100644 index 0000000..0f9b5cc --- /dev/null +++ b/native/text_preparation/tests/qualification.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import argparse +import ctypes as C +import hashlib +import json +import statistics +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from nemotron_cpp_bridge.normalization import TextPreparer +from tokenizers import Tokenizer + + +def main(): + p = argparse.ArgumentParser() + for name in ("library", "model", "config", "corpus", "output"): + p.add_argument("--" + name, type=Path, required=True) + args = p.parse_args() + lib = C.CDLL(str(args.library)) + lib.nemotron_text_create.argtypes = [ + C.c_char_p, + C.c_char_p, + C.POINTER(C.c_void_p), + C.POINTER(C.c_void_p), + ] + lib.nemotron_text_create.restype = C.c_int + lib.nemotron_text_prepare.argtypes = [ + C.c_void_p, + C.c_char_p, + C.c_size_t, + C.c_uint8, + C.POINTER(C.c_void_p), + C.POINTER(C.c_void_p), + ] + lib.nemotron_text_prepare.restype = C.c_int + lib.nemotron_text_free.argtypes = [C.c_void_p] + lib.nemotron_text_free.restype = None + lib.nemotron_text_destroy.argtypes = [C.c_void_p] + lib.nemotron_text_destroy.restype = None + context, error = C.c_void_p(), C.c_void_p() + code = lib.nemotron_text_create( + args.config.read_bytes(), + str(args.model / "tokenizer.json").encode(), + C.byref(context), + C.byref(error), + ) + if code: + message = C.string_at(error).decode() + lib.nemotron_text_free(error) + raise RuntimeError(message) + preparer = TextPreparer(args.config) + tokenizer = Tokenizer.from_file(str(args.model / "tokenizer.json")) + + def native(text, apply=True): + data = text.encode() + out, error = C.c_void_p(), C.c_void_p() + code = lib.nemotron_text_prepare( + context, data, len(data), apply, C.byref(out), C.byref(error) + ) + try: + if code: + message = C.string_at(error).decode() + raise (ValueError if code == 1 else RuntimeError)(message) + value = json.loads(C.string_at(out)) + return value["text"], value["tokens"] + finally: + lib.nemotron_text_free(out) + lib.nemotron_text_free(error) + + def reference(text, apply=True): + text = preparer.prepare(text, "en-US", apply) + ids = tokenizer.encode_batch([text], add_special_tokens=False)[0].ids + [131073] + if len(ids) > 512 or any(x >= 131075 for x in ids): + raise ValueError("invalid token IDs") + return text, ids + + corpus = [line.split("\t", 1)[1] for line in args.corpus.read_text().splitlines()] + edges = [ + " Hello world. ", + "\tHello\n", + "The cost is $12.50.", + "Dr. Smith arrived at 12:30 p.m.", + "Call 415-555-0123.", + "On 01/02/2025, 3.14 percent rose.", + "Version 2.1 and B512.", + "Café naïve résumé.", + "What ? Yes ! Fine , thanks .", + "नमस्ते ।", + "Hello 😀 world.", + "a\0b", + "\u001cHello\u001f", + "\u0085Hello\u00a0", + "\u200bHello\u200b", + "hello", + "one\ntwo\nthree", + "", + " ", + "x " * 600, + "x" * 65537, + ] + records = [] + try: + for text in corpus + edges: + for apply in (False, True): + outcomes = [] + for fn in (reference, native): + try: + outcomes.append(("ok", fn(text, apply))) + except Exception as e: + outcomes.append((type(e).__name__, None)) + if outcomes[0] != outcomes[1]: + raise AssertionError((repr(text[:100]), apply, outcomes)) + records.append( + { + "input_sha256": hashlib.sha256(text.encode()).hexdigest(), + "normalize": apply, + "status": outcomes[0][0], + } + ) + expected = [reference(text) for text in corpus] + timings = [] + with ThreadPoolExecutor(max_workers=16) as pool: + for name, fn in [ + ("python", reference), + ("native", native), + ("native", native), + ("python", reference), + ]: + list(pool.map(fn, [corpus[i % len(corpus)] for i in range(32)])) + started = time.perf_counter() + outputs = list(pool.map(fn, [corpus[i % len(corpus)] for i in range(512)])) + elapsed = (time.perf_counter() - started) * 1000 + assert outputs == [expected[i % len(corpus)] for i in range(512)] + timings.append( + { + "implementation": name, + "requests": 512, + "threads": 16, + "wall_ms": elapsed, + } + ) + result = { + "status": "PASS", + "cases": len(records), + "records": records, + "concurrent_requests": 2048, + "timings": timings, + "library_sha256": hashlib.sha256(args.library.read_bytes()).hexdigest(), + } + args.output.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps({k: v for k, v in result.items() if k != "records"})) + finally: + preparer.close() + lib.nemotron_text_destroy(context) + + +if __name__ == "__main__": + main() diff --git a/native/text_preparation/text_preparation.h b/native/text_preparation/text_preparation.h new file mode 100644 index 0000000..439d77c --- /dev/null +++ b/native/text_preparation/text_preparation.h @@ -0,0 +1,18 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#pragma once +#include +#include +#ifdef __cplusplus +extern "C" { +#endif +int nemotron_text_create( + const char* config_json, const char* tokenizer_path, void** handle, char** error); +int nemotron_text_prepare( + void* handle, const uint8_t* text, size_t text_len, uint8_t apply_normalize, char** output_json, + char** error); +void nemotron_text_free(char* value); +void nemotron_text_destroy(void* handle); +#ifdef __cplusplus +} +#endif diff --git a/native/tts-normalizer/.bazelrc b/native/tts-normalizer/.bazelrc new file mode 100644 index 0000000..7591f5f --- /dev/null +++ b/native/tts-normalizer/.bazelrc @@ -0,0 +1,8 @@ +build --announce_rc +build --color=yes +build --cxxopt=-std=c++17 +build --keep_going +build --show_timestamps +build --strip=always +build:release --compilation_mode=opt --copt=-O3 +test --test_output=errors diff --git a/native/tts-normalizer/BUILD.bazel b/native/tts-normalizer/BUILD.bazel new file mode 100644 index 0000000..31749e2 --- /dev/null +++ b/native/tts-normalizer/BUILD.bazel @@ -0,0 +1,97 @@ +package(default_visibility = ["//visibility:public"]) + +cc_test( + name = "string_matcher_test", + srcs = ["tests/string_matcher_test.cc", "//third_party:sparrowhawk_string_matcher.h"], + includes = ["third_party"], + deps = ["@openfst//:fst"], +) + +cc_library( + name = "processor_abi", + hdrs = ["include/processor_abi.h"], + includes = ["include"], +) + +cc_library( + name = "text_preprocessor", + srcs = [ + "src/normalize.cc", + "src/speech_hints.cc", + "src/text_preprocessor.cc", + ], + hdrs = [ + "src/normalize.h", + "src/speech_hints.h", + "src/text_preprocessor.h", + ], + copts = [ + "-O3", + "-fPIC", + "-Wno-deprecated-declarations", + ], + includes = ["src"], + linkstatic = True, + deps = [ + "@com_google_absl//absl/strings", + "@openfst//:fst", + "@pugixml//:pugixml", + "@sparrowhawk//:sparrowhawk", + "@thrax_1_2_7//:thrax_1_2_7", + ], +) + +cc_binary( + name = "text_preprocessor_c", + srcs = ["src/text_preprocessor_c.cc"], + linkshared = True, + linkstatic = True, + deps = [ + ":processor_abi", + ":text_preprocessor", + ], +) + +cc_test( + name = "text_preprocessor_test", + srcs = ["tests/text_preprocessor_test.cc"], + data = ["//testdata:english_assets"], + deps = [ + ":text_preprocessor", + "@googletest//:gtest_main", + ], +) + +cc_test( + name = "normalize_test", + srcs = ["tests/normalize_test.cc"], + data = ["//testdata:english_assets"], + deps = [ + ":text_preprocessor", + "@googletest//:gtest_main", + "@openfst//:far_base", + ], +) + +cc_test( + name = "speech_hints_test", + srcs = ["tests/speech_hints_test.cc"], + deps = [ + ":text_preprocessor", + "@googletest//:gtest_main", + ], +) + +cc_test( + name = "text_preprocessor_c_test", + srcs = [ + "src/text_preprocessor_c.cc", + "tests/text_preprocessor_c_test.cc", + ], + data = ["//testdata:english_assets"], + deps = [ + ":processor_abi", + ":text_preprocessor", + "@googletest//:gtest_main", + ], +) diff --git a/native/tts-normalizer/LICENSE.legacy b/native/tts-normalizer/LICENSE.legacy new file mode 100644 index 0000000..009736b --- /dev/null +++ b/native/tts-normalizer/LICENSE.legacy @@ -0,0 +1,7 @@ +Copyright (c) 2020-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. diff --git a/native/tts-normalizer/LICENSE.riva b/native/tts-normalizer/LICENSE.riva new file mode 100644 index 0000000..009736b --- /dev/null +++ b/native/tts-normalizer/LICENSE.riva @@ -0,0 +1,7 @@ +Copyright (c) 2020-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. diff --git a/native/tts-normalizer/README.md b/native/tts-normalizer/README.md new file mode 100644 index 0000000..c846216 --- /dev/null +++ b/native/tts-normalizer/README.md @@ -0,0 +1,49 @@ +# Native TTS normalizer + +This package is the source-owned build boundary for the CPU-only WFST text +normalizer used by the Rust EasyMagpie adapter. It is intentionally independent +of a NemotronSpeech checkout and exposes only processor ABI v1 from +`include/processor_abi.h`. + +The final deployment still contains a shared library because the stable native +processor ABI is the isolation boundary between the Rust adapter and the C++ +normalizer. The library is now built and tested from this directory rather than +stored as a Git LFS binary. + +## Build and test + +The supported build runs in `docker/Dockerfile.frontend`. The checked-in English +test grammars under `testdata/en-US` make a matching local Bazel 6.2.1 build +hermetic: + +```bash +bazel test --config=release \ + //:normalize_test \ + //:speech_hints_test \ + //:text_preprocessor_test \ + //:text_preprocessor_c_test +bazel build --config=release //:text_preprocessor_c +``` + +The shared library is produced as +`bazel-bin/libtext_preprocessor_c.so`. The container installs it as +`/opt/nemotron-speech/lib/processors/libnemotron_speech_tts_text_preprocessor.so`. + +## Ownership boundary + +This directory owns normalization, sentence splitting, bounded SSML handling, +abbreviation expansion, and the processor ABI wrapper. It does not own model +tokenization, session state, streaming boundary detection, language pool sizing, +or vLLM transport; those remain in the Rust adapter and frontend. + +The English FAR/proto/abbreviation files under `testdata/en-US` are test fixtures +only. They are consumed by the native Bazel tests and are not copied into the +runtime image. Production deployments use the external multilingual asset root +configured in `rust/config/example.json` (or an explicit deploy-time override); +the source build does not generate or package production grammars. + +All external source archives are revision- and SHA-256-pinned. An uncached build +therefore needs network access, but it does not depend on mutable Git branches or +an adjacent source checkout. + +See `LICENSE.nemotron_speech` and `THIRD_PARTY_NOTICES.md` for licensing information. diff --git a/native/tts-normalizer/THIRD_PARTY_NOTICES.md b/native/tts-normalizer/THIRD_PARTY_NOTICES.md new file mode 100644 index 0000000..f9d2877 --- /dev/null +++ b/native/tts-normalizer/THIRD_PARTY_NOTICES.md @@ -0,0 +1,24 @@ +# Third-party dependencies + +The native normalizer statically links the dependencies below. Their source is +fetched at the immutable revision or archive checksum declared in +`WORKSPACE.bazel`; no third-party source archive is committed here. + +| Component | Pinned revision/version | License | +|---|---|---| +| Abseil | 20220623.1 | Apache-2.0 | +| gflags | 2.2.2 | BSD-3-Clause | +| glog | 0.4.0 | BSD-3-Clause | +| gRPC build dependencies | 1.50.1 | Apache-2.0 | +| GoogleTest | 1.11.0 | BSD-3-Clause | +| OpenFst | `fc23b4cf529429284b874a26f28b15c6cc94f404` | Apache-2.0 | +| protobuf | version selected by gRPC 1.50.1 | BSD-3-Clause | +| pugixml | 1.11.4 | MIT | +| RE2 | `ccdbd348259434cca6c66c66a35c5ef71ff3ee1a` | BSD-3-Clause | +| Sparrowhawk | `d9d68e7155b2b2cd64f578144c32b4e09674f43d` | Apache-2.0 | +| Thrax | `21407b7ccbd5dbaa45f34b26f286bc28d57c47fb` | Apache-2.0 | + +Boost and zlib headers/libraries are supplied by the Ubuntu 22.04 builder under +their respective Boost Software License 1.0 and zlib license. The original +license files remain authoritative and are present in the fetched source or +builder packages. diff --git a/native/tts-normalizer/WORKSPACE.bazel b/native/tts-normalizer/WORKSPACE.bazel new file mode 100644 index 0000000..213e4d9 --- /dev/null +++ b/native/tts-normalizer/WORKSPACE.bazel @@ -0,0 +1,89 @@ +workspace(name = "nemotron_tts_normalizer") + +load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") + +# Versions match the extracted NemotronSpeech normalizer dependency set. +http_archive( + name = "googletest", + sha256 = "b4870bf121ff7795ba20d20bcdd8627b8e088f2d1dab299a031c1034eddc93d5", + strip_prefix = "googletest-release-1.11.0", + urls = ["https://github.com/google/googletest/archive/refs/tags/release-1.11.0.tar.gz"], +) + +http_archive( + name = "com_github_gflags_gflags", + sha256 = "34af2f15cf7367513b352bdcd2493ab14ce43692d2dcd9dfc499492966c64dcf", + strip_prefix = "gflags-2.2.2", + urls = ["https://github.com/gflags/gflags/archive/v2.2.2.tar.gz"], +) + +http_archive( + name = "glog", + sha256 = "f28359aeba12f30d73d9e4711ef356dc842886968112162bc73002645139c39c", + strip_prefix = "glog-0.4.0", + urls = ["https://github.com/google/glog/archive/v0.4.0.tar.gz"], +) + +http_archive( + name = "com_google_absl", + sha256 = "91ac87d30cc6d79f9ab974c51874a704de9c2647c40f6932597329a282217ba8", + strip_prefix = "abseil-cpp-20220623.1", + urls = ["https://github.com/abseil/abseil-cpp/archive/refs/tags/20220623.1.tar.gz"], +) + +http_archive( + name = "com_googlesource_code_re2", + build_file = "//third_party:BUILD.re2", + sha256 = "8675f33ef3dcac24692821edf21a56ef0112c95c4ec50b472e68efe0ff5dda1e", + strip_prefix = "re2-ccdbd348259434cca6c66c66a35c5ef71ff3ee1a", + urls = ["https://github.com/google/re2/archive/ccdbd348259434cca6c66c66a35c5ef71ff3ee1a.zip"], +) + +# gRPC's dependency macros provide the protobuf and Bazel rule versions used by +# the pinned Sparrowhawk BUILD definition. +http_archive( + name = "com_github_grpc_grpc", + sha256 = "fb1ed98eb3555877d55eb2b948caca44bc8601c6704896594de81558639709ef", + strip_prefix = "grpc-1.50.1", + urls = ["https://github.com/grpc/grpc/archive/refs/tags/v1.50.1.tar.gz"], +) + +load("@com_github_grpc_grpc//bazel:grpc_deps.bzl", "grpc_deps") + +grpc_deps() + +load("@com_github_grpc_grpc//bazel:grpc_extra_deps.bzl", "grpc_extra_deps") + +grpc_extra_deps() + +http_archive( + name = "openfst", + build_file = "//third_party:BUILD.openfst", + sha256 = "9d1afcfe72ccded0e846a5883eed9dfcd7593e6cd935ad55a94fec69270d7540", + strip_prefix = "openfst-fc23b4cf529429284b874a26f28b15c6cc94f404", + urls = ["https://github.com/sarane22/openfst/archive/fc23b4cf529429284b874a26f28b15c6cc94f404.tar.gz"], +) + +http_archive( + name = "thrax_1_2_7", + build_file = "//third_party:BUILD.thrax_1_2_7", + sha256 = "be2307b4aee07a8596e74743ef9f7a51459242b1a82bf84648782966d0ea6ec9", + strip_prefix = "thrax-21407b7ccbd5dbaa45f34b26f286bc28d57c47fb", + urls = ["https://github.com/sarane22/thrax/archive/21407b7ccbd5dbaa45f34b26f286bc28d57c47fb.tar.gz"], +) + +http_archive( + name = "sparrowhawk", + build_file = "//third_party:BUILD.sparrowhawk_0_1", + sha256 = "5168968b1a9f277b52f189e9aca1aa84f13082333d1ed3e96e9e653de7a5c3f4", + strip_prefix = "sparrowhawk-d9d68e7155b2b2cd64f578144c32b4e09674f43d", + urls = ["https://github.com/anand-nv/sparrowhawk/archive/d9d68e7155b2b2cd64f578144c32b4e09674f43d.tar.gz"], +) + +http_archive( + name = "pugixml", + build_file = "//third_party:BUILD.pugixml", + sha256 = "017139251c122dbff400a507cddc4cb74120a431a50c6c524f30edcc5b331ade", + strip_prefix = "pugixml-1.11.4", + urls = ["https://github.com/zeux/pugixml/archive/refs/tags/v1.11.4.tar.gz"], +) diff --git a/native/tts-normalizer/include/processor_abi.h b/native/tts-normalizer/include/processor_abi.h new file mode 100644 index 0000000..3f3dd63 --- /dev/null +++ b/native/tts-normalizer/include/processor_abi.h @@ -0,0 +1,129 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + */ + +#pragma once + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +#define NEMOTRON_SPEECH_PROCESSOR_ABI_VERSION 1u + +enum NemotronSpeechRecordKindV1 { + NEMOTRON_SPEECH_RECORD_UTF8_V1 = 1, + NEMOTRON_SPEECH_RECORD_BYTES_V1 = 2, + NEMOTRON_SPEECH_RECORD_AUDIO_V1 = 3, + NEMOTRON_SPEECH_RECORD_TOKEN_IDS_V1 = 4, + NEMOTRON_SPEECH_RECORD_TENSOR_V1 = 5, + NEMOTRON_SPEECH_RECORD_INTEGER_V1 = 6, + NEMOTRON_SPEECH_RECORD_FLOAT_V1 = 7, + NEMOTRON_SPEECH_RECORD_BOOLEAN_V1 = 8, +}; + +enum NemotronSpeechElementTypeV1 { + NEMOTRON_SPEECH_ELEMENT_NONE_V1 = 0, + NEMOTRON_SPEECH_ELEMENT_U8_V1 = 1, + NEMOTRON_SPEECH_ELEMENT_I32_V1 = 2, + NEMOTRON_SPEECH_ELEMENT_I64_V1 = 3, + NEMOTRON_SPEECH_ELEMENT_F32_V1 = 4, +}; + +enum NemotronSpeechAudioEncodingV1 { + NEMOTRON_SPEECH_AUDIO_NONE_V1 = 0, + NEMOTRON_SPEECH_AUDIO_PCM_S16_LE_V1 = 1, + NEMOTRON_SPEECH_AUDIO_PCM_F32_LE_V1 = 2, +}; + +typedef struct NemotronSpeechByteSpanV1 { + const uint8_t* data; + size_t len; +} NemotronSpeechByteSpanV1; + +typedef struct NemotronSpeechStatusV1 { + size_t struct_size; + int32_t code; + NemotronSpeechByteSpanV1 message; + void* context; +} NemotronSpeechStatusV1; + +typedef struct NemotronSpeechRecordViewV1 { + size_t struct_size; + NemotronSpeechByteSpanV1 name; + uint32_t kind; + uint32_t element_type; + NemotronSpeechByteSpanV1 data; + const size_t* shape; + size_t rank; + uint32_t sample_rate_hz; + uint16_t channels; + uint16_t audio_encoding; + uint64_t sequence; + uint8_t has_sequence; + uint8_t is_final; + uint8_t reserved[6]; +} NemotronSpeechRecordViewV1; + +typedef struct NemotronSpeechEnvelopeViewV1 { + size_t struct_size; + uint32_t schema_version; + NemotronSpeechByteSpanV1 request_id; + NemotronSpeechByteSpanV1 session_id; + const NemotronSpeechRecordViewV1* records; + size_t record_count; + uint64_t sequence; + uint8_t has_sequence; + uint8_t is_final; + uint8_t reserved[6]; +} NemotronSpeechEnvelopeViewV1; + +typedef struct NemotronSpeechEnvelopeOutputV1 { + size_t struct_size; + uint32_t schema_version; + NemotronSpeechByteSpanV1 request_id; + NemotronSpeechByteSpanV1 session_id; + const NemotronSpeechRecordViewV1* records; + size_t record_count; + uint64_t sequence; + uint8_t has_sequence; + uint8_t is_final; + uint8_t reserved[6]; + void* context; +} NemotronSpeechEnvelopeOutputV1; + +typedef struct NemotronSpeechProcessorCapabilitiesV1 { + size_t struct_size; + NemotronSpeechByteSpanV1 processor_id; + uint8_t batch_input; + uint8_t streaming_input; + uint8_t thread_safe; + uint8_t stateful; + uint8_t reserved[4]; +} NemotronSpeechProcessorCapabilitiesV1; + +typedef void* NemotronSpeechProcessorHandleV1; + +typedef struct NemotronSpeechProcessorApiV1 { + size_t struct_size; + uint32_t abi_version; + NemotronSpeechStatusV1 (*capabilities)(NemotronSpeechProcessorCapabilitiesV1* output); + NemotronSpeechStatusV1 (*create)(const NemotronSpeechEnvelopeViewV1* config, + NemotronSpeechProcessorHandleV1* output); + void (*destroy)(NemotronSpeechProcessorHandleV1 handle); + NemotronSpeechStatusV1 (*process)( + NemotronSpeechProcessorHandleV1 handle, const NemotronSpeechEnvelopeViewV1* request, + NemotronSpeechEnvelopeOutputV1* output); + void (*release_output)(NemotronSpeechProcessorHandleV1 handle, + NemotronSpeechEnvelopeOutputV1* output); + void (*release_status)(NemotronSpeechStatusV1* status); +} NemotronSpeechProcessorApiV1; + +const NemotronSpeechProcessorApiV1* nemotron_speech_processor_get_api( + uint32_t requested_version); + +#ifdef __cplusplus +} +#endif diff --git a/native/tts-normalizer/src/normalize.cc b/native/tts-normalizer/src/normalize.cc new file mode 100644 index 0000000..61d76f2 --- /dev/null +++ b/native/tts-normalizer/src/normalize.cc @@ -0,0 +1,221 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. + * + * NVIDIA CORPORATION and its licensors retain all intellectual property + * and proprietary rights in and to this software, related documentation + * and any modifications thereto. Any use, reproduction, disclosure or + * distribution of this software and related documentation without an express + * license agreement from NVIDIA CORPORATION is strictly prohibited. + */ + +#include "normalize.h" + +#include "absl/strings/str_replace.h" + +namespace nemotron_speech::utils { + +Normalizer::Normalizer(std::string config, std::string prefix, bool speech_hints) + : config_(config), prefix_(prefix) +{ + normalizer_ = std::make_unique(); + if (speech_hints) { + speechhints_ = std::make_unique(); + check_speech_hints_ = true; + } + if (!LoadFullGrammar()) { + throw std::runtime_error("Error loading WFST grammar"); + } +} + +Normalizer::Normalizer(const std::string& prefix, bool speech_hints) : prefix_(prefix) +{ + normalizer_ = std::make_unique(); + if (speech_hints) { + speechhints_ = std::make_unique(); + check_speech_hints_ = true; + } + if (!LoadGrammar()) { + throw std::runtime_error("Error loading WFST grammar"); + } +} + +bool +Normalizer::LoadGrammar() +{ + struct stat filestat; + if (-1 == stat(prefix_.c_str(), &filestat)) { + LOG(ERROR) << "Grammar directory missing:" << prefix_ << std::endl; + return false; + } + if (-1 == stat((prefix_ + "/speech_class.far").c_str(), &filestat)) { + check_speech_hints_ = false; + } + if (-1 == stat((prefix_ + "/post_process.far").c_str(), &filestat)) { + enable_postprocess_ = false; + } + if (-1 == stat((prefix_ + "/pre_process.far").c_str(), &filestat)) { + enable_preprocess_ = false; + } + if (check_speech_hints_) { + auto load_check = speechhints_->LoadGrammars((prefix_ + "/speech_class.far")); + if (!load_check) { + LOG(ERROR) << "Failed to load the grammar"; + return load_check; + } + } + return normalizer_->Setup(prefix_, enable_postprocess_, enable_preprocess_); +} + +bool +Normalizer::LoadFullGrammar() +{ + struct stat filestat; + if (-1 == stat(prefix_.c_str(), &filestat) || + -1 == stat((prefix_ + "/" + config_).c_str(), &filestat)) { + LOG(ERROR) << "Grammar files missing:" << prefix_ + "/" + config_ << std::endl; + return false; + } + if (-1 == stat((prefix_ + "/speech_class.far").c_str(), &filestat)) { + check_speech_hints_ = false; + } + if (check_speech_hints_) { + auto load_check = speechhints_->LoadGrammars((prefix_ + "/speech_class.far")); + if (!load_check) { + LOG(ERROR) << "Failed to load the grammar"; + return load_check; + } + } + + return normalizer_->Setup(config_, prefix_); +} + +// intentionally no sentence splitting as its pre punctuation +// remove any trailing whitepsace then reattach after normalizing. +bool +Normalizer::NormalizeAndShowAlignments( + const std::string& input, std::string& output, bool enable_preprocessing) const +{ + int terminal_position = input.find_last_not_of(" ") + 1; + std::string fixed_input = input.substr(0, terminal_position); + std::string ws_substring = input.substr(terminal_position, input.length() - 1); + fixed_input = absl::StrReplaceAll(fixed_input, {{"\"", "\\\" "}}); + auto rval = normalizer_->NormalizeAndShowLinks(fixed_input, &output, enable_preprocessing); + if (!rval) { + LOG(ERROR) << "normalizer failed for '" << fixed_input << "'. Returning output unchanged."; + output = input; + return rval; + } + + if (fixed_input != input) { + output += ws_substring; + } + + // it is possible for the normalizer to inject unicode non-breaking spaces + // we replace these with standard spaces if encountered + output = absl::StrReplaceAll(output, {{"\xC2\xA0", " "}}); + + // Strip lone trailing '<' characters left by SPARROWHAWK verbalizer grammars + // (e.g. Korean verbalize.far leaves an unclosed token-class marker at end of sentence). + // Only strip if the '<' has no matching '>'. + while (!output.empty() && output.back() == '<') { + output.pop_back(); + } + + return rval; +} + +bool +Normalizer::IsHintValid(const std::string& speechhint) const +{ + if (check_speech_hints_) { + return speechhints_->IsHintValid(speechhint); + } else { + LOG(ERROR) << "Speechhints not initiated!"; + return check_speech_hints_; + } +} + +bool +Normalizer::AreHintsValid(const std::vector& speechhints) const +{ + if (check_speech_hints_) { + return speechhints_->AreHintsValid(speechhints); + } else { + LOG(ERROR) << "Speechhints not initiated!"; + return check_speech_hints_; + } +} + +bool +Normalizer::Normalize( + const std::string& input, std::string& output, const std::vector classHints, + const bool verbatim_transcripts, const bool enable_preprocessing) const +{ + if (input.empty()) { + output.clear(); + return true; + } + static constexpr char kBoundaryWs[] = " \t\n\r\f\v"; + const auto first_non_space = input.find_first_not_of(kBoundaryWs); + if (first_non_space == std::string::npos) { + output = input; + return true; + } + const auto terminal_position = input.find_last_not_of(kBoundaryWs) + 1; + const std::string leading_ws = input.substr(0, first_non_space); + const std::string trailing_ws = input.substr(terminal_position); + std::string fixed_input = input.substr(first_non_space, terminal_position - first_non_space); + fixed_input = absl::StrReplaceAll(fixed_input, {{"\"", "\\\" "}}); + bool result = false; + if (check_speech_hints_) { + Alignments alignments; + if (!classHints.empty()) { + result = speechhints_->ApplyHints(fixed_input, classHints, output, alignments); + if (!result) { + LOG(ERROR) << "normalizer failed to apply speechhints"; + } + if (result && alignments.empty()) { + LOG(ERROR) << "normalizer failed to apply speechhints, alignments empty"; + result = false; + } + } + if (!result) { + output = input; + } else { + NormalizedSegments segments; + speechhints_->GetSegments(alignments, segments); + if (segments.empty()) { + output = leading_ws + fixed_input + trailing_ws; + return false; + } + std::string final_output; + for (auto& segment : segments) { + if (!verbatim_transcripts && !segment.second && + !normalizer_->Normalize(segment.first, &segment.first, enable_preprocessing)) { + LOG(ERROR) << "normalizer failed for segment, returning value unchanged"; + } + final_output += (segment.first + " "); + } + output = final_output; + boost::algorithm::trim(output); + } + } + if (!check_speech_hints_ || !result) { + auto rval = normalizer_->Normalize(fixed_input, &output, enable_preprocessing); + if (!rval) { + LOG(ERROR) << "normalizer failed, returning value unchanged"; + output = input; + return rval; + } + } + // A preprocessor may validly delete the entire input (for example, a final ASR segment that + // consists only of a filler word). Do not turn that result back into boundary whitespace. + if (output.empty()) { + return true; + } + output = leading_ws + output + trailing_ws; + output = absl::StrReplaceAll(output, {{"\xC2\xA0", " "}}); + return true; +} + +} // namespace nemotron_speech::utils diff --git a/native/tts-normalizer/src/normalize.h b/native/tts-normalizer/src/normalize.h new file mode 100644 index 0000000..c1772eb --- /dev/null +++ b/native/tts-normalizer/src/normalize.h @@ -0,0 +1,147 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. + * + * NVIDIA CORPORATION and its licensors retain all intellectual property + * and proprietary rights in and to this software, related documentation + * and any modifications thereto. Any use, reproduction, disclosure or + * distribution of this software and related documentation without an express + * license agreement from NVIDIA CORPORATION is strictly prohibited. + */ + +#pragma once + +#include + +#include "speech_hints.h" + +namespace nemotron_speech::utils { + +/* + * This Normalizer class exists as an abstraction over however we implement text normalization. + * + */ + +class Normalizer { + private: + std::unique_ptr normalizer_; + std::unique_ptr speechhints_; + std::string config_; + std::string prefix_; + + public: + bool check_speech_hints_ = false; + bool enable_preprocess_ = true; + bool enable_postprocess_ = true; + Normalizer() = delete; + + /** + * @brief original path into sparrowhawk normalization - requires a sparrowhawk config (config) + * and the six files in prefix_ + * + * @detail sparrowhawk expected files in the config_ directory - these have not changed in the + * first year plus of deployment. sentence_boundary_exceptions.txt + * sparrowhawk_configuration.ascii_proto tokenize_and_classify.far tokenizer.ascii_proto + * verbalize.far verbalizer.ascii_proto + * + * @param config + * @param prefix + * @param speechhints + */ + Normalizer(std::string config, std::string prefix, bool speech_hint = false); + + Normalizer(const std::string& prefix, bool speech_hint = false); + + /** + * @brief normalizes the input provided and applies speech hints. + * + * @detail Applies speechhints and then normalizing the remaining segments. The output is a + * normalized string with speechhints applied: + * Input: i work at the b b c studios + * Speechints: {"$OOV_ALPHA_SEQUENCE"} + * Output: i work at the bbc studios + * + * @param input input string to normalize + * @param output normalized string + * @param classHints speech hints + * @param verbatim_transcripts run only speech hints if verbatim_transcripts is true + * @param enable_preprocessing apply the configured pre_process.far before ITN + */ + bool Normalize( + const std::string& input, std::string& output, const std::vector classHints = {}, + const bool verbatim_transcripts = false, const bool enable_preprocessing = false) const; + + /** + * @brief normalizes the input provided and adds the alignments to the location in the original + * string. + * + * @detail Normalize input based on grammars specified and adds alignments. The output is an + * overloaded string: + * Token: 0 what 0,3 0,0 + * Token: 1 are 5,7 1,1 + * Token: 2 the 9,11 2,2 + * Token: 3 alignments 13,22 3,3 + * Token: 4 here 24,27 4,4 + * Token: 5 four three two one 29,46 5,5 + * Word: 0 what 0 + * Word: 1 are 1 + * Word: 2 the 2 + * Word: 3 alignments 3 + * Word: 4 here 4 + * Word: 5 4321 5 + * + * If we need alignments we should return the Utterance object this is dervied from or something + * containing the words/token pieces vs a string + * + * @param input + * @param output + */ + bool NormalizeAndShowAlignments( + const std::string& input, std::string& output, bool enable_preprocessing = false) const; + + /** + * @brief Uses as many defaults as possible to load grammars from a directory + * + * @detail the only two files required for this load path are tokenize_and_classify.far and + * verbalize.far, the others are still configurable, but are setup to use defaults. see the + * sparrowhawk repo for more details. + */ + bool LoadGrammar(); + /** + * @brief requires all files to be present in the directory stored in prefix_. + * + * @detail sparrowhawk by default requires six files to stand up normalization, this function + * expect all files to be present and properly setup in the directory prefix_. + * + */ + + bool LoadFullGrammar(); + + /** + * @brief Checks if the speechhints vector has valid hints. + * + * @detail This is an wrapper around the AreHintsValid speech_hints.cc function to check if the + * input speechhints vector has all the phrases as valid speechhints. + * + * @param speechhints vector of speechhints strings + */ + bool AreHintsValid(const std::vector& speechhints) const; + /** + * @brief Checks if the speechhint is a valid hint. + * + * @detail This is an wrapper around the IsHintValid speech_hints.cc function to check if the + * input speechhint phrase is valid. + * + * @param speechhint string + */ + bool IsHintValid(const std::string& speechhint) const; + /** + * @brief splits the input string based on the configuration of the sentence splitting regex as + * passed in when constructing the normalization class. + */ + std::vector SentenceSplitter(const std::string& input) const + { + return normalizer_->SentenceSplitter(input); + } +}; + +} // namespace nemotron_speech::utils diff --git a/native/tts-normalizer/src/speech_hints.cc b/native/tts-normalizer/src/speech_hints.cc new file mode 100644 index 0000000..905b583 --- /dev/null +++ b/native/tts-normalizer/src/speech_hints.cc @@ -0,0 +1,359 @@ +/* + * Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. + * + * NVIDIA CORPORATION and its licensors retain all intellectual property + * and proprietary rights in and to this software, related documentation + * and any modifications thereto. Any use, reproduction, disclosure or + * distribution of this software and related documentation without an express + * license agreement from NVIDIA CORPORATION is strictly prohibited. + */ + +#include "speech_hints.h" +namespace nemotron_speech::utils { + +SpeechHints::SpeechHints() +{ + word_sep = boost::char_separator(" \t,.!?;./\"(){}[]<>"); +} + +bool +SpeechHints::StringFstToOutputLabels( + const MutableTransducer& ifst, std::vector* labels) +{ + labels->clear(); + auto s = ifst.Start(); + + if (s == fst::kNoStateId) { + LOG(ERROR) << "StringFstToOutputLabels: Invalid start state"; + return false; + } + + while (ifst.Final(s) == fst::StdArc::Weight::Zero()) { + fst::ArcIterator aiter(ifst, s); + if (aiter.Done()) { + LOG(ERROR) << "StringFstToOutputLabels: Does not reach final state"; + return false; + } + + const auto& arc = aiter.Value(); + labels->push_back(arc.olabel); + s = arc.nextstate; + aiter.Next(); + + if (!aiter.Done()) { + LOG(ERROR) << "StringFstToOutputLabels: State " << s << " has multiple outgoing arcs"; + return false; + } + } + + if (ifst.NumArcs(s) != 0) { + LOG(ERROR) << "StringFstToOutputLabels: Final state " << s << " has outgoing arc(s)"; + return false; + } + return true; +} + + +bool +SpeechHints::LoadGrammars(const std::string& far_filename) +{ + sh_class_.reset(new GrmManager); + if (!sh_class_->LoadArchive(far_filename)) { + LOG(ERROR) << "Unable to load grammars from " << far_filename << std::endl; + return false; + } else { + for (auto& i : *sh_class_->GetFstMap()) { + if ((i.first != space_grammar) && (i.first != passthrough_grammar)) + valid_grammars.push_back(i.first); + } + if (!GetClassFst(passthrough_grammar, sent_fst_)) { + LOG(ERROR) << "Missing passthrough grammar" << far_filename << std::endl; + return false; + } + if (!GetClassFst(space_grammar, space_fst_)) { + LOG(ERROR) << "Missing whitespace grammar" << far_filename << std::endl; + return false; + } else { + space_sigma_ = *space_fst_.Copy(true); + fst::Closure(&space_sigma_, fst::CLOSURE_STAR); + } + } + grammar_loaded = true; + return true; +} + +std::vector +SpeechHints::GetGrammarList() +{ + if (grammar_loaded) { + return valid_grammars; + } + return {}; +} + +MutableTransducer +SpeechHints::Word2Fst(const std::string& word) +{ + Compiler str2fst(fst::TokenType::BYTE); + MutableTransducer word_fst; + + if ((word == passthrough_grammar) || (word == space_grammar)) { + LOG(ERROR) << "Cannot specify one of '" << passthrough_grammar << "' or '" << space_grammar + << " in speech hints"; + } else if ((!GetClassFst(word, word_fst)) && (!str2fst.operator()(word, &word_fst))) { + // Get class fst if word is class else convert word to fst + LOG(ERROR) << "Unable to create fst for word ' " << word << "'" << std::endl; + } + return word_fst; +} + +bool +SpeechHints::GetClassFst(const std::string& sh_class_name, MutableTransducer& word_fst) +{ + if ((sh_class_->GetFstMap()->count(sh_class_name) == 0)) { + if (sh_class_name[0] == '$') + LOG(ERROR) << "Invalid speech class - " << sh_class_name; + return false; + } + auto it = sh_class_->GetFstMap()->find(sh_class_name); + + /*Copy the speech hint fst for the specified class. Since this will be composed with other fsts, + needs to be a copy. Refer https://www.openfst.org/twiki/bin/view/FST/FstConventions */ + word_fst = *it->second->Copy(true); + + return true; +} + +MutableTransducer +SpeechHints::Hint2Fst(const std::string& hint) +{ + MutableTransducer hint_fst; + if ((hint != passthrough_grammar) && (hint != space_grammar)) { + if (!GetClassFst(hint, hint_fst)) { + tokenizer word_tokenizer(hint, word_sep); + tokenizer::iterator word = word_tokenizer.begin(); + hint_fst = Word2Fst(word.current_token()); + std::advance(word, 1); + for (; word != word_tokenizer.end(); ++word) { + MutableTransducer word_fst = Word2Fst(word.current_token()); + fst::Concat(&hint_fst, space_sigma_); + fst::Concat(&hint_fst, word_fst); + } + } + } else { + LOG(ERROR) << "Cannot specify one of '" << passthrough_grammar << "' or '" << space_grammar + << " in speech hints"; + } + + + return hint_fst; +} + +bool +SpeechHints::StringifyFst(MutableTransducer& ofst, std::string& output, Alignments& alignments) +{ + output.clear(); + alignments.clear(); + MutableTransducer temp; + fst::ShortestPath(ofst, &temp); + if (temp.Start() == ::fst::kNoStateId) + return false; + MutableTransducer ifst(*temp.Copy()); + fst::Invert(&ifst); + std::vector ilabels; + std::vector olabels; + std::vector::iterator it, ot; + StringFstToOutputLabels(ifst, &ilabels); + StringFstToOutputLabels(temp, &olabels); + std::string iword; + std::string oword; + + for (it = ilabels.begin(), ot = olabels.begin(); it != ilabels.end() && ot != olabels.end(); + ++it, ++ot) { + if ((*it == ' ') && (*ot == ' ')) { + boost::algorithm::trim(iword); + boost::algorithm::trim(oword); + if (iword != "" || oword != "") + alignments.push_back(std::make_pair(iword, oword)); + output += oword + ' '; + iword = ""; + oword = ""; + + + } else { + if (*it != 0) + iword.push_back(*it); + if (*ot != 0) + oword.push_back(*ot); + } + } + boost::algorithm::trim(output); + return true; +} + +bool +SpeechHints::AreHintsValid(const std::vector& speech_hints) +{ + if (!grammar_loaded) + return false; + + for (auto hint : speech_hints) { + if (!IsHintValid(hint)) { + return false; + } + } + return true; +} + +bool +SpeechHints::IsHintValid(const std::string& speech_hint) +{ + if (!grammar_loaded) + return false; + tokenizer word_tokenizer(speech_hint, word_sep); + for (auto word : word_tokenizer) { + if ((word == passthrough_grammar) || (word == space_grammar)) + return false; + if (word[0] == '$') { + if (std::find(valid_grammars.begin(), valid_grammars.end(), word) == valid_grammars.end()) { + return false; + } + } + } + return true; +} + +bool +SpeechHints::ApplyHint( + const std::string& input, const std::string& speech_hint, std::string& output, + Alignments& alignments) +{ + if (!grammar_loaded) { + return false; + } + + MutableTransducer sh_fst; + fst::Union(&sh_fst, sent_fst_); + fst::Union(&sh_fst, space_sigma_); + std::string speech_hint_ = speech_hint; + boost::algorithm::trim(speech_hint_); + if (!speech_hint_.empty()) + fst::Union(&sh_fst, Hint2Fst(speech_hint_)); + + fst::Closure(&sh_fst, fst::CLOSURE_STAR); + + return ApplyHintFst(input, sh_fst, output, alignments); +} + +bool +SpeechHints::ApplyHints( + const std::string& input, const std::vector& speech_hints, std::string& output, + Alignments& alignments) +{ + if (!grammar_loaded) + return false; + MutableTransducer hint_fst, sh_fst = sent_fst_; + fst::Union(&sh_fst, space_fst_); + + if (!speech_hints.empty()) { + std::vector::const_iterator speech_hint = speech_hints.begin(); + std::string speech_hint_ = *speech_hint; + boost::algorithm::trim(speech_hint_); + if (!speech_hint_.empty()) + fst::Union(&sh_fst, Hint2Fst(speech_hint_)); + + std::advance(speech_hint, 1); + for (; speech_hint != speech_hints.end(); ++speech_hint) { + speech_hint_ = *speech_hint; + boost::algorithm::trim(speech_hint_); + + if (!speech_hint_.empty()) + fst::Union(&sh_fst, Hint2Fst(speech_hint_)); + } + } + + fst::Closure(&sh_fst, fst::CLOSURE_STAR); + + return ApplyHintFst(input, sh_fst, output, alignments); +} + +bool +SpeechHints::ApplyHintFst( + const std::string& input, MutableTransducer speech_hint, std::string& output, + Alignments& alignments) +{ + MutableTransducer input_fst, output_fst; + std::string input_pad = " " + input + " "; + Compiler str2fst(fst::TokenType::BYTE); + if (!str2fst.operator()(input_pad, &input_fst)) { + LOG(INFO) << "Unable to create fst for input ' " << input << "'"; + } + fst::ComposeOptions opts(true, fst::ALT_SEQUENCE_FILTER); + + fst::Compose(input_fst, speech_hint, &output_fst, opts); + if (!StringifyFst(output_fst, output, alignments)) { + output = input; + + return false; + } + + return true; +} + +void +SpeechHints::GetSegments(Alignments alignments, NormalizedSegments& segments) +{ + if (!grammar_loaded) + return; + std::string segment; + + for (auto it : alignments) { + if (it.first == it.second) { + segment = segment + it.first + " "; + } else { + boost::algorithm::trim(segment); + if (!segment.empty()) { + segments.push_back(std::make_pair(segment, false)); + } + segments.push_back(std::make_pair(it.second, true)); + segment = ""; + } + } + boost::algorithm::trim(segment); + if (segments.empty()) { + if (!segment.empty()) { + segments.push_back(std::make_pair(segment, false)); + } + } else if (segment != segments.back().first && !segment.empty()) { + segments.push_back(std::make_pair(segment, false)); + } +} + +void +SpeechHints::Alignments2Segments(Alignments alignments, std::vector& segments) +{ + if (!grammar_loaded) + return; + std::string segment; + for (auto it : alignments) { + if (it.first == it.second) { + segment = segment + it.first + " "; + } else { + boost::algorithm::trim(segment); + if (!segment.empty()) { + segments.push_back(segment); + } + segments.push_back(it.second); + segment = ""; + } + } + boost::algorithm::trim(segment); + if (segments.empty()) { + if (!segment.empty()) { + segments.push_back(segment); + } + } else if (segment != segments.back() && !segment.empty()) { + segments.push_back(segment); + } +} +} // namespace nemotron_speech::utils diff --git a/native/tts-normalizer/src/speech_hints.h b/native/tts-normalizer/src/speech_hints.h new file mode 100644 index 0000000..ab4940e --- /dev/null +++ b/native/tts-normalizer/src/speech_hints.h @@ -0,0 +1,198 @@ +/* + * Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. + * + * NVIDIA CORPORATION and its licensors retain all intellectual property + * and proprietary rights in and to this software, related documentation + * and any modifications thereto. Any use, reproduction, disclosure or + * distribution of this software and related documentation without an express + * license agreement from NVIDIA CORPORATION is strictly prohibited. + */ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#ifndef NEMOTRON_SPEECH_SPEECH_HINTS_H +#define NEMOTRON_SPEECH_SPEECH_HINTS_H + +namespace nemotron_speech::utils { + +using thrax::GrmManager; +typedef fst::VectorFst MutableTransducer; +typedef fst::StringCompiler Compiler; +typedef std::vector> Alignments; +typedef std::vector> NormalizedSegments; +typedef boost::tokenizer> tokenizer; + + +/** + * @brief SpeechHints class used for supporting speech hints based normalization. + */ +class SpeechHints { + public: + SpeechHints(); + + + /** + * @brief Loads class grammars for speech hints + * + * @param[in] far_filename Filename with path of the far file to be loaded for speech hints. + * FAR should contain "$__PASSTHROUGH__" and "$__SPACE__" fsts at a + * minimum + * + * @return True if grammar far files for speech hints is loaded successfully + */ + bool LoadGrammars(const std::string& far_filename); + + /** + * @brief Apply normalization on the string using the provided speech hints + * + * @param[in] input String to be normalized + * @param[in] speech_hints List of speech hints to be applied + * @param[out] output Normalized string + * @param[out] alignments Alignments between the input word(s) and the normalized segments + * + * @return True if normalization is successful + */ + bool ApplyHints( + const std::string& input, const std::vector& speech_hints, std::string& output, + Alignments& alignments); + + /** + * @brief Apply normalization on the string using the provided speech hint + * + * @param[in] input String to be normalized + * @param[in] speech_hint Speech hint to be applied + * @param[out] output Normalized string + * @param[out] alignments Alignments between the input word(s) and the normalized output segments + * + * @return True if normalization is successful + */ + bool ApplyHint( + const std::string& input, const std::string& speech_hint, std::string& output, + Alignments& alignments); + + /** + * @brief Given alignments, split text into a list of text segments. Segments contain text and + * bool (true if segment is normalized) + * + * @param[in] alignments + * @param[out] segments + */ + void GetSegments(Alignments alignments, NormalizedSegments& segments); + + /** + * @brief Given alignments, split text into a list of text segments. + * + * @param[in] alignments + * @param[out] segments + */ + void Alignments2Segments(Alignments alignments, std::vector& segments); + + /** + * @brief Get list of valid grammars present in the far file + * + * @return list of valid grammars + */ + std::vector GetGrammarList(void); + + /** + * @brief Checks if speech hints contains valid grammars + * + * @param speech_hints List of speech hints + * @return True if all grammars in speech hints are valid + */ + bool AreHintsValid(const std::vector& speech_hints); + + /** + * @brief Checks if speech hint contains valid grammars + * + * @param speech_hints Speech hint + * @return True if all grammars in speech hints are valid + */ + bool IsHintValid(const std::string& speech_hint); + + private: + const std::string space_grammar = "$__SPACE__"; + const std::string passthrough_grammar = "$__PASSTHROUGH__"; + boost::char_separator word_sep; + std::unique_ptr sh_class_; + bool grammar_loaded = false; + std::vector valid_grammars; + + MutableTransducer space_fst_, space_sigma_; + MutableTransducer sent_fst_; + bool StringFstToOutputLabels( + const MutableTransducer& ifst, std::vector* labels); + /** + * @brief Converts a given word to an fst + * + * @param word word that has to be converted to an FST + * @return FST for word + */ + MutableTransducer Word2Fst(const std::string& word); + + /** + * @brief Given a speech hint class name, returns the FST for that class + * + * @param[in] sh_class_name Name of the speech hint class + * @param[out] word_fst FST for the speech hint class + * @return True if speech hint class exists + */ + bool GetClassFst(const std::string& sh_class_name, MutableTransducer& word_fst); + + /** + * @brief Converts the provided speech hint into an FST + * + * @param[in] hint Speech hint string + * @return MutableTransducer containing the speech hint FST + */ + MutableTransducer Hint2Fst(const std::string& hint); + + /** + * @brief Apply normalization on the string using the provided speech hint + * + * @param[in] input String to be normalized + * @param[in] speech_hint Speech hint(s) (FST) to be applied + * @param[out] output Normalized string + * @param[out] alignments Alignments between the input word(s) and the normalized output segments + * + * @return True if normalization is successful + */ + bool ApplyHintFst( + const std::string& input, MutableTransducer speech_hint, std::string& output, + Alignments& alignments); + + + /** + * @brief Given an FST (obtained after applying the speech hints FST on input FST, returns the + * string and the alignments + * + * @param[in] ofst FST obtained after composing input FST with speech hints FST + * @param[out] output Normalized string + * @param[out] alignments Alignments between the input word(s) and the normalized output segments + * + * @return True if shortest path is obtained + */ + bool StringifyFst(MutableTransducer& ofst, std::string& output, Alignments& alignments); +}; + +} // namespace nemotron_speech::utils + +#endif // NEMOTRON_SPEECH_SPEECH_HINTS_H diff --git a/native/tts-normalizer/src/text_preprocessor.cc b/native/tts-normalizer/src/text_preprocessor.cc new file mode 100644 index 0000000..5ad2bc3 --- /dev/null +++ b/native/tts-normalizer/src/text_preprocessor.cc @@ -0,0 +1,510 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * NVIDIA CORPORATION and its licensors retain all intellectual property + * and proprietary rights in and to this software, related documentation + * and any modifications thereto. Any use, reproduction, disclosure or + * distribution of this software and related documentation without an express + * license agreement from NVIDIA CORPORATION is strictly prohibited. + */ + +#include "text_preprocessor.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "normalize.h" + +namespace nemotron_speech::tts::preprocessing { +namespace { + +class AbbreviationMapping { + public: + explicit AbbreviationMapping(const std::string& path) + { + std::ifstream input(path); + if (!input.good()) { + throw std::runtime_error("Unable to load requested abbreviations"); + } + std::string line; + while (std::getline(input, line)) { + if (line.empty() || !std::isalnum(static_cast(line.front()))) { + continue; + } + const auto separator = line.find(' '); + if (separator == std::string::npos || separator + 1 >= line.size()) { + continue; + } + const auto abbreviation = line.substr(0, separator); + const auto expansion = line.substr(separator + 1); + mappings_.emplace_back( + std::regex( + "\\b" + abbreviation + "(\\.|\\b)", std::regex_constants::ECMAScript | + std::regex_constants::optimize | + std::regex_constants::icase), + expansion); + } + if (mappings_.empty()) { + throw std::runtime_error("Abbreviation mapping is empty"); + } + } + + std::string Transform(const std::string& input) const + { + std::string output = input; + for (const auto& [pattern, expansion] : mappings_) { + output = std::regex_replace(output, pattern, expansion); + } + return output; + } + + private: + std::vector> mappings_; +}; + +bool +IsAsciiWhitespace(unsigned char c) +{ + return c == ' ' || c == '\t' || c == '\n' || c == '\r' || c == '\f' || c == '\v'; +} + +std::string +CollapseWhitespace(const std::string& input, bool preserve_boundary_whitespace = false) +{ + std::string output; + output.reserve(input.size()); + bool pending_space = false; + for (unsigned char c : input) { + if (IsAsciiWhitespace(c)) { + pending_space = preserve_boundary_whitespace || !output.empty(); + continue; + } + if (pending_space) { + output.push_back(' '); + pending_space = false; + } + output.push_back(static_cast(c)); + } + if (preserve_boundary_whitespace && pending_space && !output.empty() && output.back() != ' ') { + output.push_back(' '); + } + return output; +} + +std::string +Trim(const std::string& input) +{ + std::size_t begin = 0; + while (begin < input.size() && IsAsciiWhitespace(static_cast(input[begin]))) { + ++begin; + } + std::size_t end = input.size(); + while (end > begin && IsAsciiWhitespace(static_cast(input[end - 1]))) { + --end; + } + return input.substr(begin, end - begin); +} + +bool +ContainsMarkup(const std::string& input) +{ + const std::string trimmed = Trim(input); + if (trimmed.rfind("(trimmed[5]); + return IsAsciiWhitespace(boundary) || boundary == '?'; + } + if (trimmed.rfind("(trimmed[6]); + return IsAsciiWhitespace(boundary) || boundary == '>' || boundary == '/'; +} + +bool +IsValidUtf8(const std::string& input) +{ + std::size_t index = 0; + while (index < input.size()) { + const auto lead = static_cast(input[index]); + std::size_t continuation_count = 0; + uint32_t codepoint = 0; + if (lead <= 0x7f) { + ++index; + continue; + } else if ((lead & 0xe0) == 0xc0) { + continuation_count = 1; + codepoint = lead & 0x1f; + } else if ((lead & 0xf0) == 0xe0) { + continuation_count = 2; + codepoint = lead & 0x0f; + } else if ((lead & 0xf8) == 0xf0) { + continuation_count = 3; + codepoint = lead & 0x07; + } else { + return false; + } + if (index + continuation_count >= input.size()) { + return false; + } + for (std::size_t offset = 1; offset <= continuation_count; ++offset) { + const auto continuation = static_cast(input[index + offset]); + if ((continuation & 0xc0) != 0x80) { + return false; + } + codepoint = (codepoint << 6) | (continuation & 0x3f); + } + if ((continuation_count == 1 && codepoint < 0x80) || + (continuation_count == 2 && codepoint < 0x800) || + (continuation_count == 3 && codepoint < 0x10000) || + (codepoint >= 0xd800 && codepoint <= 0xdfff) || codepoint > 0x10ffff) { + return false; + } + index += continuation_count + 1; + } + return true; +} + +void +AppendSupportedSsml(const pugi::xml_node& node, std::string& output) +{ + for (const auto& child : node.children()) { + if (child.type() == pugi::node_pcdata || child.type() == pugi::node_cdata) { + output += child.value(); + continue; + } + if (child.type() == pugi::node_comment) { + continue; + } + if (child.type() != pugi::node_element) { + throw TextPreprocessorException( + PreprocessorErrorCode::kUnsupportedSsml, "Unsupported node in SSML input."); + } + + const std::string name = child.name(); + if (name != "sub") { + throw TextPreprocessorException( + PreprocessorErrorCode::kUnsupportedSsml, + "Unsupported SSML tag <" + name + ">. Only and are supported."); + } + const auto alias = child.attribute("alias"); + if (!alias || std::string(alias.value()).empty()) { + throw TextPreprocessorException( + PreprocessorErrorCode::kUnsupportedSsml, + "SSML requires a non-empty alias attribute."); + } + for (const auto& nested : child.children()) { + if (nested.type() == pugi::node_element) { + throw TextPreprocessorException( + PreprocessorErrorCode::kUnsupportedSsml, + "Nested SSML tags are not supported inside ."); + } + } + output += alias.value(); + } +} + +std::string +ExtractSpokenText(const std::string& input) +{ + if (!ContainsMarkup(input)) { + return input; + } + + pugi::xml_document document; + const auto result = + document.load_buffer(input.data(), input.size(), pugi::parse_default | pugi::parse_ws_pcdata); + if (!result) { + throw TextPreprocessorException( + PreprocessorErrorCode::kUnsupportedSsml, + std::string("Malformed SSML input: ") + result.description()); + } + const auto root = document.document_element(); + if (!root || std::string(root.name()) != "speak") { + throw TextPreprocessorException( + PreprocessorErrorCode::kUnsupportedSsml, "SSML root element must be ."); + } + for (const auto& attribute : root.attributes()) { + const std::string name = attribute.name(); + if (name != "version" && name != "xml:lang" && name != "xmlns") { + throw TextPreprocessorException( + PreprocessorErrorCode::kUnsupportedSsml, + "Unsupported attribute on SSML : " + name); + } + } + + std::string output; + AppendSupportedSsml(root, output); + return output; +} + +bool +HasSpokenContent(const std::string& input) +{ + for (unsigned char c : input) { + if (c >= 0x80 || std::isalnum(c)) { + return true; + } + } + return false; +} + +bool +PreviousCodePoint(const std::string& text, std::size_t& end, uint32_t& codepoint) +{ + if (end == 0 || end > text.size()) { + return false; + } + + std::size_t begin = end - 1; + while (begin > 0 && (static_cast(text[begin]) & 0xc0) == 0x80) { + --begin; + } + const auto lead = static_cast(text[begin]); + if (lead <= 0x7f) { + codepoint = lead; + } + else if ((lead & 0xe0) == 0xc0) { + codepoint = lead & 0x1f; + } + else if ((lead & 0xf0) == 0xe0) { + codepoint = lead & 0x0f; + } + else { + codepoint = lead & 0x07; + } + for (std::size_t index = begin + 1; index < end; ++index) { + codepoint = (codepoint << 6) | (static_cast(text[index]) & 0x3f); + } + end = begin; + return true; +} + +bool +IsSentenceCloser(uint32_t codepoint) +{ + return codepoint == '\'' || codepoint == '"' || codepoint == ')' || codepoint == ']' || + codepoint == '}' || codepoint == 0x2019 || codepoint == 0x201d; +} + +bool +IsTerminalPunctuation(uint32_t codepoint) +{ + // Keep the legacy semicolon/colon behavior and match the multilingual + // sentence terminators recognized by the Rust streaming boundary scanner. + return codepoint == '.' || codepoint == '!' || codepoint == '?' || codepoint == ';' || + codepoint == ':' || codepoint == 0x061f || codepoint == 0x06d4 || + codepoint == 0x0964 || codepoint == 0x0965 || codepoint == 0x3002 || + codepoint == 0xff01 || codepoint == 0xff1f; +} + +bool +HasTerminalPunctuation(const std::string& text) +{ + std::size_t end = text.size(); + uint32_t codepoint = 0; + while (PreviousCodePoint(text, end, codepoint)) { + if (IsSentenceCloser(codepoint)) { + continue; + } + return IsTerminalPunctuation(codepoint); + } + return false; +} + +} // namespace + +class TextPreprocessor::Impl { + public: + explicit Impl(const TextPreprocessorConfig& config) : config_(config) + { + if (config_.languages.empty()) { + throw TextPreprocessorException( + PreprocessorErrorCode::kUnsupportedLanguage, + "At least one language must be configured for text preprocessing."); + } + if (config_.max_input_bytes == 0 || config_.max_normalized_bytes == 0 || + config_.max_segments == 0) { + throw TextPreprocessorException( + PreprocessorErrorCode::kInvalidInput, "Preprocessor limits must be greater than zero."); + } + + for (const auto& [language, language_config] : config_.languages) { + if (language.empty()) { + throw TextPreprocessorException( + PreprocessorErrorCode::kUnsupportedLanguage, + "Configured language code cannot be empty."); + } + auto resources = std::make_unique(); + resources->config = language_config; + try { + if (!language_config.normalizer_path.empty()) { + const std::filesystem::path normalizer_path(language_config.normalizer_path); + if (!std::filesystem::is_directory(normalizer_path)) { + throw TextPreprocessorException( + PreprocessorErrorCode::kMissingAsset, + "Normalizer asset directory does not exist for " + language + ": " + + normalizer_path.string()); + } + if (language_config.direct_grammar) { + resources->normalizer = + std::make_unique(normalizer_path.string(), false); + } + else { + resources->normalizer = std::make_unique( + language_config.normalizer_config, normalizer_path.string()); + } + } + if (!language_config.abbreviations_path.empty()) { + if (!std::filesystem::is_regular_file(language_config.abbreviations_path)) { + throw TextPreprocessorException( + PreprocessorErrorCode::kMissingAsset, "Abbreviation asset does not exist for " + + language + ": " + + language_config.abbreviations_path); + } + resources->abbreviations = + std::make_unique(language_config.abbreviations_path); + } + } + catch (const TextPreprocessorException&) { + throw; + } + catch (const std::exception& error) { + throw TextPreprocessorException( + PreprocessorErrorCode::kMissingAsset, + "Failed to initialize preprocessing assets for " + language + ": " + error.what()); + } + resources_[language] = std::move(resources); + } + } + + std::vector Preprocess( + const std::string& input, const std::string& language_code, + const PreprocessOptions& options) + { + const auto language_it = resources_.find(language_code); + if (language_it == resources_.end()) { + throw TextPreprocessorException( + PreprocessorErrorCode::kUnsupportedLanguage, + "Unsupported preprocessing language: " + language_code); + } + if (input.empty() || input.size() > config_.max_input_bytes) { + std::ostringstream message; + message << "Input must contain 1-" << config_.max_input_bytes << " UTF-8 bytes."; + throw TextPreprocessorException(PreprocessorErrorCode::kInvalidInput, message.str()); + } + if (!IsValidUtf8(input)) { + throw TextPreprocessorException( + PreprocessorErrorCode::kInvalidInput, "Input must be valid UTF-8."); + } + + auto& resources = *language_it->second; + std::string normalized = ExtractSpokenText(input); + if (resources.normalizer) { + std::string output; + if (!resources.normalizer->Normalize(normalized, output)) { + throw TextPreprocessorException( + PreprocessorErrorCode::kInvalidInput, + "WFST normalization failed for the supplied input."); + } + normalized = std::move(output); + } + if (resources.abbreviations) { + normalized = resources.abbreviations->Transform(normalized); + } + normalized = CollapseWhitespace(normalized, options.continuous_streaming); + if (!HasSpokenContent(normalized)) { + throw TextPreprocessorException( + PreprocessorErrorCode::kInvalidInput, "Input must contain spoken text."); + } + if (normalized.size() > config_.max_normalized_bytes) { + std::ostringstream message; + message << "Normalized text exceeds " << config_.max_normalized_bytes << " UTF-8 bytes."; + throw TextPreprocessorException(PreprocessorErrorCode::kInvalidInput, message.str()); + } + + if (options.continuous_streaming) { + // A continuous request is already a caller-certified normalization unit. + // Splitting or adding punctuation here would create model-visible tokens + // that were not present in the append-only input stream. + return {TextSegment{std::move(normalized), 0, true}}; + } + + std::vector sentences; + if (resources.normalizer) { + sentences = resources.normalizer->SentenceSplitter(normalized); + } else { + sentences = {normalized}; + } + + std::vector segments; + segments.reserve(sentences.size()); + for (auto& sentence : sentences) { + sentence = CollapseWhitespace(sentence); + if (sentence.empty()) { + continue; + } + if (resources.config.append_punctuation && !HasTerminalPunctuation(sentence)) { + sentence.push_back('.'); + } + segments.push_back(TextSegment{sentence, static_cast(segments.size()), false}); + } + if (segments.empty()) { + throw TextPreprocessorException( + PreprocessorErrorCode::kInvalidInput, "Preprocessing produced no spoken text."); + } + if (segments.size() > config_.max_segments) { + std::ostringstream message; + message << "Preprocessing produced more than " << config_.max_segments << " segments."; + throw TextPreprocessorException(PreprocessorErrorCode::kInvalidInput, message.str()); + } + segments.back().is_last_sentence = true; + return segments; + } + + private: + struct LanguageResources { + LanguageConfig config; + std::unique_ptr normalizer; + std::unique_ptr abbreviations; + }; + + TextPreprocessorConfig config_; + std::unordered_map> resources_; +}; + +TextPreprocessor::TextPreprocessor(const TextPreprocessorConfig& config) + : impl_(std::make_unique(config)) +{ +} + +TextPreprocessor::~TextPreprocessor() = default; +TextPreprocessor::TextPreprocessor(TextPreprocessor&&) noexcept = default; +TextPreprocessor& TextPreprocessor::operator=(TextPreprocessor&&) noexcept = default; + +std::vector +TextPreprocessor::Preprocess( + const std::string& text, const std::string& language_code, const PreprocessOptions& options) +{ + try { + return impl_->Preprocess(text, language_code, options); + } + catch (const TextPreprocessorException&) { + throw; + } + catch (const std::exception& error) { + throw TextPreprocessorException(PreprocessorErrorCode::kInternal, error.what()); + } +} + +} // namespace nemotron_speech::tts::preprocessing diff --git a/native/tts-normalizer/src/text_preprocessor.h b/native/tts-normalizer/src/text_preprocessor.h new file mode 100644 index 0000000..d46465d --- /dev/null +++ b/native/tts-normalizer/src/text_preprocessor.h @@ -0,0 +1,107 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * NVIDIA CORPORATION and its licensors retain all intellectual property + * and proprietary rights in and to this software, related documentation + * and any modifications thereto. Any use, reproduction, disclosure or + * distribution of this software and related documentation without an express + * license agreement from NVIDIA CORPORATION is strictly prohibited. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace nemotron_speech::tts::preprocessing { + +enum class PreprocessorErrorCode { + kInvalidInput, + kUnsupportedLanguage, + kUnsupportedSsml, + kMissingAsset, + kInternal, +}; + +class TextPreprocessorException : public std::runtime_error { + public: + TextPreprocessorException(PreprocessorErrorCode code, const std::string& message) + : std::runtime_error(message), code_(code) + { + } + + PreprocessorErrorCode code() const noexcept { return code_; } + + private: + PreprocessorErrorCode code_; +}; + +struct LanguageConfig { + std::string normalizer_path; + std::string normalizer_config = "sparrowhawk_configuration.ascii_proto"; + std::string abbreviations_path; + bool append_punctuation = true; + // Load tokenize_and_classify.far and verbalize.far directly from + // normalizer_path. This matches AggregatedTokenizer's non-Chinese path. + bool direct_grammar = false; +}; + +struct TextPreprocessorConfig { + std::unordered_map languages; + std::size_t max_input_bytes = 64 * 1024; + std::size_t max_normalized_bytes = 256 * 1024; + std::size_t max_segments = 256; +}; + +struct PreprocessOptions { + /** + * Normalize one caller-certified unit of a continuous text stream. + * + * Continuous mode deliberately does not run sentence splitting or append + * punctuation. Boundary whitespace is retained canonically as at most one + * ASCII space on either side of the returned segment. The caller must only + * submit prefixes whose normalization cannot be changed by future text. + */ + bool continuous_streaming = false; +}; + +struct TextSegment { + std::string text; + int32_t sentence_num = 0; + bool is_last_sentence = false; +}; + +/** + * CPU-only, model-agnostic TTS text preprocessing. + * + * The class deliberately returns UTF-8 text rather than model token IDs. It + * does not depend on Triton, Torch, CUDA, G2P, character mappings, or Python. + * Instances are immutable after construction but are not advertised as + * thread-safe because the underlying WFST implementation has not established + * a cross-platform concurrent-call contract. + */ +class TextPreprocessor { + public: + explicit TextPreprocessor(const TextPreprocessorConfig& config); + ~TextPreprocessor(); + + TextPreprocessor(const TextPreprocessor&) = delete; + TextPreprocessor& operator=(const TextPreprocessor&) = delete; + TextPreprocessor(TextPreprocessor&&) noexcept; + TextPreprocessor& operator=(TextPreprocessor&&) noexcept; + + std::vector Preprocess( + const std::string& text, const std::string& language_code, + const PreprocessOptions& options = {}); + + private: + class Impl; + std::unique_ptr impl_; +}; + +} // namespace nemotron_speech::tts::preprocessing diff --git a/native/tts-normalizer/src/text_preprocessor_c.cc b/native/tts-normalizer/src/text_preprocessor_c.cc new file mode 100644 index 0000000..dbcac86 --- /dev/null +++ b/native/tts-normalizer/src/text_preprocessor_c.cc @@ -0,0 +1,330 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + */ + +#include +#include +#include +#include +#include +#include + +#include "processor_abi.h" +#include "text_preprocessor.h" + +namespace { + +using nemotron_speech::tts::preprocessing::LanguageConfig; +using nemotron_speech::tts::preprocessing::PreprocessOptions; +using nemotron_speech::tts::preprocessing::PreprocessorErrorCode; +using nemotron_speech::tts::preprocessing::TextPreprocessor; +using nemotron_speech::tts::preprocessing::TextPreprocessorConfig; +using nemotron_speech::tts::preprocessing::TextPreprocessorException; + +constexpr uint32_t kSchemaVersion = 1; +constexpr char kProcessorId[] = "nemotron_speech.tts.text_preprocessor"; + +struct ProcessorState { + explicit ProcessorState(TextPreprocessorConfig config) : processor(std::move(config)) {} + TextPreprocessor processor; +}; + +struct OutputStorage { + std::string request_id; + std::string session_id; + std::vector names; + std::vector data; + std::vector views; +}; + +NemotronSpeechByteSpanV1 +Span(const std::string& value) +{ + return {reinterpret_cast(value.data()), value.size()}; +} + +NemotronSpeechByteSpanV1 +Span(const char* value) +{ + return {reinterpret_cast(value), std::strlen(value)}; +} + +NemotronSpeechStatusV1 +Status(int32_t code, const std::string& message = "") noexcept +{ + NemotronSpeechStatusV1 status{ + sizeof(NemotronSpeechStatusV1), code, {nullptr, 0}, nullptr}; + try { + if (!message.empty()) { + auto owned = std::make_unique(message); + status.message = Span(*owned); + status.context = owned.release(); + } + } + catch (...) { + status.message = Span("native processor error (message unavailable)"); + } + return status; +} + +void +ReleaseStatus(NemotronSpeechStatusV1* status) +{ + if (status == nullptr) { + return; + } + delete static_cast(status->context); + *status = {sizeof(NemotronSpeechStatusV1), 0, {nullptr, 0}, nullptr}; +} + +std::string +CopySpan(NemotronSpeechByteSpanV1 value, const char* field) +{ + if (value.len != 0 && value.data == nullptr) { + throw std::invalid_argument(std::string(field) + " pointer is null"); + } + if (value.len == 0) { + return ""; + } + return {reinterpret_cast(value.data), value.len}; +} + +const NemotronSpeechRecordViewV1* +FindRecord(const NemotronSpeechEnvelopeViewV1& envelope, const std::string& name) +{ + if (envelope.record_count != 0 && envelope.records == nullptr) { + throw std::invalid_argument("record array is null"); + } + for (size_t index = 0; index < envelope.record_count; ++index) { + const auto& record = envelope.records[index]; + if (record.struct_size < sizeof(NemotronSpeechRecordViewV1)) { + throw std::invalid_argument("record structure is too small"); + } + if (CopySpan(record.name, "record name") == name) { + return &record; + } + } + return nullptr; +} + +std::string +Utf8Record( + const NemotronSpeechEnvelopeViewV1& envelope, const std::string& name, + bool required = true) +{ + const auto* record = FindRecord(envelope, name); + if (record == nullptr) { + if (required) { + throw std::invalid_argument("missing UTF-8 record: " + name); + } + return ""; + } + if (record->kind != NEMOTRON_SPEECH_RECORD_UTF8_V1) { + throw std::invalid_argument("record is not UTF-8: " + name); + } + return CopySpan(record->data, name.c_str()); +} + +bool +BooleanRecord( + const NemotronSpeechEnvelopeViewV1& envelope, const std::string& name, + bool default_value) +{ + const auto* record = FindRecord(envelope, name); + if (record == nullptr) { + return default_value; + } + if (record->kind != NEMOTRON_SPEECH_RECORD_BOOLEAN_V1 || record->data.len != 1 || + record->data.data == nullptr) { + throw std::invalid_argument("record is not a boolean: " + name); + } + return record->data.data[0] != 0; +} + +size_t +IntegerRecord( + const NemotronSpeechEnvelopeViewV1& envelope, const std::string& name, + size_t default_value) +{ + const auto* record = FindRecord(envelope, name); + if (record == nullptr) { + return default_value; + } + if (record->kind != NEMOTRON_SPEECH_RECORD_INTEGER_V1 || + record->data.len != sizeof(int64_t) || + record->data.data == nullptr) { + throw std::invalid_argument("record is not an integer: " + name); + } + int64_t value = 0; + std::memcpy(&value, record->data.data, sizeof(value)); + if (value <= 0) { + throw std::invalid_argument(name + " must be positive"); + } + return static_cast(value); +} + +NemotronSpeechStatusV1 +Capabilities(NemotronSpeechProcessorCapabilitiesV1* output) +{ + if (output == nullptr || + output->struct_size < sizeof(NemotronSpeechProcessorCapabilitiesV1)) { + return Status(1, "capability output structure is missing or too small"); + } + output->processor_id = Span(kProcessorId); + output->batch_input = 0; + output->streaming_input = 1; + output->thread_safe = 0; + output->stateful = 0; + return Status(0); +} + +NemotronSpeechStatusV1 +Create( + const NemotronSpeechEnvelopeViewV1* config, + NemotronSpeechProcessorHandleV1* output) +{ + if (config == nullptr || + config->struct_size < sizeof(NemotronSpeechEnvelopeViewV1) || output == nullptr) { + return Status(1, "processor configuration is missing or too small"); + } + *output = nullptr; + try { + if (config->schema_version != kSchemaVersion) { + return Status(1, "unsupported configuration schema version"); + } + const auto language_code = Utf8Record(*config, "language_code"); + LanguageConfig language; + language.normalizer_path = Utf8Record(*config, "normalizer_path", false); + const auto normalizer_config = Utf8Record(*config, "normalizer_config", false); + if (!normalizer_config.empty()) { + language.normalizer_config = normalizer_config; + } + language.abbreviations_path = Utf8Record(*config, "abbreviations_path", false); + language.append_punctuation = BooleanRecord(*config, "append_punctuation", true); + language.direct_grammar = BooleanRecord(*config, "direct_grammar", false); + + TextPreprocessorConfig processor_config; + processor_config.languages.emplace(language_code, std::move(language)); + processor_config.max_input_bytes = + IntegerRecord(*config, "max_input_bytes", processor_config.max_input_bytes); + processor_config.max_normalized_bytes = + IntegerRecord(*config, "max_normalized_bytes", processor_config.max_normalized_bytes); + processor_config.max_segments = + IntegerRecord(*config, "max_segments", processor_config.max_segments); + *output = new ProcessorState(std::move(processor_config)); + return Status(0); + } + catch (const TextPreprocessorException& error) { + return Status(100 + static_cast(error.code()), error.what()); + } + catch (const std::exception& error) { + return Status(1, error.what()); + } + catch (...) { + return Status(1, "unknown exception while creating text processor"); + } +} + +void +Destroy(NemotronSpeechProcessorHandleV1 handle) +{ + delete static_cast(handle); +} + +NemotronSpeechStatusV1 +Process( + NemotronSpeechProcessorHandleV1 handle, const NemotronSpeechEnvelopeViewV1* request, + NemotronSpeechEnvelopeOutputV1* output) +{ + if (handle == nullptr || request == nullptr || output == nullptr || + request->struct_size < sizeof(NemotronSpeechEnvelopeViewV1) || + output->struct_size < sizeof(NemotronSpeechEnvelopeOutputV1)) { + return Status(1, "processor handle, request, or output is invalid"); + } + if (output->context != nullptr) { + return Status(1, "output must be released before reuse"); + } + try { + if (request->schema_version != kSchemaVersion) { + return Status(1, "unsupported request schema version"); + } + const auto text = Utf8Record(*request, "text"); + const auto language_code = Utf8Record(*request, "language_code"); + PreprocessOptions options; + options.continuous_streaming = BooleanRecord(*request, "continuous_streaming", false); + const auto segments = static_cast(handle)->processor.Preprocess( + text, language_code, options); + + auto storage = std::make_unique(); + storage->request_id = CopySpan(request->request_id, "request ID"); + storage->session_id = CopySpan(request->session_id, "session ID"); + storage->names.reserve(segments.size()); + storage->data.reserve(segments.size()); + storage->views.reserve(segments.size()); + for (size_t index = 0; index < segments.size(); ++index) { + storage->names.push_back("normalized_text." + std::to_string(index)); + storage->data.push_back(segments[index].text); + } + for (size_t index = 0; index < segments.size(); ++index) { + NemotronSpeechRecordViewV1 view{}; + view.struct_size = sizeof(NemotronSpeechRecordViewV1); + view.name = Span(storage->names[index]); + view.kind = NEMOTRON_SPEECH_RECORD_UTF8_V1; + view.element_type = NEMOTRON_SPEECH_ELEMENT_U8_V1; + view.data = Span(storage->data[index]); + view.sequence = static_cast(segments[index].sentence_num); + view.has_sequence = 1; + view.is_final = segments[index].is_last_sentence ? 1 : 0; + storage->views.push_back(view); + } + + output->schema_version = kSchemaVersion; + output->request_id = Span(storage->request_id); + output->session_id = Span(storage->session_id); + output->records = storage->views.data(); + output->record_count = storage->views.size(); + output->sequence = request->sequence; + output->has_sequence = request->has_sequence; + output->is_final = request->is_final; + output->context = storage.release(); + return Status(0); + } + catch (const TextPreprocessorException& error) { + return Status(100 + static_cast(error.code()), error.what()); + } + catch (const std::exception& error) { + return Status(1, error.what()); + } + catch (...) { + return Status(1, "unknown exception while processing text"); + } +} + +void +ReleaseOutput(NemotronSpeechProcessorHandleV1, NemotronSpeechEnvelopeOutputV1* output) +{ + if (output == nullptr) { + return; + } + delete static_cast(output->context); + *output = {}; + output->struct_size = sizeof(NemotronSpeechEnvelopeOutputV1); +} + +const NemotronSpeechProcessorApiV1 kApi = { + sizeof(NemotronSpeechProcessorApiV1), + NEMOTRON_SPEECH_PROCESSOR_ABI_VERSION, + Capabilities, + Create, + Destroy, + Process, + ReleaseOutput, + ReleaseStatus}; + +} // namespace + +extern "C" const NemotronSpeechProcessorApiV1* +nemotron_speech_processor_get_api(uint32_t requested_version) +{ + return requested_version == NEMOTRON_SPEECH_PROCESSOR_ABI_VERSION ? &kApi : nullptr; +} diff --git a/native/tts-normalizer/testdata/BUILD.bazel b/native/tts-normalizer/testdata/BUILD.bazel new file mode 100644 index 0000000..4c33f07 --- /dev/null +++ b/native/tts-normalizer/testdata/BUILD.bazel @@ -0,0 +1,5 @@ +filegroup( + name = "english_assets", + srcs = glob(["en-US/**"]), + visibility = ["//visibility:public"], +) diff --git a/native/tts-normalizer/testdata/en-US/abbr.txt b/native/tts-normalizer/testdata/en-US/abbr.txt new file mode 100644 index 0000000..b499039 --- /dev/null +++ b/native/tts-normalizer/testdata/en-US/abbr.txt @@ -0,0 +1,18 @@ +mrs misses +mr mister +dr doctor +st saint +gen general +co company +jr junior +maj major +drs doctors +rev reverend +lt lieutenant +hon honorable +sgt sergeant +capt captain +esq esquire +ltd limited +col colonel +ft fort diff --git a/native/tts-normalizer/testdata/en-US/post_process.far b/native/tts-normalizer/testdata/en-US/post_process.far new file mode 100644 index 0000000..45bc69d Binary files /dev/null and b/native/tts-normalizer/testdata/en-US/post_process.far differ diff --git a/native/tts-normalizer/testdata/en-US/postprocessor.ascii_proto b/native/tts-normalizer/testdata/en-US/postprocessor.ascii_proto new file mode 100644 index 0000000..693c56f --- /dev/null +++ b/native/tts-normalizer/testdata/en-US/postprocessor.ascii_proto @@ -0,0 +1,5 @@ +grammar_file: "post_process.far" + +grammar_name: "PostProcessor" + +rules { main: "POSTPROCESSOR" } diff --git a/native/tts-normalizer/testdata/en-US/sentence_boundary_exceptions.txt b/native/tts-normalizer/testdata/en-US/sentence_boundary_exceptions.txt new file mode 100644 index 0000000..db99a06 --- /dev/null +++ b/native/tts-normalizer/testdata/en-US/sentence_boundary_exceptions.txt @@ -0,0 +1,15 @@ +Mr. +Dr. +Mrs. +St. +Jan. +Feb. +Mar. +Apr. +Jun. +Jul. +Aug. +Sep. +Oct. +Nov. +Dec. diff --git a/native/tts-normalizer/testdata/en-US/sparrowhawk_configuration.ascii_proto b/native/tts-normalizer/testdata/en-US/sparrowhawk_configuration.ascii_proto new file mode 100644 index 0000000..0630169 --- /dev/null +++ b/native/tts-normalizer/testdata/en-US/sparrowhawk_configuration.ascii_proto @@ -0,0 +1,9 @@ +tokenizer_grammar: "tokenizer.ascii_proto" + +verbalizer_grammar: "verbalizer.ascii_proto" + +postprocessor_grammar: "postprocessor.ascii_proto" + +sentence_boundary_regexp: "[\\.:!\\?] " + +sentence_boundary_exceptions_file: "sentence_boundary_exceptions.txt" diff --git a/native/tts-normalizer/testdata/en-US/tokenize_and_classify.far b/native/tts-normalizer/testdata/en-US/tokenize_and_classify.far new file mode 100644 index 0000000..643da26 Binary files /dev/null and b/native/tts-normalizer/testdata/en-US/tokenize_and_classify.far differ diff --git a/native/tts-normalizer/testdata/en-US/tokenizer.ascii_proto b/native/tts-normalizer/testdata/en-US/tokenizer.ascii_proto new file mode 100644 index 0000000..f2480c8 --- /dev/null +++ b/native/tts-normalizer/testdata/en-US/tokenizer.ascii_proto @@ -0,0 +1,5 @@ +grammar_file: "tokenize_and_classify.far" + +grammar_name: "TokenizerClassifier" + +rules { main: "TOKENIZE_AND_CLASSIFY" } diff --git a/native/tts-normalizer/testdata/en-US/verbalize.far b/native/tts-normalizer/testdata/en-US/verbalize.far new file mode 100644 index 0000000..87c68ed Binary files /dev/null and b/native/tts-normalizer/testdata/en-US/verbalize.far differ diff --git a/native/tts-normalizer/testdata/en-US/verbalizer.ascii_proto b/native/tts-normalizer/testdata/en-US/verbalizer.ascii_proto new file mode 100644 index 0000000..a7aa9e0 --- /dev/null +++ b/native/tts-normalizer/testdata/en-US/verbalizer.ascii_proto @@ -0,0 +1,5 @@ +grammar_file: "verbalize.far" + +grammar_name: "Verbalizer" + +rules { main: "ALL" redup: "REDUP" } diff --git a/native/tts-normalizer/tests/normalize_test.cc b/native/tts-normalizer/tests/normalize_test.cc new file mode 100644 index 0000000..32444af --- /dev/null +++ b/native/tts-normalizer/tests/normalize_test.cc @@ -0,0 +1,119 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + */ + +#include "normalize.h" + +#include +#include + +#include +#include +#include +#include +#include + +namespace nemotron_speech::utils { +namespace { + +class DirectNormalizerTest : public testing::Test { + protected: + static void SetUpTestSuite() + { + normalizer_ = std::make_unique("testdata/en-US"); + } + + static void TearDownTestSuite() { normalizer_.reset(); } + + inline static std::unique_ptr normalizer_; +}; + +TEST_F(DirectNormalizerTest, NormalizesTextAndPreservesBoundaryWhitespace) +{ + std::string output; + + EXPECT_TRUE(normalizer_->Normalize(" \tI spent $2 today\r\n", output)); + EXPECT_EQ(output, " \tI spent two dollars today\r\n"); +} + +TEST_F(DirectNormalizerTest, EmptyAndWhitespaceOnlyInputsRemainStable) +{ + std::string output = "stale output"; + + EXPECT_TRUE(normalizer_->Normalize("", output)); + EXPECT_TRUE(output.empty()); + EXPECT_TRUE(normalizer_->Normalize(" \t\r\n", output)); + EXPECT_EQ(output, " \t\r\n"); +} + +TEST_F(DirectNormalizerTest, DisablesSpeechHintsWhenTheArchiveIsAbsent) +{ + EXPECT_FALSE(normalizer_->IsHintValid("$OOV_ALPHA_SEQUENCE")); + EXPECT_FALSE(normalizer_->AreHintsValid({"$OOV_ALPHA_SEQUENCE"})); +} + +TEST_F(DirectNormalizerTest, PreservesPunctuationAndLiteralSil) +{ + for (const std::string input : { + "Hello, world, this is a test.", "Stop. Look. Listen.", "The word sil remains sil."}) { + std::string output; + EXPECT_TRUE(normalizer_->Normalize(input, output)); + EXPECT_EQ(output, input); + } +} + +TEST(NormalizerTest, PreservesCommasInsideVerbalizedDates) +{ + const auto directory = std::filesystem::path(testing::TempDir()) / "date-comma-grammar"; + std::filesystem::create_directories(directory); + std::filesystem::copy_file( + "testdata/en-US/verbalize.far", directory / "verbalize.far", + std::filesystem::copy_options::overwrite_existing); + + // The deployed date grammar embeds the comma inside the year field. + const std::string input = "July 28, 1914"; + const std::string tagged = + "tokens { date { month: \"july\" day: \"twenty eight\" " + "year: \", nineteen fourteen\" preserve_order: true } }"; + fst::VectorFst tagger; + auto state = tagger.AddState(); + tagger.SetStart(state); + for (size_t i = 0; i < tagged.size(); ++i) { + const auto next = tagger.AddState(); + tagger.AddArc(state, fst::StdArc( + i < input.size() ? input[i] : 0, tagged[i], + fst::StdArc::Weight::One(), next)); + state = next; + } + tagger.SetFinal(state, fst::StdArc::Weight::One()); + { + std::unique_ptr> writer( + fst::FarWriter::Create((directory / "tokenize_and_classify.far").string())); + ASSERT_NE(writer, nullptr); + writer->Add("TOKENIZE_AND_CLASSIFY", tagger); + ASSERT_FALSE(writer->Error()); + } + + Normalizer normalizer(directory.string()); + std::string output; + ASSERT_TRUE(normalizer.Normalize(input, output)); + EXPECT_EQ(output, "july twenty eighth , nineteen fourteen"); +} + +TEST(NormalizerTest, FullGrammarConfigurationNormalizesText) +{ + Normalizer normalizer( + std::string("sparrowhawk_configuration.ascii_proto"), std::string("testdata/en-US")); + std::string output; + + EXPECT_TRUE(normalizer.Normalize("I spent $2 today", output)); + EXPECT_EQ(output, "I spent two dollars today"); +} + +TEST(NormalizerTest, MissingGrammarDirectoryIsRejected) +{ + EXPECT_THROW(Normalizer("testdata/missing-normalizer"), std::runtime_error); +} + +} // namespace +} // namespace nemotron_speech::utils diff --git a/native/tts-normalizer/tests/speech_hints_test.cc b/native/tts-normalizer/tests/speech_hints_test.cc new file mode 100644 index 0000000..bdb015e --- /dev/null +++ b/native/tts-normalizer/tests/speech_hints_test.cc @@ -0,0 +1,47 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + */ + +#include "speech_hints.h" + +#include + +#include +#include + +namespace nemotron_speech::utils { +namespace { + +TEST(SpeechHintsTest, RejectsOperationsBeforeGrammarLoading) +{ + SpeechHints hints; + std::string output = "unchanged"; + Alignments alignments = {{"input", "output"}}; + + EXPECT_TRUE(hints.GetGrammarList().empty()); + EXPECT_FALSE(hints.IsHintValid("$OOV_ALPHA_SEQUENCE")); + EXPECT_FALSE(hints.AreHintsValid({"$OOV_ALPHA_SEQUENCE"})); + EXPECT_FALSE(hints.ApplyHint("hello", "$OOV_ALPHA_SEQUENCE", output, alignments)); + EXPECT_EQ(output, "unchanged"); + EXPECT_EQ(alignments, Alignments({{"input", "output"}})); + EXPECT_FALSE(hints.ApplyHints("hello", {"$OOV_ALPHA_SEQUENCE"}, output, alignments)); + EXPECT_EQ(output, "unchanged"); + EXPECT_EQ(alignments, Alignments({{"input", "output"}})); +} + +TEST(SpeechHintsTest, SegmentConversionIsANoOpBeforeGrammarLoading) +{ + SpeechHints hints; + const Alignments alignments = {{"plain", "plain"}, {"spoken", "normalized"}}; + NormalizedSegments normalized_segments = {{"existing", false}}; + std::vector segments = {"existing"}; + + hints.GetSegments(alignments, normalized_segments); + hints.Alignments2Segments(alignments, segments); + + EXPECT_EQ(normalized_segments, NormalizedSegments({{"existing", false}})); + EXPECT_EQ(segments, std::vector({"existing"})); +} + +} // namespace +} // namespace nemotron_speech::utils diff --git a/native/tts-normalizer/tests/string_matcher_test.cc b/native/tts-normalizer/tests/string_matcher_test.cc new file mode 100644 index 0000000..76f8629 --- /dev/null +++ b/native/tts-normalizer/tests/string_matcher_test.cc @@ -0,0 +1,294 @@ +#include "sparrowhawk_string_matcher.h" + +#include +#include +#include + +using Arc = fst::StdArc; +using Fst = fst::VectorFst; +using Matcher = speech::sparrowhawk::StringMatcher; + +static int checks = 0; +static void Check(bool condition, const char* label) { + ++checks; + if (!condition) { + std::cerr << "FAIL " << label << '\n'; + std::exit(1); + } +} + +static Fst Input(const std::string& text) { + Fst result; + result.AddState(); + result.SetStart(0); + for (unsigned char byte : text) { + int next = result.AddState(); + result.AddArc(next - 1, Arc(byte, byte, 0, next)); + } + result.SetFinal(result.NumStates() - 1, 0); + return result; +} + +static Fst Grammar(int states) { + Fst result; + for (int i = 0; i < states; ++i) result.AddState(); + result.SetStart(0); + return result; +} + +static std::string Output(const Fst& result, bool input = false) { + std::string text; + int state = result.Start(); + for (int i = 0; i < result.NumStates(); ++i) { + fst::ArcIterator it(result, state); + if (it.Done()) return text; + Arc arc = it.Value(); + if (input ? arc.ilabel : arc.olabel) + text.push_back(char(input ? arc.ilabel : arc.olabel)); + state = arc.nextstate; + } + Check(false, "output not a simple path"); + return {}; +} + +static void Unchanged(const Fst& output) { + Check(Output(output) == "sentinel", "fallback leaves output untouched"); +} + +static void AddOutputPath(Fst* grammar, const std::string& text, bool separate_input) { + int state = 0; + if (separate_input) { + int next = grammar->AddState(); + grammar->AddArc(state, Arc('a', 0, 0, next)); + state = next; + } + bool first = true; + for (unsigned char byte : text) { + int next = grammar->AddState(); + grammar->AddArc(state, Arc(first && !separate_input ? 'a' : 0, byte, 0, next)); + first = false; + state = next; + } + grammar->SetFinal(state, 0); +} + +static void CanonicalCase(const std::string& first, const std::string& second, + bool accepted, const char* label) { + Fst grammar = Grammar(1); + AddOutputPath(&grammar, first, false); + AddOutputPath(&grammar, second, true); + Matcher::Limits limits; + limits.canonical_word_outputs = true; + Matcher matcher(grammar, limits); + Fst output = Input("sentinel"); + Check(matcher.Match(Input("a"), &output) == accepted, label); + if (accepted) { + Check(Output(output) == first || Output(output) == second, + "canonical identity never rewrites emitted output"); + Check(Output(output, true) == "a", "canonical mode retains actual input labels"); + } else { + Unchanged(output); + } + Check(matcher.Match(Input("a"), &output) == accepted, + "canonical prefix generations reset across calls"); +} + +int main() { + { + const std::string one = R"(tokens { name: "Alice" } tokens { name: "Bob" })"; + const std::string many = R"(tokens { name: "Alice" } tokens { name: "Bob" })"; + CanonicalCase(one, many, true, "outside repeated spaces coalesce"); + Fst grammar = Grammar(1); + AddOutputPath(&grammar, one, false); + AddOutputPath(&grammar, many, true); + Matcher original(grammar); + Fst output = Input("sentinel"); + Check(!original.Match(Input("a"), &output), "canonical mode is opt-in"); + Unchanged(output); + } + CanonicalCase(R"(tokens { name: "Alice Bob" })", R"(tokens { name: "Alice Bob" })", + false, "quoted repeated spaces remain distinct"); + CanonicalCase(R"(tokens { name: "Alice\" Bob" })", R"(tokens { name: "Alice\" Bob" })", + false, "escaped quote does not end quoted whitespace protection"); + CanonicalCase(R"(tokens { name: "A\\" } tokens { name: "B" })", + R"(tokens { name: "A\\" } tokens { name: "B" })", + true, "even backslash parity permits outside collapse"); + CanonicalCase(R"(tokens { name: "A\\\" B" })", R"(tokens { name: "A\\\" B" })", + false, "odd backslash parity retains quoted spaces"); + CanonicalCase(R"(tokens { name: "Alice" })", R"(tokens { name: "Carol" })", + false, "different word names are genuine output ambiguity"); + CanonicalCase(R"(tokens { name: "unterminated })", R"(tokens { name: "unterminated })", + false, "unbalanced normal escape syntax cannot enter WORD shortcut"); + CanonicalCase(R"(tokens { date { month: "january" day: "two" } })", + R"(tokens { date { month: "february" day: "three" } })", + false, "different date semantics never coalesce"); + { + Fst grammar = Grammar(3); + grammar.AddArc(0, Arc('a', 'A', 0, 1)); + grammar.AddArc(0, Arc('a', 'A', 0, 1)); + grammar.AddArc(0, Arc('a', 'B', 10, 2)); + grammar.SetFinal(1, 0); + grammar.SetFinal(2, 0); + Matcher matcher(grammar); + Fst output; + Check(matcher.Match(Input("a"), &output), "same-prefix parallel paths coalesce"); + Check(Output(output) == "A", "same-prefix winning output"); + Check(matcher.diagnostics().second_cost == 10, "runner-up is distinct prefix"); + } + { + Fst grammar = Grammar(5); + grammar.AddArc(0, Arc('a', 'A', 2, 1)); + grammar.AddArc(0, Arc(0, 0, 0, 2)); + grammar.AddArc(2, Arc('a', 'A', 0, 1)); + grammar.AddArc(1, Arc('b', 'B', 0, 3)); + grammar.AddArc(1, Arc('b', 'X', 5, 4)); + grammar.SetFinal(3, 0); + grammar.SetFinal(4, 0); + Matcher matcher(grammar); + Fst output; + Check(matcher.Match(Input("ab"), &output), "double-epsilon keeps canonical prefix"); + Check(Output(output) == "AB", "best weight retained when same prefix merges"); + Check(matcher.diagnostics().best_cost == 0, "better same-prefix raw weight"); + Check(matcher.diagnostics().second_cost == 5, "distinct runner-up propagates suffix"); + } + { + Fst grammar = Grammar(104); + grammar.AddArc(0, Arc('a', 'A', 0, 103)); + grammar.SetFinal(103, 0); + for (int i = 1; i <= 100; ++i) { + grammar.AddArc(0, Arc(0, 'x', -1, i)); + grammar.AddArc(i, Arc('b', 'B', 0, 102)); + } + grammar.SetFinal(102, 0); + Matcher matcher(grammar); + Fst output; + Check(matcher.Match(Input("a"), &output), "next-byte pruning preserves winner"); + Check(matcher.diagnostics().visited_states == 2, "dead epsilon states not enqueued"); + Check(Output(output) == "A", "pruned output exact"); + } + { + Fst grammar = Grammar(5); + grammar.AddArc(0, Arc('a', 'A', 0, 1)); + grammar.AddArc(0, Arc('a', 'B', -20, 2)); + grammar.AddArc(1, Arc(0, 'Z', 0, 3)); + grammar.AddArc(2, Arc('b', 'X', 0, 4)); + grammar.SetFinal(3, 0); + grammar.SetFinal(4, 0); + Matcher matcher(grammar); + Fst output; + Check(matcher.Match(Input("a"), &output), "end-of-input can-final pruning"); + Check(Output(output) == "AZ", "negative dead suffix not selected"); + } + { + Fst grammar = Grammar(2); + grammar.AddArc(0, Arc('a', 'A', 0, 1)); + grammar.AddArc(0, Arc('a', 'B', 0, 1)); + grammar.AddArc(0, Arc('a', 'C', 0, 1)); + grammar.SetFinal(1, 0); + Matcher matcher(grammar); + Fst output = Input("sentinel"); + Check(!matcher.Match(Input("a"), &output), "third distinct tie stays guarded"); + Check(matcher.diagnostics().discarded_cost == 0, "discarded witness reaches final"); + Unchanged(output); + } + { + Fst grammar = Grammar(5); + grammar.AddArc(0, Arc(0, 'x', 1, 3)); + grammar.AddArc(0, Arc(0, 'y', 5, 1)); + grammar.AddArc(1, Arc(0, 'z', -10, 3)); + grammar.AddArc(3, Arc('a', 'A', 0, 4)); + grammar.SetFinal(4, 0); + Matcher matcher(grammar); + Fst output; + Check(matcher.valid(), "negative epsilon DAG supported"); + Check(matcher.Match(Input("a"), &output), "negative epsilon path accepted"); + Check(Output(output) == "yzA", "epsilon topology updates before propagation"); + Check(Output(output, true) == "a", "input labels preserved"); + Check(matcher.diagnostics().best_cost == -5, "negative path cost"); + Check(matcher.diagnostics().second_cost == 1, "runner-up through shared suffix"); + } + { + Fst grammar = Grammar(4); + grammar.AddArc(0, Arc('a', 'A', 3, 1)); + grammar.AddArc(0, Arc('a', 'X', 1, 2)); + grammar.AddArc(1, Arc('b', 'B', -10, 3)); + grammar.AddArc(2, Arc('b', 'Y', 0, 3)); + grammar.SetFinal(3, 0); + Matcher matcher(grammar); + Fst output; + Check(matcher.Match(Input("ab"), &output), "negative consuming path"); + Check(Output(output) == "AB", "not greedy or Dijkstra early termination"); + output = Input("sentinel"); + Check(!matcher.Match(Input("a"), &output), "incomplete input has no path"); + Unchanged(output); + Check(matcher.Match(Input("ab"), &output), "generation reuse after fallback"); + Check(Output(output) == "AB", "no stale previous-position state"); + } + for (bool merge : {false, true}) { + Fst grammar = Grammar(4); + grammar.AddArc(0, Arc('a', 'A', 0, 1)); + grammar.AddArc(0, Arc('a', 'B', 0, merge ? 1 : 2)); + grammar.SetFinal(1, 0); + grammar.SetFinal(2, 0); + Matcher matcher(grammar); + Fst output = Input("sentinel"); + Check(!matcher.Match(Input("a"), &output), "exact merge/final tie falls back"); + Check(std::string(matcher.diagnostics().reason) == "ambiguous_path", "tie reason"); + Unchanged(output); + } + { + Fst grammar = Grammar(3); + grammar.AddArc(0, Arc('a', 'A', 0, 1)); + grammar.AddArc(0, Arc('a', 'B', 0.00001f, 2)); + grammar.SetFinal(1, 0); + grammar.SetFinal(2, 0); + Matcher matcher(grammar); + Fst output = Input("sentinel"); + Check(!matcher.Match(Input("a"), &output), "near tie falls back"); + Unchanged(output); + } + { + Fst grammar = Grammar(3); + grammar.AddArc(0, Arc(0, 'A', 0, 1)); + grammar.AddArc(1, Arc(0, 'B', -1, 0)); + grammar.AddArc(1, Arc('a', 'C', 0, 2)); + grammar.SetFinal(2, 0); + Matcher matcher(grammar); + Check(!matcher.valid(), "input epsilon cycle rejected at construction"); + Fst output = Input("sentinel"); + Check(!matcher.Match(Input("a"), &output), "cycle never enters relaxation"); + Unchanged(output); + } + { + Fst grammar = Grammar(2); + grammar.AddArc(0, Arc('a', 'A', 0, 1)); + grammar.SetFinal(1, 0); + Matcher::Limits limits; + limits.arc_relaxations = 0; + Matcher matcher(grammar, limits); + Fst output = Input("sentinel"); + Check(!matcher.Match(Input("a"), &output), "work budget falls back"); + Unchanged(output); + limits.arc_relaxations = 10; + limits.backpointers = 0; + Matcher no_memory(grammar, limits); + Check(!no_memory.Match(Input("a"), &output), "backpointer budget falls back"); + Unchanged(output); + Matcher normal(grammar); + Fst weighted = Input("a"); + weighted.SetFinal(1, 1); + Check(!normal.Match(weighted, &output), "weighted/preprocessed input rejected"); + Fst branching = Input("a"); + branching.AddArc(0, Arc('a', 'a', 0, 1)); + Check(!normal.Match(branching, &output), "nonlinear input rejected"); + Fst nul = Input(std::string(1, '\0')); + Check(!normal.Match(nul, &output), "NUL epsilon input rejected"); + Unchanged(output); + limits.input_bytes = 8; + Matcher short_limit(grammar, limits); + Fst cyclic = Grammar(1); + cyclic.AddArc(0, Arc('a', 'a', 0, 0)); + Check(!short_limit.Match(cyclic, &output), "cyclic input bounded"); + } + std::cout << "PASS " << checks << " synthetic CPU checks\n"; +} diff --git a/native/tts-normalizer/tests/text_preprocessor_c_test.cc b/native/tts-normalizer/tests/text_preprocessor_c_test.cc new file mode 100644 index 0000000..96b5711 --- /dev/null +++ b/native/tts-normalizer/tests/text_preprocessor_c_test.cc @@ -0,0 +1,246 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + */ + +#include + +#include +#include + +#include "processor_abi.h" + +namespace { + +NemotronSpeechByteSpanV1 +Span(const std::string& value) +{ + return {reinterpret_cast(value.data()), value.size()}; +} + +NemotronSpeechRecordViewV1 +Utf8(const std::string& name, const std::string& value) +{ + NemotronSpeechRecordViewV1 record{}; + record.struct_size = sizeof(record); + record.name = Span(name); + record.kind = NEMOTRON_SPEECH_RECORD_UTF8_V1; + record.element_type = NEMOTRON_SPEECH_ELEMENT_U8_V1; + record.data = Span(value); + return record; +} + +NemotronSpeechRecordViewV1 +Boolean(const std::string& name, const std::string& value) +{ + NemotronSpeechRecordViewV1 record{}; + record.struct_size = sizeof(record); + record.name = Span(name); + record.kind = NEMOTRON_SPEECH_RECORD_BOOLEAN_V1; + record.element_type = NEMOTRON_SPEECH_ELEMENT_U8_V1; + record.data = Span(value); + return record; +} + +TEST(TextPreprocessorC, NegotiatesAbiAndProcessesText) +{ + EXPECT_EQ(nemotron_speech_processor_get_api(999), nullptr); + const auto* api = + nemotron_speech_processor_get_api(NEMOTRON_SPEECH_PROCESSOR_ABI_VERSION); + ASSERT_NE(api, nullptr); + EXPECT_EQ(api->abi_version, NEMOTRON_SPEECH_PROCESSOR_ABI_VERSION); + + NemotronSpeechProcessorCapabilitiesV1 capabilities{}; + capabilities.struct_size = sizeof(capabilities); + auto status = api->capabilities(&capabilities); + ASSERT_EQ(status.code, 0); + api->release_status(&status); + EXPECT_EQ(capabilities.thread_safe, 0); + EXPECT_EQ(capabilities.streaming_input, 1); + EXPECT_EQ(capabilities.stateful, 0); + + const std::string language_name = "language_code"; + const std::string language = "en-US"; + const auto config_record = Utf8(language_name, language); + NemotronSpeechEnvelopeViewV1 config{}; + config.struct_size = sizeof(config); + config.schema_version = 1; + config.records = &config_record; + config.record_count = 1; + config.is_final = 1; + + NemotronSpeechProcessorHandleV1 handle = nullptr; + status = api->create(&config, &handle); + ASSERT_EQ(status.code, 0); + api->release_status(&status); + ASSERT_NE(handle, nullptr); + + const std::string text_name = "text"; + const std::string text = "Hello world"; + const auto text_record = Utf8(text_name, text); + const auto language_record = Utf8(language_name, language); + const NemotronSpeechRecordViewV1 request_records[] = {text_record, language_record}; + const std::string request_id = "request-1"; + NemotronSpeechEnvelopeViewV1 request{}; + request.struct_size = sizeof(request); + request.schema_version = 1; + request.request_id = Span(request_id); + request.records = request_records; + request.record_count = 2; + request.is_final = 1; + + NemotronSpeechEnvelopeOutputV1 output{}; + output.struct_size = sizeof(output); + status = api->process(handle, &request, &output); + ASSERT_EQ(status.code, 0); + api->release_status(&status); + ASSERT_EQ(output.record_count, 1); + EXPECT_EQ( + std::string( + reinterpret_cast(output.records[0].data.data), output.records[0].data.len), + "Hello world."); + EXPECT_EQ(output.records[0].is_final, 1); + + api->release_output(handle, &output); + api->destroy(handle); +} + +TEST(TextPreprocessorC, ContinuousRequestReturnsOneUnpunctuatedCanonicalSegment) +{ + const auto* api = + nemotron_speech_processor_get_api(NEMOTRON_SPEECH_PROCESSOR_ABI_VERSION); + ASSERT_NE(api, nullptr); + + const std::string language_name = "language_code"; + const std::string language = "en-US"; + const auto config_record = Utf8(language_name, language); + NemotronSpeechEnvelopeViewV1 config{}; + config.struct_size = sizeof(config); + config.schema_version = 1; + config.records = &config_record; + config.record_count = 1; + + NemotronSpeechProcessorHandleV1 handle = nullptr; + auto status = api->create(&config, &handle); + ASSERT_EQ(status.code, 0); + api->release_status(&status); + + const std::string text_name = "text"; + const std::string text = " \tFirst unit. Second unit \n"; + const std::string continuous_name = "continuous_streaming"; + const std::string continuous_value(1, '\1'); + const NemotronSpeechRecordViewV1 request_records[] = { + Utf8(text_name, text), Utf8(language_name, language), + Boolean(continuous_name, continuous_value)}; + NemotronSpeechEnvelopeViewV1 request{}; + request.struct_size = sizeof(request); + request.schema_version = 1; + request.records = request_records; + request.record_count = 3; + + NemotronSpeechEnvelopeOutputV1 output{}; + output.struct_size = sizeof(output); + status = api->process(handle, &request, &output); + ASSERT_EQ(status.code, 0); + api->release_status(&status); + ASSERT_EQ(output.record_count, 1); + EXPECT_EQ( + std::string( + reinterpret_cast(output.records[0].data.data), output.records[0].data.len), + " First unit. Second unit "); + EXPECT_EQ(output.records[0].sequence, 0); + EXPECT_EQ(output.records[0].is_final, 1); + + api->release_output(handle, &output); + api->destroy(handle); +} + +TEST(TextPreprocessorC, DirectGrammarConfigurationLoadsFarPair) +{ + const auto* api = + nemotron_speech_processor_get_api(NEMOTRON_SPEECH_PROCESSOR_ABI_VERSION); + ASSERT_NE(api, nullptr); + + const std::string language_name = "language_code"; + const std::string language = "en-US"; + const std::string path_name = "normalizer_path"; + const std::string path = "testdata/en-US"; + const std::string direct_name = "direct_grammar"; + const std::string true_value(1, '\1'); + const NemotronSpeechRecordViewV1 config_records[] = { + Utf8(language_name, language), Utf8(path_name, path), Boolean(direct_name, true_value)}; + NemotronSpeechEnvelopeViewV1 config{}; + config.struct_size = sizeof(config); + config.schema_version = 1; + config.records = config_records; + config.record_count = 3; + + NemotronSpeechProcessorHandleV1 handle = nullptr; + auto status = api->create(&config, &handle); + ASSERT_EQ(status.code, 0); + api->release_status(&status); + + const std::string text_name = "text"; + const std::string text = "I spent $2 today"; + const std::string continuous_name = "continuous_streaming"; + const NemotronSpeechRecordViewV1 request_records[] = { + Utf8(text_name, text), Utf8(language_name, language), + Boolean(continuous_name, true_value)}; + NemotronSpeechEnvelopeViewV1 request{}; + request.struct_size = sizeof(request); + request.schema_version = 1; + request.records = request_records; + request.record_count = 3; + + NemotronSpeechEnvelopeOutputV1 output{}; + output.struct_size = sizeof(output); + status = api->process(handle, &request, &output); + ASSERT_EQ(status.code, 0); + api->release_status(&status); + ASSERT_EQ(output.record_count, 1); + EXPECT_EQ( + std::string( + reinterpret_cast(output.records[0].data.data), output.records[0].data.len), + "I spent two dollars today"); + + api->release_output(handle, &output); + api->destroy(handle); +} + +TEST(TextPreprocessorC, ConvertsExceptionsToStatus) +{ + const auto* api = + nemotron_speech_processor_get_api(NEMOTRON_SPEECH_PROCESSOR_ABI_VERSION); + const std::string language_name = "language_code"; + const std::string language = "en-US"; + const auto config_record = Utf8(language_name, language); + NemotronSpeechEnvelopeViewV1 config{}; + config.struct_size = sizeof(config); + config.schema_version = 1; + config.records = &config_record; + config.record_count = 1; + + NemotronSpeechProcessorHandleV1 handle = nullptr; + auto status = api->create(&config, &handle); + ASSERT_EQ(status.code, 0); + api->release_status(&status); + + const std::string text_name = "text"; + const std::string empty; + const auto text_record = Utf8(text_name, empty); + const auto language_record = Utf8(language_name, language); + const NemotronSpeechRecordViewV1 records[] = {text_record, language_record}; + NemotronSpeechEnvelopeViewV1 request{}; + request.struct_size = sizeof(request); + request.schema_version = 1; + request.records = records; + request.record_count = 2; + NemotronSpeechEnvelopeOutputV1 output{}; + output.struct_size = sizeof(output); + status = api->process(handle, &request, &output); + EXPECT_NE(status.code, 0); + EXPECT_GT(status.message.len, 0); + api->release_status(&status); + api->destroy(handle); +} + +} // namespace diff --git a/native/tts-normalizer/tests/text_preprocessor_test.cc b/native/tts-normalizer/tests/text_preprocessor_test.cc new file mode 100644 index 0000000..80bb372 --- /dev/null +++ b/native/tts-normalizer/tests/text_preprocessor_test.cc @@ -0,0 +1,236 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + */ + +#include "text_preprocessor.h" + +#include + +#include +#include +#include + +namespace nemotron_speech::tts::preprocessing { +namespace { + +TextPreprocessor +MakePlainPreprocessor() +{ + TextPreprocessorConfig config; + config.languages.emplace("en-US", LanguageConfig{}); + return TextPreprocessor(config); +} + +PreprocessOptions +ContinuousOptions() +{ + PreprocessOptions options; + options.continuous_streaming = true; + return options; +} + +TEST(TextPreprocessorTest, NormalizesWhitespaceAndAppendsPunctuation) +{ + auto preprocessor = MakePlainPreprocessor(); + const auto segments = preprocessor.Preprocess(" Hello \n\t portable world ", "en-US"); + ASSERT_EQ(segments.size(), 1); + EXPECT_EQ(segments[0].text, "Hello portable world."); + EXPECT_EQ(segments[0].sentence_num, 0); + EXPECT_TRUE(segments[0].is_last_sentence); +} + +TEST(TextPreprocessorTest, PreservesMultilingualTerminalPunctuation) +{ + auto preprocessor = MakePlainPreprocessor(); + const std::vector terminators = { + ".", "!", "?", ";", ":", "؟", "۔", "।", "॥", "。", "！", "？"}; + + for (const auto& terminator : terminators) { + const std::string input = "Spoken text" + terminator; + const auto segments = preprocessor.Preprocess(input, "en-US"); + ASSERT_EQ(segments.size(), 1) << "terminator: " << terminator; + EXPECT_EQ(segments[0].text, input) << "terminator: " << terminator; + } +} + +TEST(TextPreprocessorTest, DetectsTerminalPunctuationBeforeSentenceClosers) +{ + auto preprocessor = MakePlainPreprocessor(); + const std::vector> cases = { + {"Spoken text.\"", "Spoken text.\""}, + {"Spoken text?']", "Spoken text?']"}, + {"Spoken text。”", "Spoken text。”"}, + {"Spoken text؟’", "Spoken text؟’"}, + {"\"Spoken text\"", "\"Spoken text\"."}, + }; + + for (const auto& [input, expected] : cases) { + const auto segments = preprocessor.Preprocess(input, "en-US"); + ASSERT_EQ(segments.size(), 1) << "input: " << input; + EXPECT_EQ(segments[0].text, expected) << "input: " << input; + } +} + +TEST(TextPreprocessorTest, PreservesUnicodeTerminatorAfterTrailingWhitespaceIsCollapsed) +{ + auto preprocessor = MakePlainPreprocessor(); + const auto segments = preprocessor.Preprocess("Spoken text। \n\t", "en-US"); + + ASSERT_EQ(segments.size(), 1); + EXPECT_EQ(segments[0].text, "Spoken text।"); +} + +TEST(TextPreprocessorTest, ContinuousModeReturnsOneSegmentWithoutAppendingPunctuation) +{ + auto preprocessor = MakePlainPreprocessor(); + const auto segments = preprocessor.Preprocess( + "First unit. Second unit without a terminator", "en-US", ContinuousOptions()); + + ASSERT_EQ(segments.size(), 1); + EXPECT_EQ(segments[0].text, "First unit. Second unit without a terminator"); + EXPECT_EQ(segments[0].sentence_num, 0); + EXPECT_TRUE(segments[0].is_last_sentence); +} + +TEST(TextPreprocessorTest, ContinuousModeKeepsWfstAndAbbreviationsButSkipsSentenceSplitting) +{ + TextPreprocessorConfig config; + LanguageConfig language; + language.normalizer_path = "testdata/en-US"; + language.abbreviations_path = "testdata/en-US/abbr.txt"; + config.languages.emplace("en-US", std::move(language)); + TextPreprocessor preprocessor(config); + + const auto segments = preprocessor.Preprocess( + "Dr Jones spent $2 today. Second unit", "en-US", ContinuousOptions()); + + ASSERT_EQ(segments.size(), 1); + EXPECT_EQ(segments[0].text, "doctor Jones spent two dollars today. Second unit"); +} + +TEST(TextPreprocessorTest, DirectGrammarModeMatchesLegacyNonChineseNormalizerLoading) +{ + TextPreprocessorConfig config; + LanguageConfig language; + language.normalizer_path = "testdata/en-US"; + language.direct_grammar = true; + config.languages.emplace("en-US", std::move(language)); + TextPreprocessor preprocessor(config); + + const auto segments = + preprocessor.Preprocess("I spent $2 today", "en-US", ContinuousOptions()); + + ASSERT_EQ(segments.size(), 1); + EXPECT_EQ(segments[0].text, "I spent two dollars today"); +} + +TEST(TextPreprocessorTest, ContinuousModeCanonicalizesBoundaryWhitespace) +{ + auto preprocessor = MakePlainPreprocessor(); + const auto segments = preprocessor.Preprocess( + " \tHello\n\tcontinuous world \r\n", "en-US", ContinuousOptions()); + + ASSERT_EQ(segments.size(), 1); + EXPECT_EQ(segments[0].text, " Hello continuous world "); +} + +TEST(TextPreprocessorTest, LegacyModeRemainsSentenceOrientedAndPunctuated) +{ + auto preprocessor = MakePlainPreprocessor(); + const auto segments = preprocessor.Preprocess(" Legacy behavior ", "en-US"); + + ASSERT_EQ(segments.size(), 1); + EXPECT_EQ(segments[0].text, "Legacy behavior."); +} + +TEST(TextPreprocessorTest, ExpandsSupportedSubSsml) +{ + auto preprocessor = MakePlainPreprocessor(); + const auto segments = preprocessor.Preprocess( + "Visit the WWW", "en-US"); + ASSERT_EQ(segments.size(), 1); + EXPECT_EQ(segments[0].text, "Visit the World Wide Web."); +} + +TEST(TextPreprocessorTest, TreatsLeadingLessThanTextAsPlainText) +{ + auto preprocessor = MakePlainPreprocessor(); + const std::vector inputs = {"< 3 is spoken", "( + preprocessor.Preprocess("Hello", "en-US")); + FAIL() << "Expected TextPreprocessorException"; + } + catch (const TextPreprocessorException& error) { + EXPECT_EQ(error.code(), PreprocessorErrorCode::kUnsupportedSsml); + } +} + +TEST(TextPreprocessorTest, RejectsUnsupportedLanguage) +{ + auto preprocessor = MakePlainPreprocessor(); + EXPECT_THROW(preprocessor.Preprocess("Hello", "fr-FR"), TextPreprocessorException); +} + +TEST(TextPreprocessorTest, RejectsPunctuationOnlyInput) +{ + auto preprocessor = MakePlainPreprocessor(); + EXPECT_THROW(preprocessor.Preprocess("...?!", "en-US"), TextPreprocessorException); +} + +TEST(TextPreprocessorTest, RejectsInvalidUtf8) +{ + auto preprocessor = MakePlainPreprocessor(); + const std::string invalid("\xc3\x28", 2); + EXPECT_THROW(preprocessor.Preprocess(invalid, "en-US"), TextPreprocessorException); +} + +TEST(TextPreprocessorTest, ContinuousModeRejectsInvalidUtf8) +{ + auto preprocessor = MakePlainPreprocessor(); + const std::string invalid("\xc3\x28", 2); + EXPECT_THROW( + preprocessor.Preprocess(invalid, "en-US", ContinuousOptions()), TextPreprocessorException); +} + +TEST(TextPreprocessorTest, EnforcesConfiguredLimits) +{ + TextPreprocessorConfig config; + config.max_input_bytes = 4; + config.languages.emplace("en-US", LanguageConfig{}); + TextPreprocessor preprocessor(config); + EXPECT_THROW(preprocessor.Preprocess("hello", "en-US"), TextPreprocessorException); +} + +TEST(TextPreprocessorTest, ContinuousModeEnforcesInputAndNormalizedLimits) +{ + TextPreprocessorConfig input_limited_config; + input_limited_config.max_input_bytes = 4; + input_limited_config.languages.emplace("en-US", LanguageConfig{}); + TextPreprocessor input_limited(input_limited_config); + EXPECT_THROW( + input_limited.Preprocess("hello", "en-US", ContinuousOptions()), + TextPreprocessorException); + + TextPreprocessorConfig output_limited_config; + output_limited_config.max_normalized_bytes = 4; + output_limited_config.languages.emplace("en-US", LanguageConfig{}); + TextPreprocessor output_limited(output_limited_config); + EXPECT_THROW( + output_limited.Preprocess("hello", "en-US", ContinuousOptions()), + TextPreprocessorException); +} + +} // namespace +} // namespace nemotron_speech::tts::preprocessing diff --git a/native/tts-normalizer/third_party/BUILD.bazel b/native/tts-normalizer/third_party/BUILD.bazel new file mode 100644 index 0000000..ef737a8 --- /dev/null +++ b/native/tts-normalizer/third_party/BUILD.bazel @@ -0,0 +1,3 @@ +# Repository BUILD definitions referenced by WORKSPACE.bazel live beside this file. + +exports_files(["sparrowhawk_normalizer.cc", "sparrowhawk_rule_system.cc", "sparrowhawk_string_matcher.h"]) diff --git a/native/tts-normalizer/third_party/BUILD.openfst b/native/tts-normalizer/third_party/BUILD.openfst new file mode 100644 index 0000000..b8207e2 --- /dev/null +++ b/native/tts-normalizer/third_party/BUILD.openfst @@ -0,0 +1,1212 @@ +""" +Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +""" + +package( + default_visibility = ["//visibility:public"], +) + +COPTS = select({ + "@bazel_tools//src/conditions:windows": ["/wd4388"], + "//conditions:default": ["-Wno-sign-compare"], +}) + +DEFINES = select({ + "@bazel_tools//src/conditions:windows": ["FST_NO_DYNAMIC_LINKING"], + "//conditions:default": [], +}) + +config_setting( + name = "has_absl", + values = {"define": "absl=1"}, +) + +cc_library( + name = "openfst", + srcs = glob([ + "src/lib/*.cc", + "src/lib/**/*.cc", + "src/lib/**/*.inc", + ]), + hdrs = glob([ + "src/include/*.h", + "src/include/**/*.h", + "src/include/**/*.inc", + ]), + copts = ["-O3 -Wno-unused-local-typedefs -Wno-sign-compare -Wno-unused-variable -Wno-reorder -Wno-return-type"], + includes = ["src/include"], + deps = [ + "@com_github_gflags_gflags//:gflags", + "@glog//:glog" + ] +) + +licenses(["notice"]) # Apache 2.0 + +exports_files(["COPYING"]) + +prefix_dir = "src/" + +static_binary = 1 + +# Core library (lib/) + +cc_library( + name = "fst-decl", + hdrs = [prefix_dir + "include/fst/fst-decl.h"], + includes = [prefix_dir + "include"], + deps = [":base"], +) + +PUBLIC_HEADERS = [ + # One-stop header file which includes the remaining headers below: + prefix_dir + "include/fst/fstlib.h", + + # Fine-grained headers: + prefix_dir + "include/fst/accumulator.h", + prefix_dir + "include/fst/add-on.h", + prefix_dir + "include/fst/arc-arena.h", + prefix_dir + "include/fst/arc-map.h", + prefix_dir + "include/fst/arc.h", + prefix_dir + "include/fst/arcfilter.h", + prefix_dir + "include/fst/arcsort.h", + prefix_dir + "include/fst/bi-table.h", + prefix_dir + "include/fst/cache.h", + prefix_dir + "include/fst/cc-visitors.h", + prefix_dir + "include/fst/closure.h", + prefix_dir + "include/fst/compact-fst.h", + prefix_dir + "include/fst/complement.h", + prefix_dir + "include/fst/compose-filter.h", + prefix_dir + "include/fst/compose.h", + prefix_dir + "include/fst/concat.h", + prefix_dir + "include/fst/connect.h", + prefix_dir + "include/fst/const-fst.h", + prefix_dir + "include/fst/determinize.h", + prefix_dir + "include/fst/dfs-visit.h", + prefix_dir + "include/fst/difference.h", + prefix_dir + "include/fst/disambiguate.h", + prefix_dir + "include/fst/edit-fst.h", + prefix_dir + "include/fst/encode.h", + prefix_dir + "include/fst/epsnormalize.h", + prefix_dir + "include/fst/equal.h", + prefix_dir + "include/fst/equivalent.h", + prefix_dir + "include/fst/expanded-fst.h", + prefix_dir + "include/fst/factor-weight.h", + prefix_dir + "include/fst/filter-state.h", + prefix_dir + "include/fst/fst.h", + prefix_dir + "include/fst/heap.h", + prefix_dir + "include/fst/impl-to-fst.h", + prefix_dir + "include/fst/intersect.h", + prefix_dir + "include/fst/invert.h", + prefix_dir + "include/fst/isomorphic.h", + prefix_dir + "include/fst/label-reachable.h", + prefix_dir + "include/fst/lookahead-filter.h", + prefix_dir + "include/fst/lookahead-matcher.h", + prefix_dir + "include/fst/matcher-fst.h", + prefix_dir + "include/fst/matcher.h", + prefix_dir + "include/fst/memory.h", + prefix_dir + "include/fst/minimize.h", + prefix_dir + "include/fst/mutable-fst.h", + prefix_dir + "include/fst/partition.h", + prefix_dir + "include/fst/project.h", + prefix_dir + "include/fst/properties.h", + prefix_dir + "include/fst/prune.h", + prefix_dir + "include/fst/push.h", + prefix_dir + "include/fst/queue.h", + prefix_dir + "include/fst/randequivalent.h", + prefix_dir + "include/fst/randgen.h", + prefix_dir + "include/fst/rational.h", + prefix_dir + "include/fst/relabel.h", + prefix_dir + "include/fst/replace-util.h", + prefix_dir + "include/fst/replace.h", + prefix_dir + "include/fst/reverse.h", + prefix_dir + "include/fst/reweight.h", + prefix_dir + "include/fst/rmepsilon.h", + prefix_dir + "include/fst/rmfinalepsilon.h", + prefix_dir + "include/fst/shortest-distance.h", + prefix_dir + "include/fst/shortest-path.h", + prefix_dir + "include/fst/state-map.h", + prefix_dir + "include/fst/state-reachable.h", + prefix_dir + "include/fst/state-table.h", + prefix_dir + "include/fst/statesort.h", + prefix_dir + "include/fst/string.h", + prefix_dir + "include/fst/symbol-table-ops.h", + prefix_dir + "include/fst/synchronize.h", + prefix_dir + "include/fst/test-properties.h", + prefix_dir + "include/fst/topsort.h", + prefix_dir + "include/fst/union.h", + prefix_dir + "include/fst/vector-fst.h", + prefix_dir + "include/fst/verify.h", + prefix_dir + "include/fst/visit.h", + prefix_dir + "include/fst/flags.h", +] + +# This version does not have the export-dynamic flag set and should not be +# used to load dynamic-shared object FST extensions. Please see the +# "lib_export_dynamic" target below for binaries that need DSO loading. +cc_library( + name = "lib_lite", + srcs = [ + prefix_dir + "lib/encode.cc", + prefix_dir + "lib/fst.cc", + prefix_dir + "lib/properties.cc", + prefix_dir + "lib/symbol-table-ops.cc", + ], + hdrs = PUBLIC_HEADERS, + copts = COPTS, + defines = DEFINES, + includes = [prefix_dir + "include"], + linkopts = ["-lm"], + deps = [ + ":base", + ":fst-decl", + ":icu", + ":interval-set", + ":register", + ":symbol-table", + ":union-find", + ":util", + ":weight", + ], +) + +cc_library( + name = "fst", + hdrs = PUBLIC_HEADERS, + defines = DEFINES, + includes = [prefix_dir + "include"], + deps = [ + ":fst-types", + ":lib_lite", + ], +) + +cc_library( + name = "lib_export_dynamic", + linkopts = ["-Wl,--export-dynamic"], + deps = [":fst"], +) + +cc_library( + name = "fst-types", + srcs = [prefix_dir + "lib/fst-types.cc"], + defines = DEFINES, + deps = [":lib_lite"], + alwayslink = 1, # because of registration +) + +cc_library( + name = "symbol-table", + srcs = [prefix_dir + "lib/symbol-table.cc"], + hdrs = [prefix_dir + "include/fst/symbol-table.h"], + copts = COPTS, + defines = DEFINES, + includes = [prefix_dir + "include"], + deps = [ + ":base", + ":util", + ], +) + +cc_library( + name = "weight", + srcs = [prefix_dir + "lib/weight.cc"], + hdrs = [ + prefix_dir + "include/fst/error-weight.h", + prefix_dir + "include/fst/expectation-weight.h", + prefix_dir + "include/fst/float-weight.h", + prefix_dir + "include/fst/lexicographic-weight.h", + prefix_dir + "include/fst/pair-weight.h", + prefix_dir + "include/fst/power-weight.h", + prefix_dir + "include/fst/product-weight.h", + prefix_dir + "include/fst/set-weight.h", + prefix_dir + "include/fst/signed-log-weight.h", + prefix_dir + "include/fst/sparse-power-weight.h", + prefix_dir + "include/fst/sparse-tuple-weight.h", + prefix_dir + "include/fst/string-weight.h", + prefix_dir + "include/fst/tuple-weight.h", + prefix_dir + "include/fst/union-weight.h", + prefix_dir + "include/fst/weight.h", + ], + defines = DEFINES, + includes = [prefix_dir + "include"], + linkopts = ["-lm"], + deps = [ + ":base", + ":util", + ], +) + +cc_library( + name = "interval-set", + hdrs = [prefix_dir + "include/fst/interval-set.h"], + includes = [prefix_dir + "include"], + deps = [ + ":base", + ":util", + ], +) + +cc_library( + name = "register", + hdrs = [ + prefix_dir + "include/fst/generic-register.h", + prefix_dir + "include/fst/register.h", + ], + includes = [prefix_dir + "include"], + linkopts = ["-ldl"], + deps = [ + ":base", + ":util", + ], +) + +cc_library( + name = "icu", + hdrs = [ + prefix_dir + "include/fst/icu.h", + ], +) + +cc_library( + name = "union-find", + hdrs = [prefix_dir + "include/fst/union-find.h"], + includes = [prefix_dir + "include"], + deps = [":base"], +) + +cc_library( + name = "util", + srcs = [ + prefix_dir + "lib/mapped-file.cc", + prefix_dir + "lib/util.cc", + ], + hdrs = [ + prefix_dir + "include/fst/mapped-file.h", + prefix_dir + "include/fst/util.h", + ], + includes = [prefix_dir + "include"], + deps = [":base"], +) + +cc_library( + name = "base", + srcs = [ + prefix_dir + "lib/compat.cc", + prefix_dir + "lib/flags.cc", + ], + hdrs = [ + prefix_dir + "include/fst/compat.h", + prefix_dir + "include/fst/flags.h", + prefix_dir + "include/fst/icu.h", + prefix_dir + "include/fst/lock.h", + prefix_dir + "include/fst/log.h", + prefix_dir + "include/fst/windows_defs.inc", + prefix_dir + "include/fst/types.h", + ], + defines = select({ + ":has_absl": ["OPENFST_HAS_ABSL=1"], + "//conditions:default": [], + }), + includes = [prefix_dir + "include"], + deps = select({ + ":has_absl": ["@com_google_absl//absl/synchronization"], + "//conditions:default": [], + }), +) + +# Core library tests (test/) + +cc_test( + name = "fst_test", + timeout = "short", + srcs = [ + prefix_dir + "test/fst_test.cc", + prefix_dir + "include/fst/test/compactors.h", + prefix_dir + "include/fst/test/fst_test.h", + ], + includes = [prefix_dir + "include"], + deps = [":fst"], +) + +cc_library( + name = "weight-tester", + testonly = 1, + hdrs = [prefix_dir + "include/fst/test/weight-tester.h"], + includes = [prefix_dir], + deps = [":weight"], +) + +cc_test( + name = "weight_test", + timeout = "short", + srcs = [prefix_dir + "test/weight_test.cc"], + copts = COPTS + select({ + "@bazel_tools//src/conditions:windows": [], + "//conditions:default": ["-Wno-unused-local-typedefs"], + }), + deps = [ + ":fst", + ":weight-tester", + ], +) + +[ + cc_test( + name = "algo_test_%s" % weight, + srcs = [ + prefix_dir + "test/algo_test.cc", + prefix_dir + "include/fst/test/algo_test.h", + prefix_dir + "include/fst/test/rand-fst.h", + ], + defines = ["TEST_%s" % weight.upper()], + deps = [":fst"], + ) + for weight in [ + "tropical", + "log", + "minmax", + "lexicographic", + "power", + ] +] + +# Non-template scripting-language integration (script/) + +cc_library( + name = "fstscript_base", + srcs = [ + prefix_dir + "script/arciterator-class.cc", + prefix_dir + "script/encodemapper-class.cc", + prefix_dir + "script/fst-class.cc", + prefix_dir + "script/getters.cc", + prefix_dir + "script/stateiterator-class.cc", + prefix_dir + "script/text-io.cc", + prefix_dir + "script/weight-class.cc", + ], + hdrs = [ + prefix_dir + "include/fst/script/arc-class.h", + prefix_dir + "include/fst/script/arcfilter-impl.h", + prefix_dir + "include/fst/script/arciterator-class.h", + prefix_dir + "include/fst/script/arg-packs.h", + prefix_dir + "include/fst/script/encodemapper-class.h", + prefix_dir + "include/fst/script/fst-class.h", + prefix_dir + "include/fst/script/fstscript.h", + prefix_dir + "include/fst/script/getters.h", + prefix_dir + "include/fst/script/script-impl.h", + prefix_dir + "include/fst/script/stateiterator-class.h", + prefix_dir + "include/fst/script/text-io.h", + prefix_dir + "include/fst/script/weight-class.h", + # + prefix_dir + "include/fst/script/arcsort.h", + #prefix_dir + "include/fst/script/cc-visitors.h", + prefix_dir + "include/fst/script/closure.h", + prefix_dir + "include/fst/script/compile.h", + prefix_dir + "include/fst/script/compile-impl.h", + prefix_dir + "include/fst/script/compose.h", + prefix_dir + "include/fst/script/concat.h", + prefix_dir + "include/fst/script/connect.h", + prefix_dir + "include/fst/script/convert.h", + prefix_dir + "include/fst/script/decode.h", + prefix_dir + "include/fst/script/determinize.h", + prefix_dir + "include/fst/script/difference.h", + prefix_dir + "include/fst/script/disambiguate.h", + prefix_dir + "include/fst/script/draw.h", + prefix_dir + "include/fst/script/draw-impl.h", + prefix_dir + "include/fst/script/encode.h", + prefix_dir + "include/fst/script/epsnormalize.h", + prefix_dir + "include/fst/script/equal.h", + prefix_dir + "include/fst/script/equivalent.h", + prefix_dir + "include/fst/script/info.h", + prefix_dir + "include/fst/script/info-impl.h", + prefix_dir + "include/fst/script/intersect.h", + prefix_dir + "include/fst/script/invert.h", + prefix_dir + "include/fst/script/isomorphic.h", + prefix_dir + "include/fst/script/map.h", + prefix_dir + "include/fst/script/minimize.h", + prefix_dir + "include/fst/script/print.h", + prefix_dir + "include/fst/script/print-impl.h", + prefix_dir + "include/fst/script/project.h", + prefix_dir + "include/fst/script/prune.h", + prefix_dir + "include/fst/script/push.h", + prefix_dir + "include/fst/script/randequivalent.h", + prefix_dir + "include/fst/script/randgen.h", + prefix_dir + "include/fst/script/relabel.h", + prefix_dir + "include/fst/script/replace.h", + prefix_dir + "include/fst/script/reverse.h", + prefix_dir + "include/fst/script/reweight.h", + prefix_dir + "include/fst/script/rmepsilon.h", + prefix_dir + "include/fst/script/shortest-distance.h", + prefix_dir + "include/fst/script/shortest-path.h", + prefix_dir + "include/fst/script/synchronize.h", + prefix_dir + "include/fst/script/topsort.h", + prefix_dir + "include/fst/script/union.h", + prefix_dir + "include/fst/script/verify.h", + ], + copts = COPTS, + includes = [prefix_dir + "include"], + deps = [":fst"], +) + +[ + cc_library( + name = "fstscript_%s" % operation, + srcs = [prefix_dir + "script/%s.cc" % operation], + hdrs = [prefix_dir + "include/fst/script/%s.h" % operation], + includes = [prefix_dir + "include"], + deps = [":fstscript_base"], + ) + for operation in [ + "arcsort", + "closure", + "concat", + "connect", + "convert", + "decode", + "determinize", + "disambiguate", + "encode", + "epsnormalize", + "equal", + "equivalent", + "invert", + "isomorphic", + "map", + "minimize", + "project", + "prune", + "push", + "randgen", + "relabel", + "replace", + "reverse", + "reweight", + "synchronize", + "topsort", + "union", + "verify", + ] +] + +cc_library( + name = "fstscript_compile", + srcs = [prefix_dir + "script/compile.cc"], + hdrs = [ + prefix_dir + "include/fst/script/compile.h", + prefix_dir + "include/fst/script/compile-impl.h", + ], + includes = [prefix_dir + "include"], + deps = [":fstscript_base"], +) + +cc_library( + name = "fstscript_compose", + srcs = [prefix_dir + "script/compose.cc"], + hdrs = [prefix_dir + "include/fst/script/compose.h"], + includes = [prefix_dir + "include"], + deps = [ + ":fstscript_base", + ":fstscript_connect", + ], +) + +cc_library( + name = "fstscript_difference", + srcs = [prefix_dir + "script/difference.cc"], + hdrs = [prefix_dir + "include/fst/script/difference.h"], + includes = [prefix_dir + "include"], + deps = [ + ":fstscript_base", + ":fstscript_compose", + ], +) + +cc_library( + name = "fstscript_draw", + srcs = [prefix_dir + "script/draw.cc"], + hdrs = [ + prefix_dir + "include/fst/script/draw.h", + prefix_dir + "include/fst/script/draw-impl.h", + ], + includes = [prefix_dir + "include"], + deps = [":fstscript_base"], +) + +cc_library( + name = "fstscript_info", + srcs = [ + prefix_dir + "script/info.cc", + prefix_dir + "script/info-impl.cc", + ], + hdrs = [ + prefix_dir + "include/fst/script/info.h", + prefix_dir + "include/fst/script/info-impl.h", + ], + includes = [prefix_dir + "include"], + deps = [":fstscript_base"], +) + +cc_library( + name = "fstscript_intersect", + srcs = [prefix_dir + "script/intersect.cc"], + hdrs = [prefix_dir + "include/fst/script/intersect.h"], + includes = [prefix_dir + "include"], + deps = [ + ":fstscript_base", + ":fstscript_compose", + ], +) + +cc_library( + name = "fstscript_print", + srcs = [prefix_dir + "script/print.cc"], + hdrs = [ + prefix_dir + "include/fst/script/print.h", + prefix_dir + "include/fst/script/print-impl.h", + ], + includes = [prefix_dir + "include"], + deps = [":fstscript_base"], +) + +cc_library( + name = "fstscript_randequivalent", + srcs = [prefix_dir + "script/randequivalent.cc"], + hdrs = [prefix_dir + "include/fst/script/randequivalent.h"], + includes = [prefix_dir + "include"], + deps = [ + ":fstscript_base", + ":fstscript_randgen", + ], +) + +cc_library( + name = "fstscript_rmepsilon", + srcs = [prefix_dir + "script/rmepsilon.cc"], + hdrs = [prefix_dir + "include/fst/script/rmepsilon.h"], + includes = [prefix_dir + "include"], + deps = [ + ":fstscript_base", + ":fstscript_shortest_distance", + ], +) + +cc_library( + name = "fstscript_shortest_distance", + srcs = [prefix_dir + "script/shortest-distance.cc"], + hdrs = [prefix_dir + "include/fst/script/shortest-distance.h"], + includes = [prefix_dir + "include"], + deps = [ + ":fstscript_base", + ":fstscript_prune", + ], +) + +cc_library( + name = "fstscript_shortest_path", + srcs = [prefix_dir + "script/shortest-path.cc"], + hdrs = [prefix_dir + "include/fst/script/shortest-path.h"], + includes = [prefix_dir + "include"], + deps = [ + ":fstscript_base", + ":fstscript_shortest_distance", + ], +) + +cc_library( + name = "fstscript", + deps = [ + ":fstscript_arcsort", + ":fstscript_closure", + ":fstscript_compile", + ":fstscript_compose", + ":fstscript_concat", + ":fstscript_connect", + ":fstscript_convert", + ":fstscript_decode", + ":fstscript_determinize", + ":fstscript_difference", + ":fstscript_disambiguate", + ":fstscript_draw", + ":fstscript_encode", + ":fstscript_epsnormalize", + ":fstscript_equal", + ":fstscript_equivalent", + ":fstscript_info", + ":fstscript_intersect", + ":fstscript_invert", + ":fstscript_isomorphic", + ":fstscript_map", + ":fstscript_minimize", + ":fstscript_print", + ":fstscript_project", + ":fstscript_prune", + ":fstscript_push", + ":fstscript_randequivalent", + ":fstscript_randgen", + ":fstscript_relabel", + ":fstscript_replace", + ":fstscript_reverse", + ":fstscript_reweight", + ":fstscript_rmepsilon", + ":fstscript_shortest_distance", + ":fstscript_shortest_path", + ":fstscript_synchronize", + ":fstscript_topsort", + ":fstscript_union", + ":fstscript_verify", + ], +) + +# Command-line binaries (bin/) + +[ + cc_binary( + name = "fst%s" % operation.replace("_", ""), + srcs = [ + prefix_dir + "bin/fst%s.cc" % operation.replace("_", ""), + prefix_dir + "bin/fst%s-main.cc" % operation.replace("_", ""), + ], + linkstatic = static_binary, + deps = [":fstscript_%s" % operation], + ) + for operation in [ + "arcsort", + "closure", + "compile", + "compose", + "concat", + "connect", + "convert", + "determinize", + "difference", + "disambiguate", + "draw", + "epsnormalize", + "equal", + "info", + "intersect", + "invert", + "isomorphic", + "map", + "minimize", + "print", + "project", + "prune", + "push", + "randgen", + "relabel", + "replace", + "reverse", + "reweight", + "rmepsilon", + "shortest_distance", + "shortest_path", + "synchronize", + "topsort", + "union", + ] +] + +cc_binary( + name = "fstencode", + srcs = [ + prefix_dir + "bin/fstencode.cc", + prefix_dir + "bin/fstencode-main.cc", + ], + linkstatic = static_binary, + deps = [ + ":fstscript_decode", + ":fstscript_encode", + ], +) + +cc_binary( + name = "fstequivalent", + srcs = [ + prefix_dir + "bin/fstequivalent.cc", + prefix_dir + "bin/fstequivalent-main.cc", + ], + linkstatic = static_binary, + deps = [ + ":fstscript_equivalent", + ":fstscript_randequivalent", + ], +) + +cc_binary( + name = "fstsymbols", + srcs = [ + prefix_dir + "bin/fstsymbols.cc", + prefix_dir + "bin/fstsymbols-main.cc", + ], + linkstatic = static_binary, + deps = [":fstscript_verify"], +) + +# Extension: Fst ARchive a/k/a FAR (extensions/far/) + +cc_library( + name = "sttable", + srcs = [ + prefix_dir + "extensions/far/stlist.cc", + prefix_dir + "extensions/far/sttable.cc", + ], + hdrs = [ + prefix_dir + "include/fst/extensions/far/stlist.h", + prefix_dir + "include/fst/extensions/far/sttable.h", + ], + includes = [prefix_dir + "include"], + deps = [":util"], +) + +cc_library( + name = "far_base", + hdrs = glob([ + prefix_dir + "include/fst/extensions/far/far.h", + prefix_dir + "include/fst/types.h", + prefix_dir + "include/fst/log.h", + prefix_dir + "include/fst/flags.h", + prefix_dir + "include/fst/lock.h", + prefix_dir + "include/fst/script/*.h", + prefix_dir + "include/fst/script/arcfilter-impl.h" + ]), + includes = [prefix_dir + "include"], + deps = [ + ":fst", + ":sttable", + ], +) + +cc_library( + name = "far", + hdrs = [ + prefix_dir + "include/fst/extensions/far/compile-strings.h", + prefix_dir + "include/fst/extensions/far/convert.h", + prefix_dir + "include/fst/extensions/far/create.h", + prefix_dir + "include/fst/extensions/far/encode.h", + prefix_dir + "include/fst/extensions/far/equal.h", + prefix_dir + "include/fst/extensions/far/extract.h", + prefix_dir + "include/fst/extensions/far/farlib.h", + prefix_dir + "include/fst/extensions/far/info.h", + prefix_dir + "include/fst/extensions/far/isomorphic.h", + prefix_dir + "include/fst/extensions/far/map-reduce.h", + prefix_dir + "include/fst/extensions/far/print-strings.h", + ], + includes = [prefix_dir + "include"], + deps = [ + ":far_base", + ":fst", + ], +) + +cc_library( + name = "farscript", + srcs = [ + prefix_dir + "extensions/far/compile-strings.cc", + prefix_dir + "extensions/far/far-class.cc", + prefix_dir + "extensions/far/farscript.cc", + prefix_dir + "extensions/far/getters.cc", + prefix_dir + "extensions/far/script-impl.cc", + ], + hdrs = [ + prefix_dir + "include/fst/extensions/far/compile-strings.h", + prefix_dir + "include/fst/extensions/far/far-class.h", + prefix_dir + "include/fst/extensions/far/farscript.h", + prefix_dir + "include/fst/extensions/far/getters.h", + prefix_dir + "include/fst/extensions/far/script-impl.h", + ], + includes = [prefix_dir + "include"], + deps = [ + ":base", + ":far", + ":fstscript_base", + ], +) + +[ + cc_binary( + name = "far%s" % operation, + srcs = [ + prefix_dir + "extensions/far/far%s.cc" % operation, + prefix_dir + "extensions/far/far%s-main.cc" % operation, + ], + linkstatic = static_binary, + deps = [":farscript"], + ) + for operation in [ + "compilestrings", + "convert", + "create", + "encode", + "equal", + "extract", + "info", + "isomorphic", + "printstrings", + ] +] + +# Extension: PushDown Transducers a/k/a PDT (extensions/pdt/) + +cc_library( + name = "pdt", + hdrs = [ + prefix_dir + "include/fst/extensions/pdt/collection.h", + prefix_dir + "include/fst/extensions/pdt/compose.h", + prefix_dir + "include/fst/extensions/pdt/expand.h", + prefix_dir + "include/fst/extensions/pdt/getters.h", + prefix_dir + "include/fst/extensions/pdt/info.h", + prefix_dir + "include/fst/extensions/pdt/paren.h", + prefix_dir + "include/fst/extensions/pdt/pdt.h", + prefix_dir + "include/fst/extensions/pdt/pdtlib.h", + prefix_dir + "include/fst/extensions/pdt/replace.h", + prefix_dir + "include/fst/extensions/pdt/reverse.h", + prefix_dir + "include/fst/extensions/pdt/shortest-path.h", + ], + includes = [prefix_dir + "include"], + deps = [ + ":fst", + ], +) + +cc_library( + name = "pdtscript", + srcs = [ + prefix_dir + "extensions/pdt/getters.cc", + prefix_dir + "extensions/pdt/pdtscript.cc", + ], + hdrs = [ + prefix_dir + "include/fst/extensions/pdt/getters.h", + prefix_dir + "include/fst/extensions/pdt/pdtscript.h", + ], + includes = [prefix_dir + "include"], + deps = [ + ":fst", + ":fstscript_base", + ":pdt", + ], +) + +cc_binary( + name = "pdtcompose", + srcs = [ + prefix_dir + "extensions/pdt/pdtcompose.cc", + prefix_dir + "extensions/pdt/pdtcompose-main.cc", + ], + linkstatic = static_binary, + deps = [ + ":fstscript_connect", + ":pdtscript", + ], +) + +[ + cc_binary( + name = "pdt%s" % operation, + srcs = [ + prefix_dir + "extensions/pdt/pdt%s.cc" % operation, + prefix_dir + "extensions/pdt/pdt%s-main.cc" % operation, + ], + linkstatic = static_binary, + deps = [":pdtscript"], + ) + for operation in [ + "expand", + "info", + "replace", + "reverse", + "shortestpath", + ] +] + +# Extension: Multi PushDown Transducers a/k/a MPDT (extensions/mpdt/) + +cc_library( + name = "mpdt", + hdrs = [ + prefix_dir + "include/fst/extensions/mpdt/compose.h", + prefix_dir + "include/fst/extensions/mpdt/expand.h", + prefix_dir + "include/fst/extensions/mpdt/info.h", + prefix_dir + "include/fst/extensions/mpdt/mpdt.h", + prefix_dir + "include/fst/extensions/mpdt/mpdtlib.h", + prefix_dir + "include/fst/extensions/mpdt/read_write_utils.h", + prefix_dir + "include/fst/extensions/mpdt/reverse.h", + ], + includes = [prefix_dir + "include"], + deps = [ + ":fst", + ":pdt", + ], +) + +cc_library( + name = "mpdtscript", + srcs = [prefix_dir + "extensions/mpdt/mpdtscript.cc"], + hdrs = [prefix_dir + "include/fst/extensions/mpdt/mpdtscript.h"], + includes = [prefix_dir + "include"], + deps = [ + ":fst", + ":fstscript_base", + ":mpdt", + ":pdtscript", + ], +) + +cc_binary( + name = "mpdtcompose", + srcs = [ + prefix_dir + "extensions/mpdt/mpdtcompose.cc", + prefix_dir + "extensions/mpdt/mpdtcompose-main.cc", + ], + linkstatic = static_binary, + deps = [ + ":fstscript_connect", + ":mpdtscript", + ], +) + +[ + cc_binary( + name = "mpdt%s" % operation, + srcs = [ + prefix_dir + "extensions/mpdt/mpdt%s.cc" % operation, + prefix_dir + "extensions/mpdt/mpdt%s-main.cc" % operation, + ], + linkstatic = static_binary, + deps = [":mpdtscript"], + ) + for operation in [ + "expand", + "info", + "reverse", + ] +] + +# Extension: LOUDS compressed n-gram language models (extensions/ngram/) + +cc_library( + name = "ngram", + srcs = [ + prefix_dir + "extensions/ngram/bitmap-index.cc", + prefix_dir + "extensions/ngram/ngram-fst.cc", + prefix_dir + "extensions/ngram/nthbit.cc", + ], + hdrs = [ + prefix_dir + "include/fst/extensions/ngram/bitmap-index.h", + prefix_dir + "include/fst/extensions/ngram/ngram-fst.h", + prefix_dir + "include/fst/extensions/ngram/nthbit.h", + ], + includes = [prefix_dir + "include"], + deps = [ + ":fst", + ], +) + +# Extension: Compact FSTs and FSAs (extensions/compact) + +[ + cc_library( + name = "%s-fst" % type, + srcs = [ + prefix_dir + "extensions/compact/%s-fst.cc" % type, + ], + linkstatic = 1, + deps = [":fst"], + ) + for type in [ + "compact16_acceptor", + "compact16_string", + "compact16_unweighted", + "compact16_unweighted_acceptor", + "compact16_weighted_string", + "compact64_acceptor", + "compact64_string", + "compact64_unweighted", + "compact64_unweighted_acceptor", + "compact64_weighted_string", + "compact8_acceptor", + "compact8_string", + "compact8_unweighted", + "compact8_unweighted_acceptor", + "compact8_weighted_string", + ] +] + +# Extension: Compressed FSTs (extensions/compress) + +cc_library( + name = "compress-fst", + hdrs = [ + prefix_dir + "include/fst/extensions/compress/compress.h", + prefix_dir + "include/fst/extensions/compress/elias.h", + ], + includes = [prefix_dir + "include"], + deps = [ + ":fst", + ], +) + +cc_library( + name = "compressscript", + srcs = [ + prefix_dir + "extensions/compress/compressscript.cc", + ], + hdrs = [ + prefix_dir + "include/fst/extensions/compress/compressscript.h", + ], + includes = [prefix_dir + "include"], + deps = [ + ":base", + ":compress-fst", + ":fstscript_base", + ], +) + +cc_binary( + name = "fstcompress", + srcs = [ + prefix_dir + "extensions/compress/fstcompress.cc", + prefix_dir + "extensions/compress/fstcompress-main.cc", + ], + linkstatic = static_binary, + deps = [":compressscript"], +) + +# Extension: Constant FSTs (extensions/const) + +[ + cc_library( + name = "const%s-fst" % type, + srcs = [ + prefix_dir + "extensions/const/const%s-fst.cc" % type, + ], + linkstatic = 1, + deps = [":fst"], + ) + for type in [ + "8", + "16", + "64", + ] +] + +# Extension: Linear FSTs (extensions/linear) + +cc_library( + name = "linear-fst", + srcs = [ + prefix_dir + "extensions/linear/linear-classifier-fst.cc", + prefix_dir + "extensions/linear/linear-tagger-fst.cc", + ], + hdrs = [ + prefix_dir + "include/fst/extensions/linear/linear-fst-data-builder.h", + prefix_dir + "include/fst/extensions/linear/linear-fst-data.h", + prefix_dir + "include/fst/extensions/linear/linear-fst.h", + prefix_dir + "include/fst/extensions/linear/loglinear-apply.h", + prefix_dir + "include/fst/extensions/linear/trie.h", + ], + linkstatic = 1, + deps = [ + ":fst", + ":pdt", + ], +) + +cc_library( + name = "linearscript", + srcs = [ + prefix_dir + "extensions/linear/linearscript.cc", + ], + hdrs = [ + prefix_dir + "include/fst/extensions/linear/linearscript.h", + ], + includes = [prefix_dir + "include"], + deps = [ + ":base", + ":fstscript_base", + ":linear-fst", + ":pdt", + ], +) + +[ + cc_binary( + name = "fst%s" % name, + srcs = [ + prefix_dir + "extensions/linear/fst%s.cc" % name, + prefix_dir + "extensions/linear/fst%s-main.cc" % name, + ], + linkstatic = static_binary, + deps = [":linearscript"], + ) + for name in [ + "linear", + "loglinearapply", + ] +] + +# Extension: Lookahead FSTs (extensions/lookahead) + +[ + cc_library( + name = "%s_lookahead-fst" % type, + srcs = [ + prefix_dir + "extensions/lookahead/%s_lookahead-fst.cc" % type, + ], + linkstatic = 1, + deps = [":fst"], + ) + for type in [ + "arc", + "ilabel", + "olabel", + ] +] + +# Extension: Special FSTs (extensions/special) + +[ + cc_library( + name = "%s-fst" % type, + srcs = [ + prefix_dir + "extensions/special/%s-fst.cc" % type, + ], + hdrs = [ + prefix_dir + "include/fst/extensions/special/%s-fst.h" % type, + ], + linkstatic = 1, + deps = [":fst"], + ) + for type in [ + "phi", + "rho", + "sigma", + ] +] + +cc_binary( + name = "fstspecial", + srcs = [ + prefix_dir + "bin/fstconvert.cc", + prefix_dir + "bin/fstconvert-main.cc", + ], + linkstatic = static_binary, + deps = [ + ":fstscript_convert", + ":phi-fst", + ":rho-fst", + ":sigma-fst", + ], +) + +# Miscellaneous helper Build rules. +exports_files([ + prefix_dir + "extensions/far/build_defs.bzl", +]) + +# TODO: Extensions: python diff --git a/native/tts-normalizer/third_party/BUILD.pugixml b/native/tts-normalizer/third_party/BUILD.pugixml new file mode 100644 index 0000000..8ba8824 --- /dev/null +++ b/native/tts-normalizer/third_party/BUILD.pugixml @@ -0,0 +1,24 @@ +""" +Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +""" + +package( + default_visibility = ["//visibility:public"], +) + +cc_library( + name = "pugixml", + srcs = [ + "src/pugixml.cpp" + ], + hdrs = [ + "src/pugixml.hpp", "src/pugiconfig.hpp" + ], + strip_include_prefix = "src", + visibility = ["//visibility:public"] +) diff --git a/native/tts-normalizer/third_party/BUILD.re2 b/native/tts-normalizer/third_party/BUILD.re2 new file mode 100644 index 0000000..14f27e2 --- /dev/null +++ b/native/tts-normalizer/third_party/BUILD.re2 @@ -0,0 +1,282 @@ +# Copyright 2009 The RE2 Authors. All Rights Reserved. +# Use of this source code is governed by a BSD-style +# license that can be found in the LICENSE file. + +# Bazel (http://bazel.io/) BUILD file for RE2. + +licenses(["notice"]) + +exports_files(["LICENSE"]) + +config_setting( + name = "macos", + values = {"cpu": "darwin"}, +) + +config_setting( + name = "wasm", + values = {"cpu": "wasm32"}, +) + +config_setting( + name = "windows", + values = {"cpu": "x64_windows"}, +) + +load("@rules_cc//cc:defs.bzl", "cc_binary", "cc_library", "cc_test") + +cc_library( + name = "re2", + srcs = [ + "re2/bitmap256.h", + "re2/bitstate.cc", + "re2/compile.cc", + "re2/dfa.cc", + "re2/filtered_re2.cc", + "re2/mimics_pcre.cc", + "re2/nfa.cc", + "re2/onepass.cc", + "re2/parse.cc", + "re2/perl_groups.cc", + "re2/pod_array.h", + "re2/prefilter.cc", + "re2/prefilter.h", + "re2/prefilter_tree.cc", + "re2/prefilter_tree.h", + "re2/prog.cc", + "re2/prog.h", + "re2/re2.cc", + "re2/regexp.cc", + "re2/regexp.h", + "re2/set.cc", + "re2/simplify.cc", + "re2/sparse_array.h", + "re2/sparse_set.h", + "re2/stringpiece.cc", + "re2/tostring.cc", + "re2/unicode_casefold.cc", + "re2/unicode_casefold.h", + "re2/unicode_groups.cc", + "re2/unicode_groups.h", + "re2/walker-inl.h", + "util/logging.h", + "util/mix.h", + "util/mutex.h", + "util/rune.cc", + "util/strutil.cc", + "util/strutil.h", + "util/utf.h", + "util/util.h", + ], + hdrs = [ + "re2/filtered_re2.h", + "re2/re2.h", + "re2/set.h", + "re2/stringpiece.h", + ], + copts = select({ + ":wasm": [], + ":windows": [], + "//conditions:default": ["-pthread"], + }), + includes = ["."], + linkopts = select({ + # macOS doesn't need `-pthread' when linking and it appears that + # older versions of Clang will warn about the unused command line + # argument, so just don't pass it. + ":macos": [], + ":wasm": [], + ":windows": [], + "//conditions:default": ["-pthread"], + }), + visibility = ["//visibility:public"], + linkstatic = 1, +) + +cc_library( + name = "testing", + testonly = 1, + srcs = [ + "re2/testing/backtrack.cc", + "re2/testing/dump.cc", + "re2/testing/exhaustive_tester.cc", + "re2/testing/null_walker.cc", + "re2/testing/regexp_generator.cc", + "re2/testing/string_generator.cc", + "re2/testing/tester.cc", + "util/pcre.cc", + ], + hdrs = [ + "re2/testing/exhaustive_tester.h", + "re2/testing/regexp_generator.h", + "re2/testing/string_generator.h", + "re2/testing/tester.h", + "util/benchmark.h", + "util/flags.h", + "util/malloc_counter.h", + "util/pcre.h", + "util/test.h", + ], + deps = [":re2"], +) + +cc_library( + name = "test", + testonly = 1, + srcs = ["util/test.cc"], + deps = [":testing"], +) + +cc_test( + name = "charclass_test", + size = "small", + srcs = ["re2/testing/charclass_test.cc"], + deps = [":test"], +) + +cc_test( + name = "compile_test", + size = "small", + srcs = ["re2/testing/compile_test.cc"], + deps = [":test"], +) + +cc_test( + name = "filtered_re2_test", + size = "small", + srcs = ["re2/testing/filtered_re2_test.cc"], + deps = [":test"], +) + +cc_test( + name = "mimics_pcre_test", + size = "small", + srcs = ["re2/testing/mimics_pcre_test.cc"], + deps = [":test"], +) + +cc_test( + name = "parse_test", + size = "small", + srcs = ["re2/testing/parse_test.cc"], + deps = [":test"], +) + +cc_test( + name = "possible_match_test", + size = "small", + srcs = ["re2/testing/possible_match_test.cc"], + deps = [":test"], +) + +cc_test( + name = "re2_arg_test", + size = "small", + srcs = ["re2/testing/re2_arg_test.cc"], + deps = [":test"], +) + +cc_test( + name = "re2_test", + size = "small", + srcs = ["re2/testing/re2_test.cc"], + deps = [":test"], +) + +cc_test( + name = "regexp_test", + size = "small", + srcs = ["re2/testing/regexp_test.cc"], + deps = [":test"], +) + +cc_test( + name = "required_prefix_test", + size = "small", + srcs = ["re2/testing/required_prefix_test.cc"], + deps = [":test"], +) + +cc_test( + name = "search_test", + size = "small", + srcs = ["re2/testing/search_test.cc"], + deps = [":test"], +) + +cc_test( + name = "set_test", + size = "small", + srcs = ["re2/testing/set_test.cc"], + deps = [":test"], +) + +cc_test( + name = "simplify_test", + size = "small", + srcs = ["re2/testing/simplify_test.cc"], + deps = [":test"], +) + +cc_test( + name = "string_generator_test", + size = "small", + srcs = ["re2/testing/string_generator_test.cc"], + deps = [":test"], +) + +cc_test( + name = "dfa_test", + size = "large", + srcs = ["re2/testing/dfa_test.cc"], + deps = [":test"], +) + +cc_test( + name = "exhaustive1_test", + size = "large", + srcs = ["re2/testing/exhaustive1_test.cc"], + deps = [":test"], +) + +cc_test( + name = "exhaustive2_test", + size = "large", + srcs = ["re2/testing/exhaustive2_test.cc"], + deps = [":test"], +) + +cc_test( + name = "exhaustive3_test", + size = "large", + srcs = ["re2/testing/exhaustive3_test.cc"], + deps = [":test"], +) + +cc_test( + name = "exhaustive_test", + size = "large", + srcs = ["re2/testing/exhaustive_test.cc"], + deps = [":test"], +) + +cc_test( + name = "random_test", + size = "large", + srcs = ["re2/testing/random_test.cc"], + deps = [":test"], +) + +cc_library( + name = "benchmark", + testonly = 1, + srcs = ["util/benchmark.cc"], + deps = [":testing"], +) + +cc_binary( + name = "regexp_benchmark", + testonly = 1, + srcs = ["re2/testing/regexp_benchmark.cc"], + deps = [":benchmark"], +) diff --git a/native/tts-normalizer/third_party/BUILD.sparrowhawk_0_1 b/native/tts-normalizer/third_party/BUILD.sparrowhawk_0_1 new file mode 100644 index 0000000..49e87ca --- /dev/null +++ b/native/tts-normalizer/third_party/BUILD.sparrowhawk_0_1 @@ -0,0 +1,121 @@ +""" +Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +""" + +package( + default_visibility = ["//visibility:public"], +) + +load("@com_google_protobuf//:protobuf.bzl", "cc_proto_library") + +cc_proto_library( + name = "sparrowhawk_configuration_cc_proto", + srcs = ["src/proto/sparrowhawk_configuration.proto"], + include = "src/proto", + default_runtime = "@com_google_protobuf//:protobuf", + protoc = "@com_google_protobuf//:protoc" +) + +cc_proto_library( + name = "rule_order_cc_proto", + srcs = ["src/proto/rule_order.proto"], + include = "src/proto", + default_runtime = "@com_google_protobuf//:protobuf", + protoc = "@com_google_protobuf//:protoc" +) + +cc_proto_library( + name = "links_cc_proto", + srcs = ["src/proto/links.proto"], + include = "src/proto", + default_runtime = "@com_google_protobuf//:protobuf", + protoc = "@com_google_protobuf//:protoc" +) + +cc_proto_library( + name = "semiotic_classes_cc_proto", + srcs = ["src/proto/semiotic_classes.proto"], + include = "src/proto", + default_runtime = "@com_google_protobuf//:protobuf", + protoc = "@com_google_protobuf//:protoc", +) +cc_proto_library( + name = "items_cc_proto", + srcs = ["src/proto/items.proto"], + include = "src/proto", + default_runtime = "@com_google_protobuf//:protobuf", + protoc = "@com_google_protobuf//:protoc", + deps = [ + ":links_cc_proto", + ":semiotic_classes_cc_proto", + ], +) + +cc_proto_library( + name = "serialization_spec_cc_proto", + srcs = ["src/proto/serialization_spec.proto"], + include = "src/proto", + default_runtime = "@com_google_protobuf//:protobuf", + protoc = "@com_google_protobuf//:protoc", + deps = [ + ":links_cc_proto", + ":semiotic_classes_cc_proto", + ], +) + +genrule( + name = "make_compatibility_headers", + outs = [ + "sparrowhawk/items.pb.h", + "sparrowhawk/rule_order.pb.h", + "sparrowhawk/sparrowhawk_configuration.pb.h", + "sparrowhawk/serialization_spec.pb.h", + ], + cmd = """ + for f in $(OUTS); do + echo "#include " > $$f + done + """, +) + +cc_library( + name = "protos", + hdrs = [ + "sparrowhawk/items.pb.h", + "sparrowhawk/rule_order.pb.h", + "sparrowhawk/serialization_spec.pb.h", + "sparrowhawk/sparrowhawk_configuration.pb.h", + ], + includes = ["."], + deps = [ + ":items_cc_proto", + ":rule_order_cc_proto", + ":serialization_spec_cc_proto", + ":sparrowhawk_configuration_cc_proto", + ], +) + +cc_library( + name = "sparrowhawk", + srcs = glob ([ + "src/lib/*.cc", + ], exclude = ["src/lib/normalizer.cc", "src/lib/rule_system.cc"]) + ["@nemotron_tts_normalizer//third_party:sparrowhawk_normalizer.cc", "@nemotron_tts_normalizer//third_party:sparrowhawk_rule_system.cc"], + hdrs = glob ([ + "src/include/sparrowhawk/*.h", + ]) + ["@nemotron_tts_normalizer//third_party:sparrowhawk_string_matcher.h"], + includes = ["src/include/sparrowhawk", "src/include", "."], + copts = ["-funsigned-char", "-Wno-sign-compare", "-Wno-maybe-uninitialized", "-Wno-deprecated-declarations"], + deps = [ + "@openfst//:fst", + "@openfst//:far_base", + "@thrax_1_2_7//:thrax_1_2_7", + ":protos", + "@com_googlesource_code_re2//:re2", + ], + linkstatic=True, +) diff --git a/native/tts-normalizer/third_party/BUILD.thrax_1_2_7 b/native/tts-normalizer/third_party/BUILD.thrax_1_2_7 new file mode 100644 index 0000000..7697e37 --- /dev/null +++ b/native/tts-normalizer/third_party/BUILD.thrax_1_2_7 @@ -0,0 +1,167 @@ +""" +Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +""" + +package( + default_visibility = ["//visibility:public"], +) + +prefix_dir = "src/" + +# Separate flags library - compiled in isolation to use FST_FLAGS_* macros +# This library MUST NOT depend on glog or anything that brings in gflags + +cc_library( + name = "thrax_1_2_7", + srcs = [ + prefix_dir + "lib/ast/collection-node.cc", + prefix_dir + "lib/ast/fst-node.cc", + prefix_dir + "lib/ast/function-node.cc", + prefix_dir + "lib/ast/grammar-node.cc", + prefix_dir + "lib/ast/identifier-node.cc", + prefix_dir + "lib/ast/import-node.cc", + prefix_dir + "lib/ast/node.cc", + prefix_dir + "lib/ast/return-node.cc", + prefix_dir + "lib/ast/rule-node.cc", + prefix_dir + "lib/ast/statement-node.cc", + prefix_dir + "lib/ast/string-node.cc", + prefix_dir + "lib/flags/flags.cc", + prefix_dir + "lib/main/compiler-log.cc", + prefix_dir + "lib/main/compiler-log64.cc", + prefix_dir + "lib/main/compiler-stdarc.cc", + prefix_dir + "lib/main/grm-compiler.cc", + prefix_dir + "lib/main/lexer.cc", + prefix_dir + "lib/main/parser.cc", + prefix_dir + "lib/util/getters.cc", + prefix_dir + "lib/util/stringcompile.cc", + prefix_dir + "lib/util/stringfile.cc", + prefix_dir + "lib/util/stringutil.cc", + prefix_dir + "lib/util/utils.cc", + prefix_dir + "lib/walker/assert-empty.cc", + prefix_dir + "lib/walker/assert-equal.cc", + prefix_dir + "lib/walker/assert-null.cc", + prefix_dir + "lib/walker/cdrewrite.cc", + prefix_dir + "lib/walker/closure.cc", + prefix_dir + "lib/walker/compose.cc", + prefix_dir + "lib/walker/concat.cc", + prefix_dir + "lib/walker/determinize.cc", + prefix_dir + "lib/walker/difference.cc", + prefix_dir + "lib/walker/evaluator.cc", + prefix_dir + "lib/walker/expand.cc", + prefix_dir + "lib/walker/features.cc", + prefix_dir + "lib/walker/identifier-counter.cc", + prefix_dir + "lib/walker/invert.cc", + prefix_dir + "lib/walker/lenientlycompose.cc", + prefix_dir + "lib/walker/loadfst.cc", + prefix_dir + "lib/walker/loadfstfromfar.cc", + prefix_dir + "lib/walker/minimize.cc", + prefix_dir + "lib/walker/mpdtcompose.cc", + prefix_dir + "lib/walker/namespace.cc", + prefix_dir + "lib/walker/optimize.cc", + prefix_dir + "lib/walker/paradigm.cc", + prefix_dir + "lib/walker/pdtcompose.cc", + prefix_dir + "lib/walker/printer.cc", + prefix_dir + "lib/walker/project.cc", + prefix_dir + "lib/walker/replace.cc", + prefix_dir + "lib/walker/reverse.cc", + prefix_dir + "lib/walker/rewrite.cc", + prefix_dir + "lib/walker/rmepsilon.cc", + prefix_dir + "lib/walker/rmweight.cc", + prefix_dir + "lib/walker/stringfile.cc", + prefix_dir + "lib/walker/stringfst.cc", + prefix_dir + "lib/walker/symbols.cc", + prefix_dir + "lib/walker/symboltable.cc", + prefix_dir + "lib/walker/union.cc", + prefix_dir + "lib/walker/walker.cc", + ], + hdrs = [ + prefix_dir + "include/thrax/abstract-grm-manager.h", + prefix_dir + "include/thrax/algo/cdrewrite.h", + prefix_dir + "include/thrax/algo/checkprops.h", + prefix_dir + "include/thrax/algo/concatrange.h", + prefix_dir + "include/thrax/algo/cross.h", + prefix_dir + "include/thrax/algo/getters.h", + prefix_dir + "include/thrax/algo/lenientlycompose.h", + prefix_dir + "include/thrax/algo/optimize.h", + prefix_dir + "include/thrax/algo/paths.h", + prefix_dir + "include/thrax/algo/prefix_tree.h", + prefix_dir + "include/thrax/algo/stringcompile.h", + prefix_dir + "include/thrax/algo/stringfile.h", + prefix_dir + "include/thrax/algo/stringmap.h", + prefix_dir + "include/thrax/algo/stringprint.h", + prefix_dir + "include/thrax/algo/stringutil.h", + prefix_dir + "include/thrax/assert-empty.h", + prefix_dir + "include/thrax/assert-equal.h", + prefix_dir + "include/thrax/assert-null.h", + prefix_dir + "include/thrax/cdrewrite.h", + prefix_dir + "include/thrax/closure.h", + prefix_dir + "include/thrax/collection-node.h", + prefix_dir + "include/thrax/compat/compat.h", + prefix_dir + "include/thrax/compat/stlfunctions.h", + prefix_dir + "include/thrax/compat/utils.h", + prefix_dir + "include/thrax/compose.h", + prefix_dir + "include/thrax/compiler.h", + prefix_dir + "include/thrax/concat.h", + prefix_dir + "include/thrax/datatype.h", + prefix_dir + "include/thrax/determinize.h", + prefix_dir + "include/thrax/difference.h", + prefix_dir + "include/thrax/evaluator.h", + prefix_dir + "include/thrax/expand.h", + prefix_dir + "include/thrax/features.h", + prefix_dir + "include/thrax/fst-node.h", + prefix_dir + "include/thrax/function-node.h", + prefix_dir + "include/thrax/function.h", + prefix_dir + "include/thrax/grammar-node.h", + prefix_dir + "include/thrax/grm-compiler.h", + prefix_dir + "include/thrax/grm-manager.h", + prefix_dir + "include/thrax/identifier-counter.h", + prefix_dir + "include/thrax/identifier-node.h", + prefix_dir + "include/thrax/import-node.h", + prefix_dir + "include/thrax/invert.h", + prefix_dir + "include/thrax/lenientlycompose.h", + prefix_dir + "include/thrax/lexer.h", + prefix_dir + "include/thrax/loadfst.h", + prefix_dir + "include/thrax/loadfstfromfar.h", + prefix_dir + "include/thrax/make-parens-pair-vector.h", + prefix_dir + "include/thrax/minimize.h", + prefix_dir + "include/thrax/mpdtcompose.h", + prefix_dir + "include/thrax/namespace.h", + prefix_dir + "include/thrax/node.h", + prefix_dir + "include/thrax/optimize.h", + prefix_dir + "include/thrax/paradigm.h", + prefix_dir + "include/thrax/pdtcompose.h", + prefix_dir + "include/thrax/printer.h", + prefix_dir + "include/thrax/project.h", + prefix_dir + "include/thrax/replace.h", + prefix_dir + "include/thrax/register.h", + prefix_dir + "include/thrax/resource-map.h", + prefix_dir + "include/thrax/return-node.h", + prefix_dir + "include/thrax/reverse.h", + prefix_dir + "include/thrax/rewrite.h", + prefix_dir + "include/thrax/rmepsilon.h", + prefix_dir + "include/thrax/rmweight.h", + prefix_dir + "include/thrax/rule-node.h", + prefix_dir + "include/thrax/statement-node.h", + prefix_dir + "include/thrax/string-node.h", + prefix_dir + "include/thrax/stringfile.h", + prefix_dir + "include/thrax/stringfst.h", + prefix_dir + "include/thrax/symbols.h", + prefix_dir + "include/thrax/symboltable.h", + prefix_dir + "include/thrax/thrax.h", + prefix_dir + "include/thrax/union.h", + prefix_dir + "include/thrax/walker.h", + ], + includes = [prefix_dir + "include"], + deps = [ + "@openfst//:far", + "@openfst//:fst", + "@openfst//:fstscript", + "@openfst//:mpdt", + "@openfst//:pdt", + ], +) diff --git a/native/tts-normalizer/third_party/sparrowhawk_normalizer.cc b/native/tts-normalizer/third_party/sparrowhawk_normalizer.cc new file mode 100644 index 0000000..202792e --- /dev/null +++ b/native/tts-normalizer/third_party/sparrowhawk_normalizer.cc @@ -0,0 +1,437 @@ +// 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. +// +// Copyright 2015 and onwards Google, Inc. +#include + +#include +#include +#include +#include +#include +#include +using std::string; + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace speech { +namespace sparrowhawk { +extern thread_local bool plain_string_normalization; +namespace { +class PlainStringScope { + public: + explicit PlainStringScope(bool enabled) : previous_(plain_string_normalization) { + plain_string_normalization = enabled; + } + ~PlainStringScope() { plain_string_normalization = previous_; } + private: + bool previous_; +}; + +class StageTimer { + public: + StageTimer(const char* stage, const void* instance) : stage_(stage), instance_(instance) { + static const bool enabled = std::getenv("NEMOTRON_TN_PROFILE") != nullptr; + enabled_ = enabled; + if (enabled_) { + started_ = std::chrono::steady_clock::now(); + clock_gettime(CLOCK_THREAD_CPUTIME_ID, &cpu_started_); + } + } + ~StageTimer() { + if (!enabled_) return; + const auto ended = std::chrono::steady_clock::now(); + timespec cpu_ended; + clock_gettime(CLOCK_THREAD_CPUTIME_ID, &cpu_ended); + const double wall_us = std::chrono::duration(ended - started_).count(); + const double cpu_us = (cpu_ended.tv_sec - cpu_started_.tv_sec) * 1e6 + + (cpu_ended.tv_nsec - cpu_started_.tv_nsec) * 1e-3; + std::fprintf(stderr, "TN_STAGE %s instance=%p wall_us=%.3f cpu_us=%.3f\n", stage_, instance_, wall_us, cpu_us); + } + private: + const char* stage_; + const void* instance_; + bool enabled_; + std::chrono::steady_clock::time_point started_; + timespec cpu_started_; +}; +} + + +const char kDefaultSentenceBoundaryRegexp[] = "[\\.:!\\?] "; + +const std::string kDefaultVerbalizerProto = + "grammar_file: \"verbalize.far\"\ngrammar_name: \"Verbalizer\"\nrules { main: \"ALL\" redup: \"REDUP\" }\n"; +const std::string kDefaultTokenizerProto = + "grammar_file: \"tokenize_and_classify.far\"\ngrammar_name: \"TokenizerClassifier\"\nrules { main: \"TOKENIZE_AND_CLASSIFY\" }\n"; +const std::string kDefaultPostProcessorProto = + "grammar_file: \"post_process.far\"\ngrammar_name: \"PostProcessor\"\nrules { main: \"POSTPROCESSOR\" }\n"; +const std::string kDefaultPreProcessorProto = + "grammar_file: \"pre_process.far\"\ngrammar_name: \"PreProcessor\"\nrules { main: \"PREPROCESSOR\" }\n"; + + +std::vector kDefaultSentenceBoundaryExceptions = {"Mr.", "Dr.", "Mrs.", "St.", "Jan.", "Feb.", "Mar.", "Apr.", "Jun.", "Jul.", "Aug.", "Sep.", "Oct.", "Nov.", "Dec." }; + +Normalizer::Normalizer() { } + +Normalizer::~Normalizer() { } + + +bool Normalizer::Setup(const string &pathname, bool post_process, bool pre_process) +{ + SparrowhawkConfiguration configuration; + string proto_string = "tokenizer_grammar: \"tokenizer.ascii_proto\"\n" + "verbalizer_grammar: \"verbalizer.ascii_proto\"\n" + "sentence_boundary_regexp: \"[\\\\.:!\\\\?] \"\n" + "sentence_boundary_exceptions_file: \"sentence_boundary_exceptions.txt\"\n"; + if (pre_process) + proto_string += "preprocessor_grammar: \"preprocessor.ascii_proto\"\n"; + if (post_process) + proto_string += "postprocessor_grammar: \"postprocessor.ascii_proto\"\n"; + LOG(INFO) << "Proto String: " << proto_string; + try{ + if (!google::protobuf::TextFormat::ParseFromString(proto_string, &configuration)) + return false; + } + catch (...){ + LOG(ERROR) << "Failed to parse proto string" << proto_string; + } + + tokenizer_classifier_rules_.reset(new RuleSystem); + if (!tokenizer_classifier_rules_->LoadGrammarProtoFromString( + kDefaultTokenizerProto, + pathname)) + return false; + + verbalizer_rules_.reset(new RuleSystem); + if (!verbalizer_rules_->LoadGrammarProtoFromString(kDefaultVerbalizerProto, + pathname)) + return false; + + string sentence_boundary_regexp = kDefaultSentenceBoundaryRegexp; + sentence_boundary_.reset(new SentenceBoundary(sentence_boundary_regexp)); + sentence_boundary_->AddSentenceBoundaryExceptions(kDefaultSentenceBoundaryExceptions); + + if (configuration.has_preprocessor_grammar()) { + try{ + pre_processor_rules_.reset(new RuleSystem); + if (pre_processor_rules_->LoadGrammarProtoFromString( + kDefaultPreProcessorProto, + pathname)) + this->do_preprocess = true; + else + LoggerWarn("Unable to load pre_processor_grammar from: "); + } + catch (...){ + LOG(ERROR) << "Failed to load preprocessor" << proto_string; + } + } + if (configuration.has_postprocessor_grammar()) { + try{ + post_processor_rules_.reset(new RuleSystem); + if (post_processor_rules_->LoadGrammarProtoFromString( + kDefaultPostProcessorProto, + pathname)) + this->do_postprocess = true; + else + LoggerWarn("Unable to load post_processor_grammar"); + } + catch (...){ + LOG(ERROR) << "Failed to load postprocessor" << proto_string; + } + } + + return true; +} + +bool Normalizer::Setup(const string &configuration_proto, + const string &pathname_prefix) { + SparrowhawkConfiguration configuration; + string proto_string = IOStream::LoadFileToString(pathname_prefix + + "/" + configuration_proto); + if (!google::protobuf::TextFormat::ParseFromString(proto_string, &configuration)) + return false; + if (!(configuration.has_tokenizer_grammar())) + LoggerError("Configuration does not define a tokenizer-classifier grammar"); + if (!(configuration.has_verbalizer_grammar())) + LoggerError("Configuration does not define a verbalizer grammar"); + tokenizer_classifier_rules_.reset(new RuleSystem); + if (!tokenizer_classifier_rules_->LoadGrammar( + configuration.tokenizer_grammar(), + pathname_prefix)) + return false; + verbalizer_rules_.reset(new RuleSystem); + if (!verbalizer_rules_->LoadGrammar(configuration.verbalizer_grammar(), + pathname_prefix)) + return false; + string sentence_boundary_regexp; + if (configuration.has_sentence_boundary_regexp()) { + sentence_boundary_regexp = configuration.sentence_boundary_regexp(); + } else { + sentence_boundary_regexp = kDefaultSentenceBoundaryRegexp; + } + sentence_boundary_.reset(new SentenceBoundary(sentence_boundary_regexp)); + if (configuration.has_sentence_boundary_exceptions_file()) { + if (!sentence_boundary_->LoadSentenceBoundaryExceptions( + configuration.sentence_boundary_exceptions_file(), pathname_prefix)) { + LoggerError("Cannot load sentence boundary exceptions file: %s", + configuration.sentence_boundary_exceptions_file().c_str()); + } + } + if (configuration.has_serialization_spec()) { + string spec_string = IOStream::LoadFileToString( + pathname_prefix + "/" + configuration.serialization_spec()); + SerializeSpec spec; + if (spec_string.empty() || + !google::protobuf::TextFormat::ParseFromString(spec_string, &spec) || + (spec_serializer_ = Serializer::Create(spec)) == nullptr) { + LoggerError("Failed to load a valid serialization spec from file: %s", + configuration.serialization_spec().c_str()); + return false; + } + } + if (configuration.has_preprocessor_grammar()) { + pre_processor_rules_.reset(new RuleSystem); + if (pre_processor_rules_->LoadGrammar( + configuration.preprocessor_grammar(), + pathname_prefix)) + this->do_preprocess = true; + else + LoggerWarn("Unable to load pre_processor_grammar from: "); + } + if (configuration.has_postprocessor_grammar()) { + post_processor_rules_.reset(new RuleSystem); + if (post_processor_rules_->LoadGrammar( + configuration.postprocessor_grammar(), + pathname_prefix)) + this->do_postprocess = true; + else + LoggerWarn("Unable to load post_processor_grammar"); + } + return true; +} + + +bool Normalizer::Normalize( + const string &input, string *output, bool enable_preprocessing) const { + PlainStringScope scope(true); + std::unique_ptr utt; + utt.reset(new Utterance); + bool preprocessing_deleted_input = false; + if (!Normalize( + utt.get(), input, enable_preprocessing, + &preprocessing_deleted_input)) { + return false; + } + if (preprocessing_deleted_input) { + output->clear(); + return true; + } + *output = LinearizeWords(utt.get()); + + if (this->do_postprocess){ + const string pp_input = *output; + StageTimer timer("postprocess", this); + post_processor_rules_->ApplyRules(pp_input,output,false); + } + + + return true; +} + +bool Normalizer::Normalize( + Utterance *utt, const string &input, bool enable_preprocessing, + bool *preprocessing_deleted_input) const { + return TokenizeAndClassifyUtt( + utt, input, enable_preprocessing, + preprocessing_deleted_input) && + (*preprocessing_deleted_input || VerbalizeUtt(utt)); +} + +bool Normalizer::NormalizeAndShowLinks( + const string &input, string *output, bool enable_preprocessing) const { + PlainStringScope scope(false); + std::unique_ptr utt; + utt.reset(new Utterance); + bool preprocessing_deleted_input = false; + if (!Normalize( + utt.get(), input, enable_preprocessing, + &preprocessing_deleted_input)) { + return false; + } + if (preprocessing_deleted_input) { + output->clear(); + return true; + } + *output = ShowLinks(utt.get()); + return true; +} + +bool Normalizer::TokenizeAndClassifyUtt( + Utterance *utt, const string &input, bool enable_preprocessing, + bool *preprocessing_deleted_input) const { + StageTimer timer("classify", this); + typedef fst::StringCompiler Compiler; + typedef fst::StringPrinter Printer; + Compiler compiler(fst::TokenType::BYTE); + MutableTransducer input_fst; + if (!compiler(input, &input_fst)) { + LoggerError("Failed to compile input string \"%s\"", input.c_str()); + return false; + } + + *preprocessing_deleted_input = false; + MutableTransducer tokenizer_input(input_fst); + if (this->do_preprocess && enable_preprocessing) { + MutableTransducer preprocess_lattice; + if (!pre_processor_rules_->ApplyRules( + input_fst, &preprocess_lattice, false)) { + return false; + } + + + + + + fst::ShortestPath(preprocess_lattice, &tokenizer_input); + + MutableTransducer preprocessed_output(tokenizer_input); + fst::Project(&preprocessed_output, fst::PROJECT_OUTPUT); + fst::RmEpsilon(&preprocessed_output); + Printer printer(fst::TokenType::BYTE); + string preprocessed_text; + if (!printer(preprocessed_output, &preprocessed_text)) { + LoggerError("Failed to print preprocessed input"); + return false; + } + if (preprocessed_text.find_first_not_of(" \t\n\r\f\v") == string::npos) { + *preprocessing_deleted_input = true; + return true; + } + } + + MutableTransducer output; + if (!tokenizer_classifier_rules_->ApplyRules(tokenizer_input, + &output, + true)) { + LoggerError("Failed to tokenize \"%s\"", input.c_str()); + return false; + } + MutableTransducer shortest_path; + fst::ShortestPath(output, &shortest_path); + ProtobufParser parser(&shortest_path); + if (!parser.ParseTokensFromFST(utt, true )) { + LoggerError("Failed to parse tokens from FST for \"%s\"", input.c_str()); + return false; + } + return true; +} + + + + + + + + +bool Normalizer::VerbalizeUtt(Utterance *utt) const { + StageTimer timer("verbalize", this); + for (int i = 0; i < utt->linguistic().tokens_size(); ++i) { + Token *token = utt->mutable_linguistic()->mutable_tokens(i); + string token_form = ToString(*token); + token->set_first_daughter(-1); + token->set_last_daughter(-1); + + + + if (token->type() == Token::PUNCT) { + if (token->phrase_break() && + (utt->linguistic().words_size() == 0 || + utt->linguistic().words( + utt->linguistic().words_size() - 1).id() != "sil")) { + AddWord(utt, token, "sil"); + } + } else if (token->type() == Token::SEMIOTIC_CLASS) { + if (!token->skip()) { + LoggerDebug("Verbalizing: [%s]\n", token_form.c_str()); + string words; + if (VerbalizeSemioticClass(*token, &words)) { + AddWords(utt, token, words); + } else { + LoggerWarn("First-pass verbalization FAILED for [%s]", + token_form.c_str()); + + string original_token = token->name(); + token->Clear(); + token->set_name(original_token); + token->set_verbatim(original_token); + if (VerbalizeSemioticClass(*token, &words)) { + LoggerWarn("Reversion to verbatim succeeded for [%s]", + original_token.c_str()); + AddWords(utt, token, words); + } else { + + LoggerError("Verbalization FAILED for [%s]", token_form.c_str()); + } + } + } + } else if (token->type() == Token::WORD) { + if (token->has_wordid()) { + AddWord(utt, token, token->wordid()); + } else { + LoggerError("Token [%s] has type WORD but there is no word id", + token_form.c_str()); + } + } else { + LoggerError("No type found for [%s]", token_form.c_str()); + } + } + LoggerDebug("Verbalize output: Words\n%s\n\n", LinearizeWords(utt).c_str()); + return true; +} + +bool Normalizer::VerbalizeSemioticClass(const Token &markup, + string *words) const { + Token local(markup); + CleanFields(&local); + MutableTransducer input_fst; + if (spec_serializer_ == nullptr) { + ProtobufSerializer serializer(&local, &input_fst); + serializer.SerializeToFst(); + } else { + input_fst = spec_serializer_->Serialize(local); + } + if (!verbalizer_rules_->ApplyRules(input_fst, + words, + false )) { + LoggerError("Failed to verbalize \"%s\"", ToString(local).c_str()); + return false; + } + return true; +} + +std::vector Normalizer::SentenceSplitter(const string &input) const { + return sentence_boundary_->ExtractSentences(input); +} + +} +} diff --git a/native/tts-normalizer/third_party/sparrowhawk_rule_system.cc b/native/tts-normalizer/third_party/sparrowhawk_rule_system.cc new file mode 100644 index 0000000..6c9c64c --- /dev/null +++ b/native/tts-normalizer/third_party/sparrowhawk_rule_system.cc @@ -0,0 +1,347 @@ +// 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. +// +// Copyright 2015 and onwards Google, Inc. +#include +#include +#include +#include +#include +#include +#include "sparrowhawk_string_matcher.h" + +#include +#include +#include +#include +#include + +namespace speech { +namespace sparrowhawk { +thread_local bool plain_string_normalization = false; +namespace { +std::mutex matcher_mutex; +std::map>> string_matchers; + +StringMatcher* FindStringMatcher(const RuleSystem* owner, const string& name, + const Transducer& grammar) { + std::lock_guard lock(matcher_mutex); + auto& matcher = string_matchers[owner][name]; + if (!matcher) { + StringMatcher::Limits limits; + limits.canonical_word_outputs = true; + matcher.reset(new StringMatcher(grammar, limits)); + } + return matcher.get(); +} + +bool HasStrictWordSyntax(const MutableTransducer& path) { + string text; + auto state = path.Start(); + for (size_t steps = 0; steps <= path.NumStates(); ++steps) { + fst::ArcIterator arcs(path, state); + if (arcs.Done()) { + if (path.Final(state) == fst::StdArc::Weight::Zero()) return false; + break; + } + const auto arc = arcs.Value(); + arcs.Next(); + if (!arcs.Done() || arc.olabel < 0 || arc.olabel > 255) return false; + if (arc.olabel) text.push_back(static_cast(arc.olabel)); + state = arc.nextstate; + if (steps == path.NumStates()) return false; + } + size_t position = 0; + auto whitespace = [&]() { + while (position < text.size() && text[position] == ' ') ++position; + }; + auto literal = [&](const string& value) { + whitespace(); + if (text.compare(position, value.size(), value) != 0) return false; + position += value.size(); + return true; + }; + size_t tokens = 0; + while (true) { + whitespace(); + if (position == text.size()) return tokens != 0; + if (!literal("tokens") || !literal("{") || !literal("name") || + !literal(":") || !literal("\"")) return false; + bool closed = false; + while (position < text.size()) { + const auto byte = static_cast(text[position++]); + if (byte == '"') { closed = true; break; } + if (byte == '\\') { + if (position == text.size()) return false; + ++position; + } + } + if (!closed || !literal("}")) return false; + ++tokens; + } +} + +bool IsExplicitWordPath(const MutableTransducer& path) { + if (path.Start() == fst::kNoStateId || !HasStrictWordSyntax(path)) return false; + Utterance utterance; + ProtobufParser parser(&path); + if (!parser.ParseTokensFromFST(&utterance, true) || + utterance.linguistic().tokens_size() == 0) return false; + for (const auto& token : utterance.linguistic().tokens()) { + if (token.type() != Token::WORD || !token.has_wordid() || !token.has_name()) + return false; + std::vector fields; + token.GetReflection()->ListFields(token, &fields); + for (const auto* field : fields) { + const auto& name = field->name(); + if (name != "name" && name != "wordid" && name != "start_index" && + name != "end_index") return false; + } + } + return true; +} + +class StageTimer { + public: + StageTimer(const char* stage, const void* instance) : stage_(stage), instance_(instance) { + static const bool enabled = std::getenv("NEMOTRON_TN_PROFILE") != nullptr; + enabled_ = enabled; + if (enabled_) { + started_ = std::chrono::steady_clock::now(); + clock_gettime(CLOCK_THREAD_CPUTIME_ID, &cpu_started_); + } + } + ~StageTimer() { + if (!enabled_) return; + const auto ended = std::chrono::steady_clock::now(); + timespec cpu_ended; + clock_gettime(CLOCK_THREAD_CPUTIME_ID, &cpu_ended); + const double wall_us = std::chrono::duration(ended - started_).count(); + const double cpu_us = (cpu_ended.tv_sec - cpu_started_.tv_sec) * 1e6 + + (cpu_ended.tv_nsec - cpu_started_.tv_nsec) * 1e-3; + std::fprintf(stderr, "TN_STAGE %s instance=%p wall_us=%.3f cpu_us=%.3f\n", stage_, instance_, wall_us, cpu_us); + } + private: + const char* stage_; + const void* instance_; + bool enabled_; + std::chrono::steady_clock::time_point started_; + timespec cpu_started_; +}; +} + + + +using fst::LabelLookAheadRelabeler; +using fst::StdArc; + +RuleSystem::~RuleSystem() { + { + std::lock_guard lock(matcher_mutex); + string_matchers.erase(this); + } + std::map::iterator iter; + for (iter = lookaheads_.begin(); iter != lookaheads_.end(); iter++) { + delete iter->second; + } +} + +bool RuleSystem::LoadGrammar(const string& filename, const string& prefix) { + + std::string sep = ""; + if (prefix[prefix.size() - 1] != '/') { + sep = "/"; + } + + + string proto_string = IOStream::LoadFileToString(prefix + sep + filename); + return LoadGrammarProtoFromString(proto_string, prefix + sep); +} + +bool RuleSystem::LoadGrammarProtoFromString(const string& proto_string, const string& prefix) +{ + { + std::lock_guard lock(matcher_mutex); + string_matchers.erase(this); + } + std::string sep = ""; + if (prefix[prefix.size() - 1] != '/') { + sep = "/"; + } + + if (!google::protobuf::TextFormat::ParseFromString(proto_string, &grammar_)) + return false; + string grm_file = prefix + sep+ grammar_.grammar_file(); + grammar_name_ = grammar_.grammar_name(); + grm_.reset(new GrmManager); + if (!grm_->LoadArchive(grm_file)) { + LoggerError("Error loading archive \"%s\" from \"%s\"", + grammar_name_.c_str(), grm_file.c_str()); + return false; + } + + + for (int i = 0; i < grammar_.rules_size(); ++i) { + Rule rule = grammar_.rules(i); + if (grm_->GetFst(rule.main()) == NULL) { + LoggerError("Rule \"%s\" not found in \"%s\"", + rule.main().c_str(), grammar_name_.c_str()); + return false; + } + if (rule.has_parens() && grm_->GetFst(rule.parens()) == NULL) { + LoggerError("Rule \"%s\" not found in \"%s\"", + rule.parens().c_str(), grammar_name_.c_str()); + return false; + } + if (rule.has_redup() && grm_->GetFst(rule.redup()) == NULL) { + LoggerError("Rule \"%s\" not found in \"%s\"", + rule.redup().c_str(), grammar_name_.c_str()); + return false; + } + } + return true; +} + +bool RuleSystem::ApplyRules(const Transducer& input, + MutableTransducer* output, + bool use_lookahead) const { + StageTimer all_timer(grammar_name_.c_str(), this); + if (std::getenv("NEMOTRON_TN_STRING_MATCHER") != nullptr && plain_string_normalization && use_lookahead && + grammar_.rules_size() == 1 && !grammar_.rules(0).has_redup() && + !grammar_.rules(0).has_parens()) { + const auto& name = grammar_.rules(0).main(); + auto* matcher = FindStringMatcher(this, name, *grm_->GetFst(name)); + MutableTransducer candidate; + const bool matched = matcher->Match(input, &candidate); + const bool accepted = matched && IsExplicitWordPath(candidate); + if (std::getenv("NEMOTRON_TN_PROFILE") != nullptr) { + const auto& d = matcher->diagnostics(); + std::fprintf(stderr, "TN_MATCH reason=%s states=%zu arcs=%zu nodes=%zu best=%.9g second=%.9g guard=%.9g\n", + accepted ? "matched_word_only" : matched ? "non_word_path" : d.reason, + d.visited_states, d.arc_relaxations, d.backpointers, + d.best_cost, d.second_cost, d.ambiguity_guard); + } + if (accepted) { + *output = std::move(candidate); + return true; + } + } + MutableTransducer mutable_input(input); + for (int i = 0; i < grammar_.rules_size(); ++i) { + Rule rule = grammar_.rules(i); + if (rule.has_redup()) { + const string& redup_rule = rule.redup(); + MutableTransducer redup1; + + if (grm_->Rewrite(redup_rule, mutable_input, &redup1, "")) { + MutableTransducer redup2(redup1); + fst::Concat(redup1, &redup2); + fst::Union(&mutable_input, redup2); + fst::RmEpsilon(&mutable_input); + } + } + const string& rule_name = rule.main(); + string parens_rule = rule.has_parens() ? rule.parens() : ""; + + bool success = true; + if (parens_rule.empty() + && use_lookahead) { + std::map::iterator iter = + lookaheads_.find(rule_name); + LookaheadFst *lookahead_rule_fst; + if (iter == lookaheads_.end()) { + const Transducer *rule_fst = grm_->GetFst(rule_name); + lookahead_rule_fst = new LookaheadFst(*rule_fst); + lookaheads_[rule_name] = lookahead_rule_fst; + } else { + lookahead_rule_fst = iter->second; + } + LabelLookAheadRelabeler::Relabel(&mutable_input, + *lookahead_rule_fst, + false); + fst::ComposeFst tmp_output(mutable_input, + *lookahead_rule_fst); + *output = tmp_output; + if (output->NumStates() == 0) { + success = false; + } + + } else if (!grm_->Rewrite(rule_name, + mutable_input, + output, + parens_rule + ) + || output->NumStates() == 0) { + success = false; + } + if (!success) { + LoggerError("Application of rule \"%s\" failed", rule_name.c_str()); + return false; + } + if (i + 1 < grammar_.rules_size()) mutable_input = *output; + } + + + + { StageTimer epsilon_timer("rm_epsilon", this); fst::RmEpsilon(output); } + return true; +} + +typedef fst::StringCompiler Compiler; +typedef fst::StringPrinter Printer; + +bool RuleSystem::ApplyRules(const string& input, + string* output, + bool use_lookahead) const { + Compiler compiler(fst::TokenType::BYTE); + MutableTransducer input_fst, output_fst; + if (!compiler.operator()(input, &input_fst)) { + LoggerError("Failed to compile input string \"%s\"", input.c_str()); + return false; + } + if (!ApplyRules(input_fst, &output_fst, use_lookahead)) return false; + MutableTransducer shortest_path; + fst::ShortestPath(output_fst, &shortest_path); + fst::Project(&shortest_path, fst::PROJECT_OUTPUT); + fst::RmEpsilon(&shortest_path); + Printer printer(fst::TokenType::BYTE); + if (!printer.operator()(shortest_path, output)) { + LoggerError("Failed to print output string"); + return false; + } + return true; +} + +bool RuleSystem::ApplyRules(const Transducer& input, + string* output, + bool use_lookahead) const { + MutableTransducer output_fst; + if (!ApplyRules(input, &output_fst, use_lookahead)) return false; + MutableTransducer shortest_path; + fst::ShortestPath(output_fst, &shortest_path); + fst::Project(&shortest_path, fst::PROJECT_OUTPUT); + fst::RmEpsilon(&shortest_path); + Printer printer(fst::TokenType::BYTE); + if (!printer.operator()(shortest_path, output)) { + LoggerError("Failed to print to output string"); + return false; + } + return true; +} + +const Transducer* RuleSystem::FindRule(const string& name) const { + return grm_->GetFst(name); +} + + +} +} diff --git a/native/tts-normalizer/third_party/sparrowhawk_string_matcher.h b/native/tts-normalizer/third_party/sparrowhawk_string_matcher.h new file mode 100644 index 0000000..1892bdd --- /dev/null +++ b/native/tts-normalizer/third_party/sparrowhawk_string_matcher.h @@ -0,0 +1,583 @@ +#ifndef NEMOTRON_SPARROWHAWK_STRING_MATCHER_H_ +#define NEMOTRON_SPARROWHAWK_STRING_MATCHER_H_ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace speech { +namespace sparrowhawk { + +class StringMatcher { + public: + using Arc = fst::StdArc; + using StateId = Arc::StateId; + using Transducer = fst::Fst; + using MutableTransducer = fst::VectorFst; + + struct Limits { + bool canonical_word_outputs = false; + std::size_t grammar_states = 1000000; + std::size_t grammar_arcs = 4000000; + std::size_t input_bytes = 65536; + std::size_t visited_states = 1000000; + std::size_t arc_relaxations = 4000000; + std::size_t backpointers = 1000000; + double absolute_margin = 0.001; + double relative_margin = 0.000001; + }; + + struct Diagnostics { + const char* reason = "not_called"; + std::size_t visited_states = 0; + std::size_t arc_relaxations = 0; + std::size_t backpointers = 0; + float best_cost = std::numeric_limits::infinity(); + float second_cost = std::numeric_limits::infinity(); + float discarded_cost = std::numeric_limits::infinity(); + double raw_path_roundoff_bound = 0; + double ambiguity_guard = 0; + }; + + explicit StringMatcher(const Transducer& grammar) + : StringMatcher(grammar, Limits{}) {} + + StringMatcher(const Transducer& grammar, Limits limits) : limits_(limits) { + try { + valid_ = Build(grammar); + } catch (const std::bad_alloc&) { + build_reason_ = "grammar_allocation"; + } catch (const std::length_error&) { + build_reason_ = "grammar_allocation"; + } + } + + bool valid() const { return valid_; } + const Diagnostics& diagnostics() const { return diagnostics_; } + + bool Match(const Transducer& input, MutableTransducer* output) { + diagnostics_ = Diagnostics{}; + try { + if (!output) return Fail("null_output"); + if (!valid_) return Fail(build_reason_); + return MatchImpl(input, output); + } catch (const std::bad_alloc&) { + return Fail("request_allocation"); + } catch (const std::length_error&) { + return Fail("request_allocation"); + } + } + + private: + static constexpr std::uint32_t kNoNode = + std::numeric_limits::max(); + static constexpr float kInf = std::numeric_limits::infinity(); + + struct FlatArc { + std::int32_t input; + std::int32_t output; + StateId next; + float weight; + }; + struct State { + std::size_t epsilon_begin; + std::size_t epsilon_end; + std::size_t consume_begin; + std::size_t consume_end; + float final; + std::uint32_t rank = 0; + std::array next_bytes{}; + bool can_final = false; + }; + struct Choice { + float cost = kInf; + std::uint32_t prefix = 0; + std::uint32_t node = kNoNode; + float final_weight = 0; + }; + struct Cell { + std::uint64_t generation = 0; + Choice best; + Choice second; + float discarded = kInf; + double absolute_sum = 0; + std::uint32_t steps = 0; + }; + struct PrefixKey { + std::uint32_t parent; + std::int32_t input; + std::int32_t output; + bool operator==(const PrefixKey& other) const { + return parent == other.parent && input == other.input && output == other.output; + } + }; + struct PrefixMetadata { + bool in_quotes = false; + bool escaped = false; + bool outside_space = false; + }; + struct PrefixHash { + std::size_t operator()(const PrefixKey& key) const { + std::size_t value = key.parent; + value = value * 1000003 + std::size_t(key.input); + return value * 1000003 + std::size_t(key.output); + } + }; + using QueueItem = std::pair; + struct Layer { + std::vector cells; + std::vector queue; + std::uint64_t generation = 0; + + void Begin() { + queue.clear(); + if (++generation == 0) { + for (auto& cell : cells) cell.generation = 0; + generation = 1; + } + } + + Cell& Touch(StateId state, std::uint32_t rank) { + Cell& cell = cells[state]; + if (cell.generation != generation) { + cell = Cell{}; + cell.generation = generation; + queue.emplace_back(rank, state); + std::push_heap(queue.begin(), queue.end(), std::greater()); + } + return cell; + } + + StateId Pop() { + std::pop_heap(queue.begin(), queue.end(), std::greater()); + StateId result = queue.back().second; + queue.pop_back(); + return result; + } + }; + struct Backpointer { + std::uint32_t parent; + FlatArc arc; + }; + + bool Fail(const char* reason) { + diagnostics_.reason = reason; + diagnostics_.backpointers = path_.size(); + return false; + } + + bool Build(const Transducer& grammar) { + if (!std::isfinite(limits_.absolute_margin) || + !std::isfinite(limits_.relative_margin) || + limits_.absolute_margin < 0 || limits_.relative_margin < 0 || + limits_.grammar_states >= kNoNode || limits_.backpointers >= kNoNode) { + build_reason_ = "invalid_limits"; + return false; + } + start_ = grammar.Start(); + for (fst::StateIterator it(grammar); !it.Done(); it.Next()) { + const StateId id = it.Value(); + if (id < 0 || static_cast(id) != states_.size() || + states_.size() >= limits_.grammar_states) return false; + const float final = grammar.Final(id).Value(); + if (std::isnan(final) || final == -kInf) return false; + State state{epsilons_.size(), 0, consuming_.size(), 0, final, 0}; + for (fst::ArcIterator arcs(grammar, id); !arcs.Done(); + arcs.Next()) { + const Arc& arc = arcs.Value(); + const float weight = arc.weight.Value(); + if (arc.ilabel < 0 || arc.ilabel > 255 || arc.olabel < 0 || + arc.olabel > 255 || arc.nextstate < 0 || !std::isfinite(weight)) { + return false; + } + if (epsilons_.size() + consuming_.size() >= limits_.grammar_arcs) + return false; + FlatArc flat{arc.ilabel, arc.olabel, arc.nextstate, weight}; + if (arc.ilabel == 0) epsilons_.push_back(flat); + else consuming_.push_back(flat); + } + state.epsilon_end = epsilons_.size(); + state.consume_end = consuming_.size(); + std::stable_sort(consuming_.begin() + state.consume_begin, + consuming_.end(), + [](const FlatArc& a, const FlatArc& b) { + return a.input < b.input; + }); + states_.push_back(state); + } + if (start_ < 0 || static_cast(start_) >= states_.size()) + return false; + std::vector indegree(states_.size(), 0); + for (const auto& arc : consuming_) + if (static_cast(arc.next) >= states_.size()) return false; + for (const auto& arc : epsilons_) { + if (static_cast(arc.next) >= states_.size()) return false; + ++indegree[arc.next]; + } + std::vector ready; + ready.reserve(states_.size()); + for (std::size_t i = 0; i < states_.size(); ++i) + if (indegree[i] == 0) ready.push_back(static_cast(i)); + for (std::size_t i = 0; i < ready.size(); ++i) { + const StateId id = ready[i]; + states_[id].rank = static_cast(i); + for (std::size_t a = states_[id].epsilon_begin; + a < states_[id].epsilon_end; ++a) + if (--indegree[epsilons_[a].next] == 0) + ready.push_back(epsilons_[a].next); + } + if (ready.size() != states_.size()) { + build_reason_ = "input_epsilon_cycle"; + return false; + } + for (auto it = ready.rbegin(); it != ready.rend(); ++it) { + State& state = states_[*it]; + state.can_final = state.final != kInf; + for (std::size_t a = state.consume_begin; a < state.consume_end; ++a) { + const unsigned label = consuming_[a].input; + state.next_bytes[label / 64] |= std::uint64_t(1) << (label % 64); + } + for (std::size_t a = state.epsilon_begin; a < state.epsilon_end; ++a) { + const State& next = states_[epsilons_[a].next]; + state.can_final = state.can_final || next.can_final; + for (unsigned word = 0; word < 4; ++word) + state.next_bytes[word] |= next.next_bytes[word]; + } + } + layers_[0].cells.resize(states_.size()); + layers_[1].cells.resize(states_.size()); + return true; + } + + bool ReadInput(const Transducer& input) { + bytes_.clear(); + StateId state = input.Start(); + if (state == fst::kNoStateId) return false; + for (std::size_t count = 0; count <= limits_.input_bytes; ++count) { + const float final = input.Final(state).Value(); + fst::ArcIterator it(input, state); + if (it.Done()) return final == 0; + if (final != kInf || count == limits_.input_bytes) return false; + const Arc arc = it.Value(); + it.Next(); + if (!it.Done() || arc.ilabel < 1 || arc.ilabel > 255 || + arc.ilabel != arc.olabel || arc.weight.Value() != 0 || + arc.nextstate == fst::kNoStateId) return false; + bytes_.push_back(arc.ilabel); + state = arc.nextstate; + } + return false; + } + + bool Reachable(StateId state, std::size_t position) const { + const State& value = states_[state]; + if (position == bytes_.size()) return value.can_final; + const unsigned label = bytes_[position]; + return (value.next_bytes[label / 64] & (std::uint64_t(1) << (label % 64))) != 0; + } + + bool Prefix(std::uint32_t parent, const FlatArc& arc, std::uint32_t* output) { + PrefixMetadata metadata; + if (limits_.canonical_word_outputs) { + metadata = prefix_metadata_[parent]; + if (arc.output == 0 || + (!metadata.in_quotes && metadata.outside_space && arc.output == ' ')) { + *output = parent; + return true; + } + if (metadata.in_quotes) { + if (metadata.escaped) metadata.escaped = false; + else if (arc.output == '\\') metadata.escaped = true; + else if (arc.output == '"') metadata.in_quotes = false; + metadata.outside_space = false; + } else { + metadata.in_quotes = arc.output == '"'; + metadata.escaped = false; + metadata.outside_space = arc.output == ' '; + } + } else if (arc.input == 0 && arc.output == 0) { + *output = parent; + return true; + } + const PrefixKey key{parent, limits_.canonical_word_outputs ? 0 : arc.input, arc.output}; + const auto found = prefixes_.find(key); + if (found != prefixes_.end()) { + *output = found->second; + return true; + } + if (prefixes_.size() >= limits_.backpointers) return Fail("prefix_budget"); + const auto id = static_cast(prefixes_.size() + 1); + prefixes_.emplace(key, id); + prefix_metadata_.push_back(metadata); + *output = id; + return true; + } + + bool Consider(Cell* target, Choice candidate, const FlatArc* arc) { + Choice* slot = nullptr; + if (target->best.cost != kInf && candidate.prefix == target->best.prefix) { + if (candidate.cost >= target->best.cost) return true; + slot = &target->best; + } else if (target->second.cost != kInf && candidate.prefix == target->second.prefix) { + if (candidate.cost >= target->second.cost) return true; + slot = &target->second; + } else if (candidate.cost < target->best.cost) { + target->discarded = std::min(target->discarded, target->second.cost); + target->second = target->best; + slot = &target->best; + } else if (candidate.cost < target->second.cost) { + target->discarded = std::min(target->discarded, target->second.cost); + slot = &target->second; + } else { + target->discarded = std::min(target->discarded, candidate.cost); + return true; + } + if (arc) { + if (path_.size() >= limits_.backpointers) return Fail("backpointer_budget"); + const auto node = static_cast(path_.size()); + path_.push_back(Backpointer{candidate.node, *arc}); + candidate.node = node; + } + *slot = candidate; + if (target->second.cost < target->best.cost) + std::swap(target->best, target->second); + return true; + } + + bool Relax(const Cell& source, const FlatArc& arc, Layer* destination, + std::size_t destination_position) { + if (++diagnostics_.arc_relaxations > limits_.arc_relaxations) + return Fail("arc_budget"); + if (!Reachable(arc.next, destination_position)) return true; + Cell& target = destination->Touch(arc.next, states_[arc.next].rank); + if (source.steps == std::numeric_limits::max()) + return Fail("path_length"); + target.steps = std::max(target.steps, source.steps + 1); + target.absolute_sum = std::max(target.absolute_sum, + source.absolute_sum + std::abs(double(arc.weight))); + for (const Choice* choice : {&source.best, &source.second}) { + if (choice->cost == kInf) continue; + Choice candidate = *choice; + candidate.cost = choice->cost + arc.weight; + if (!std::isfinite(candidate.cost)) return Fail("nonfinite_path"); + if (!Prefix(choice->prefix, arc, &candidate.prefix)) return false; + if (!Consider(&target, candidate, &arc)) return false; + } + if (source.discarded != kInf) { + const float candidate = source.discarded + arc.weight; + if (!std::isfinite(candidate)) return Fail("nonfinite_discarded_path"); + target.discarded = std::min(target.discarded, candidate); + } + return true; + } + + std::vector DebugWitness(std::uint32_t node) const { + std::vector result; + while (node != kNoNode) { + if (node >= path_.size() || result.size() >= limits_.backpointers) + throw std::length_error("invalid diagnostic witness"); + const auto& value = path_[node]; + if (value.arc.input != 0 || value.arc.output != 0) + result.push_back(value.arc); + node = value.parent; + } + std::reverse(result.begin(), result.end()); + return result; + } + + static void DebugJsonString(const std::string& value) { + std::fputc('"', stderr); + for (unsigned char byte : value) { + if (byte == '"' || byte == '\\') std::fprintf(stderr, "\\%c", byte); + else if (byte < 32) std::fprintf(stderr, "\\u%04x", unsigned(byte)); + else std::fputc(byte, stderr); + } + std::fputc('"', stderr); + } + + static void DebugPath(const Choice& choice, const std::vector& arcs, + std::size_t first_difference) { + if (choice.cost == kInf) { + std::fprintf(stderr, "null"); + return; + } + std::string input, output; + for (const auto& arc : arcs) { + if (arc.input) input.push_back(char(arc.input)); + if (arc.output) output.push_back(char(arc.output)); + } + std::fprintf(stderr, "{\"cost\":%.9g,\"prefix_id\":%u,\"final_weight\":%.9g,\"input_utf8\":", + choice.cost, choice.prefix, choice.final_weight); + DebugJsonString(input); + std::fprintf(stderr, ",\"output_utf8\":"); + DebugJsonString(output); + std::fprintf(stderr, ",\"pairs\":["); + for (std::size_t i = 0; i < arcs.size(); ++i) { + if (i) std::fputc(',', stderr); + std::fprintf(stderr, "[%d,%d]", arcs[i].input, arcs[i].output); + } + std::fprintf(stderr, "],\"difference_context\":["); + const std::size_t begin = first_difference > 8 ? first_difference - 8 : 0; + const std::size_t end = std::min(arcs.size(), first_difference + 9); + std::size_t input_offset = 0, output_offset = 0; + bool first = true; + for (std::size_t i = 0; i < arcs.size(); ++i) { + if (i >= begin && i < end) { + if (!first) std::fputc(',', stderr); + first = false; + std::fprintf(stderr, "{\"pair\":%zu,\"input_offset\":%zu,\"output_offset\":%zu,\"ilabel\":%d,\"olabel\":%d}", + i, input_offset, output_offset, arcs[i].input, arcs[i].output); + } + input_offset += arcs[i].input != 0; + output_offset += arcs[i].output != 0; + } + std::fprintf(stderr, "]}"); + } + + void DebugPaths(const Cell& final) const { + if (!std::getenv("NEMOTRON_TN_PATH_DEBUG")) return; + const auto best = DebugWitness(final.best.node); + const auto second = DebugWitness(final.second.node); + std::size_t difference = 0; + while (difference < best.size() && difference < second.size() && + best[difference].input == second[difference].input && + best[difference].output == second[difference].output) ++difference; + std::fprintf(stderr, "TN_PATH {\"scope\":\"diagnostic_only\",\"first_differing_pair\":%zu,\"guard\":%.17g,\"best\":", + difference, diagnostics_.ambiguity_guard); + DebugPath(final.best, best, difference); + std::fprintf(stderr, ",\"second\":"); + DebugPath(final.second, second, difference); + std::fprintf(stderr, "}\n"); + } + + bool MatchImpl(const Transducer& input, MutableTransducer* output) { + path_.clear(); + prefixes_.clear(); + prefix_metadata_.clear(); + prefix_metadata_.emplace_back(); // Canonical empty prefix is ID zero. + if (!ReadInput(input)) return Fail("unsupported_input"); + if (!Reachable(start_, 0)) return Fail("no_path"); + Layer* current = &layers_[0]; + Layer* next = &layers_[1]; + current->Begin(); + next->Begin(); + current->Touch(start_, states_[start_].rank).best.cost = 0; + Cell final; + for (std::size_t position = 0; position <= bytes_.size(); ++position) { + while (!current->queue.empty()) { + const StateId id = current->Pop(); + if (++diagnostics_.visited_states > limits_.visited_states) + return Fail("state_budget"); + const Cell cell = current->cells[id]; + const State& state = states_[id]; + for (std::size_t a = state.epsilon_begin; a < state.epsilon_end; ++a) + if (!Relax(cell, epsilons_[a], current, position)) return false; + if (position < bytes_.size()) { + auto begin = consuming_.begin() + state.consume_begin; + auto end = consuming_.begin() + state.consume_end; + auto it = std::lower_bound(begin, end, bytes_[position], + [](const FlatArc& arc, std::int32_t byte) { + return arc.input < byte; + }); + for (; it != end && it->input == bytes_[position]; ++it) + if (!Relax(cell, *it, next, position + 1)) return false; + } else if (state.final != kInf) { + final.steps = std::max(final.steps, cell.steps + 1); + final.absolute_sum = std::max( + final.absolute_sum, cell.absolute_sum + std::abs(double(state.final))); + for (const Choice* choice : {&cell.best, &cell.second}) { + if (choice->cost == kInf) continue; + Choice candidate = *choice; + candidate.cost = choice->cost + state.final; + candidate.final_weight = state.final; + if (!std::isfinite(candidate.cost)) return Fail("nonfinite_final"); + if (!Consider(&final, candidate, nullptr)) return false; + } + if (cell.discarded != kInf) { + const float candidate = cell.discarded + state.final; + if (!std::isfinite(candidate)) return Fail("nonfinite_discarded_final"); + final.discarded = std::min(final.discarded, candidate); + } + } + } + if (position < bytes_.size()) { + std::swap(current, next); + next->Begin(); + } + } + diagnostics_.best_cost = final.best.cost; + diagnostics_.second_cost = final.second.cost; + diagnostics_.discarded_cost = final.discarded; + if (final.best.cost == kInf) return Fail("no_path"); + const double nu = final.steps * (double(std::numeric_limits::epsilon()) / 2); + if (nu >= 0.125) return Fail("roundoff_budget"); + const double error = nu / (1 - nu) * final.absolute_sum; + const double guard = 2 * error + limits_.absolute_margin + + limits_.relative_margin * std::abs(double(final.best.cost)); + diagnostics_.raw_path_roundoff_bound = error; + diagnostics_.ambiguity_guard = guard; + DebugPaths(final); + const float alternative = std::min(final.second.cost, final.discarded); + if (alternative != kInf && double(alternative) - final.best.cost <= guard) + return Fail("ambiguous_path"); + if (limits_.canonical_word_outputs && + (prefix_metadata_[final.best.prefix].in_quotes || + prefix_metadata_[final.best.prefix].escaped)) + return Fail("unbalanced_output_quotes"); + reverse_.clear(); + for (std::uint32_t node = final.best.node; node != kNoNode; + node = path_[node].parent) { + if (node >= path_.size() || reverse_.size() >= limits_.backpointers) + return Fail("invalid_backpointer"); + reverse_.push_back(node); + } + MutableTransducer result; + StateId from = result.AddState(); + result.SetStart(from); + for (auto it = reverse_.rbegin(); it != reverse_.rend(); ++it) { + const FlatArc& arc = path_[*it].arc; + const StateId to = result.AddState(); + result.AddArc(from, Arc(arc.input, arc.output, Arc::Weight(arc.weight), to)); + from = to; + } + result.SetFinal(from, Arc::Weight(final.best.final_weight)); + *output = std::move(result); + diagnostics_.reason = "matched_experimental"; + diagnostics_.backpointers = path_.size(); + return true; + } + + Limits limits_; + bool valid_ = false; + const char* build_reason_ = "unsupported_grammar"; + StateId start_ = fst::kNoStateId; + std::vector states_; + std::vector epsilons_; + std::vector consuming_; + Layer layers_[2]; + std::vector path_; + std::unordered_map prefixes_; + std::vector prefix_metadata_; + std::vector reverse_; + std::vector bytes_; + Diagnostics diagnostics_; +}; + +} // namespace sparrowhawk +} // namespace speech + +#endif // NEMOTRON_SPARROWHAWK_STRING_MATCHER_H_ diff --git a/python/nemotron_cpp_bridge/__init__.py b/python/nemotron_cpp_bridge/__init__.py new file mode 100644 index 0000000..948c558 --- /dev/null +++ b/python/nemotron_cpp_bridge/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unix socket adapter for the native Nemotron TTS controller.""" diff --git a/python/nemotron_cpp_bridge/__main__.py b/python/nemotron_cpp_bridge/__main__.py new file mode 100644 index 0000000..e9eb558 --- /dev/null +++ b/python/nemotron_cpp_bridge/__main__.py @@ -0,0 +1,5 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from .server import main + +main() diff --git a/python/nemotron_cpp_bridge/native_preparation.py b/python/nemotron_cpp_bridge/native_preparation.py new file mode 100644 index 0000000..720399c --- /dev/null +++ b/python/nemotron_cpp_bridge/native_preparation.py @@ -0,0 +1,122 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import ctypes as C +import json +import threading +from pathlib import Path + + +class FrontendPreparation: + def __init__(self, library, tokenizer_path, config_path=None): + config = json.loads(Path(config_path).read_text()) if config_path is not None else None + if config_path is not None and not isinstance(config, dict): + raise ValueError("normalization configuration must be an object") + self.enabled = config is not None and config.get("enabled", True) + if type(self.enabled) is not bool: + raise ValueError("normalization enabled must be Boolean") + self.metadata = { + "library": str(Path(library).resolve(strict=True)), + "tokenizer": str(tokenizer_path), + "normalization_config": json.dumps(config), + } + + def close(self): + pass + + +class NativePreparation: + def __init__(self, library, tokenizer_path, config_path=None): + config = json.loads(Path(config_path).read_text()) if config_path is not None else None + if config_path is not None and not isinstance(config, dict): + raise ValueError("normalization configuration must be an object") + self.enabled = config is not None and config.get("enabled", True) + self._condition = threading.Condition() + self._active = 0 + self._closed = False + self._library = C.CDLL(str(Path(library).resolve(strict=True))) + self._create = self._library.nemotron_text_create + self._create.argtypes = [ + C.c_char_p, + C.c_char_p, + C.POINTER(C.c_void_p), + C.POINTER(C.c_void_p), + ] + self._create.restype = C.c_int + self._prepare = self._library.nemotron_text_prepare + self._prepare.argtypes = [ + C.c_void_p, + C.c_char_p, + C.c_size_t, + C.c_uint8, + C.POINTER(C.c_void_p), + C.POINTER(C.c_void_p), + ] + self._prepare.restype = C.c_int + self._free = self._library.nemotron_text_free + self._free.argtypes = [C.c_void_p] + self._free.restype = None + self._destroy = self._library.nemotron_text_destroy + self._destroy.argtypes = [C.c_void_p] + self._destroy.restype = None + self._handle = C.c_void_p() + error = C.c_void_p() + status = self._create( + json.dumps(config).encode(), + str(tokenizer_path).encode(), + C.byref(self._handle), + C.byref(error), + ) + self._check(status, error) + if not self._handle.value: + raise RuntimeError("native preparation returned an empty handle") + + def _check(self, status, error): + try: + if status: + message = ( + C.string_at(error).decode("utf-8", errors="replace") + if error.value + else "native preparation failed" + ) + raise (ValueError if status == 1 else RuntimeError)(message) + finally: + if error.value: + self._free(error) + + def prepare_tokens(self, text, apply_normalize=True): + if not isinstance(text, str) or type(apply_normalize) is not bool: + raise ValueError("invalid text preparation input") + data = text.strip().encode("utf-8") + with self._condition: + if self._closed: + raise RuntimeError("text preparer is closed") + self._active += 1 + output = C.c_void_p() + error = C.c_void_p() + try: + status = self._prepare( + self._handle, data, len(data), apply_normalize, C.byref(output), C.byref(error) + ) + self._check(status, error) + if not output.value: + raise RuntimeError("native preparation returned an empty result") + result = json.loads(C.string_at(output)) + return result["text"], result["tokens"] + finally: + if output.value: + self._free(output) + with self._condition: + self._active -= 1 + self._condition.notify_all() + + def close(self): + with self._condition: + if self._closed: + self._condition.wait_for(lambda: not self._handle.value) + return + self._closed = True + self._condition.wait_for(lambda: self._active == 0) + self._destroy(self._handle) + self._handle = C.c_void_p() + self._condition.notify_all() diff --git a/python/nemotron_cpp_bridge/normalization.py b/python/nemotron_cpp_bridge/normalization.py new file mode 100644 index 0000000..e514ba8 --- /dev/null +++ b/python/nemotron_cpp_bridge/normalization.py @@ -0,0 +1,344 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import ctypes as C +import json +import re +import struct +import threading +from pathlib import Path + + +class _Span(C.Structure): + _fields_ = [("data", C.POINTER(C.c_uint8)), ("len", C.c_size_t)] + + +class _Status(C.Structure): + _fields_ = [ + ("struct_size", C.c_size_t), + ("code", C.c_int32), + ("message", _Span), + ("context", C.c_void_p), + ] + + +class _Record(C.Structure): + _fields_ = [ + ("struct_size", C.c_size_t), + ("name", _Span), + ("kind", C.c_uint32), + ("element_type", C.c_uint32), + ("data", _Span), + ("shape", C.POINTER(C.c_size_t)), + ("rank", C.c_size_t), + ("sample_rate_hz", C.c_uint32), + ("channels", C.c_uint16), + ("audio_encoding", C.c_uint16), + ("sequence", C.c_uint64), + ("has_sequence", C.c_uint8), + ("is_final", C.c_uint8), + ("reserved", C.c_uint8 * 6), + ] + + +class _Envelope(C.Structure): + _fields_ = [ + ("struct_size", C.c_size_t), + ("schema_version", C.c_uint32), + ("request_id", _Span), + ("session_id", _Span), + ("records", C.POINTER(_Record)), + ("record_count", C.c_size_t), + ("sequence", C.c_uint64), + ("has_sequence", C.c_uint8), + ("is_final", C.c_uint8), + ("reserved", C.c_uint8 * 6), + ] + + +class _Output(C.Structure): + _fields_ = _Envelope._fields_ + [("context", C.c_void_p)] + + +class _Capabilities(C.Structure): + _fields_ = [ + ("struct_size", C.c_size_t), + ("processor_id", _Span), + ("batch_input", C.c_uint8), + ("streaming_input", C.c_uint8), + ("thread_safe", C.c_uint8), + ("stateful", C.c_uint8), + ("reserved", C.c_uint8 * 4), + ] + + +class _Api(C.Structure): + _fields_ = [ + ("struct_size", C.c_size_t), + ("abi_version", C.c_uint32), + ("capabilities", C.CFUNCTYPE(_Status, C.POINTER(_Capabilities))), + ("create", C.CFUNCTYPE(_Status, C.POINTER(_Envelope), C.POINTER(C.c_void_p))), + ("destroy", C.CFUNCTYPE(None, C.c_void_p)), + ("process", C.CFUNCTYPE(_Status, C.c_void_p, C.POINTER(_Envelope), C.POINTER(_Output))), + ("release_output", C.CFUNCTYPE(None, C.c_void_p, C.POINTER(_Output))), + ("release_status", C.CFUNCTYPE(None, C.POINTER(_Status))), + ] + + +def _bytes(span, limit=1024 * 1024): + if span.len > limit or (span.len and not span.data): + raise RuntimeError("native normalizer returned an invalid byte span") + return C.string_at(span.data, span.len) if span.len else b"" + + +def _envelope(values, identity=b"nemotron-cpp-normalization"): + storage = [] + + def span(value): + buffer = C.create_string_buffer(value) + storage.append(buffer) + return _Span(C.cast(buffer, C.POINTER(C.c_uint8)), len(value)) + + records = (_Record * len(values))() + for record, (name, value) in zip(records, values.items()): + record.struct_size = C.sizeof(_Record) + record.name = span(name.encode()) + if isinstance(value, bool): + record.kind, record.element_type, data = 8, 1, bytes([value]) + elif isinstance(value, int): + record.kind, record.element_type, data = 6, 3, struct.pack("=q", value) + else: + record.kind, record.element_type, data = 1, 1, value.encode("utf-8") + record.data = span(data) + envelope = _Envelope( + struct_size=C.sizeof(_Envelope), + schema_version=1, + request_id=span(identity), + records=records, + record_count=len(records), + is_final=1, + ) + storage.append(records) + return envelope, storage + + +class TextPreparer: + def __init__(self, config_path=None): + self.enabled = False + self._condition = threading.Condition() + self._handles = [] + self._active = 0 + self._closed = False + self._api = None + self._max_input = 65536 + self._max_output = 262144 + self._max_segments = 256 + if config_path is None: + return + config = json.loads(Path(config_path).read_text()) + allowed = { + "enabled", + "native_library_root", + "native_processor", + "normalizer_root", + "normalizer_path", + "language_code", + "normalizer_config", + "direct_grammar", + "abbreviations_path", + "append_punctuation", + "pool_size", + "bounded_preparation", + "max_input_bytes", + "max_normalized_bytes", + "max_segments", + } + if not isinstance(config, dict) or set(config) - allowed: + raise ValueError("invalid normalization configuration fields") + for key in ("enabled", "direct_grammar", "append_punctuation", "bounded_preparation"): + if key in config and type(config[key]) is not bool: + raise ValueError(f"{key} must be Boolean") + if config.get("language_code", "en-US") != "en-US": + raise ValueError("normalization language_code must be en-US") + self.enabled = config.get("enabled", True) + limits = { + "max_input_bytes": 65536, + "max_normalized_bytes": 262144, + "max_segments": 256, + "pool_size": 2, + } + for key, default in limits.items(): + value = config.get(key, default) + if type(value) is not int or not 0 < value < 2**63: + raise ValueError(f"{key} must fit a positive ABI integer") + limits[key] = value + self._max_input = limits["max_input_bytes"] + self._max_output = limits["max_normalized_bytes"] + self._max_segments = limits["max_segments"] + library_root = Path(config["native_library_root"]).resolve(strict=True) + asset_root = Path(config["normalizer_root"]).resolve(strict=True) + library = Path(config["native_processor"]).resolve(strict=True) + language_path = (asset_root / config.get("normalizer_path", "en-US")).resolve(strict=True) + if not library_root.is_dir() or not asset_root.is_dir() or not language_path.is_dir(): + raise ValueError("normalizer roots and language path must be directories") + if library_root.is_relative_to(asset_root) or asset_root.is_relative_to(library_root): + raise ValueError("normalizer libraries and assets must have separate roots") + if not library.is_file() or not library.is_relative_to(library_root): + raise ValueError("native_processor must be a file inside native_library_root") + if not language_path.is_relative_to(asset_root): + raise ValueError("normalizer_path is outside normalizer_root") + normalizer_config = config.get("normalizer_config", "sparrowhawk_configuration.ascii_proto") + if ( + not isinstance(normalizer_config, str) + or Path(normalizer_config).name != normalizer_config + ): + raise ValueError("normalizer_config must name a file in normalizer_path") + required = ( + ("tokenize_and_classify.far", "verbalize.far") + if config.get("direct_grammar", False) + else (normalizer_config,) + ) + for name in required: + path = (language_path / name).resolve(strict=True) + if not path.is_file() or not path.is_relative_to(language_path): + raise ValueError("normalizer input escapes language directory") + if config.get("direct_grammar", False): + for name in ("pre_process.far", "post_process.far", "speech_class.far"): + path = language_path / name + if path.exists() or path.is_symlink(): + path = path.resolve(strict=True) + if not path.is_file() or not path.is_relative_to(language_path): + raise ValueError("normalizer input escapes language directory") + values = { + "language_code": "en-US", + "normalizer_path": str(language_path), + "normalizer_config": normalizer_config, + "direct_grammar": config.get("direct_grammar", False), + "append_punctuation": config.get("append_punctuation", True), + **{key: value for key, value in limits.items() if key != "pool_size"}, + } + if config.get("abbreviations_path") is not None: + path = (language_path / config["abbreviations_path"]).resolve(strict=True) + if not path.is_file() or not path.is_relative_to(asset_root): + raise ValueError("abbreviations_path must be inside normalizer_root") + values["abbreviations_path"] = str(path) + self._library = C.CDLL(str(library)) + getter = self._library.nemotron_speech_processor_get_api + getter.argtypes = [C.c_uint32] + getter.restype = C.POINTER(_Api) + api = getter(1) + if not api or api.contents.struct_size < C.sizeof(_Api) or api.contents.abi_version != 1: + raise RuntimeError("unsupported native normalizer ABI") + self._api = api.contents + if not all( + ( + self._api.capabilities, + self._api.create, + self._api.destroy, + self._api.process, + self._api.release_output, + self._api.release_status, + ) + ): + raise RuntimeError("native normalizer ABI has missing functions") + capabilities = _Capabilities(struct_size=C.sizeof(_Capabilities)) + self._check(self._api.capabilities(C.byref(capabilities))) + if ( + _bytes(capabilities.processor_id) != b"nemotron_speech.tts.text_preprocessor" + or capabilities.stateful + ): + raise RuntimeError("native processor is not the stateless TTS normalizer") + envelope, _storage = _envelope(values) + try: + for _ in range(limits["pool_size"]): + handle = C.c_void_p() + self._check(self._api.create(C.byref(envelope), C.byref(handle))) + if not handle.value: + raise RuntimeError("native normalizer returned an empty handle") + self._handles.append(handle) + except BaseException: + self.close() + raise + + def _check(self, status): + try: + if status.code: + message = _bytes(status.message).decode("utf-8", errors="replace") + error = ValueError if 100 <= status.code <= 102 else RuntimeError + raise error(f"native normalizer error {status.code}: {message}") + finally: + self._api.release_status(C.byref(status)) + + def prepare(self, text, language, apply_normalize=True): + if not isinstance(text, str) or not isinstance(language, str): + raise ValueError("text and language must be strings") + if language.lower() != "en-us": + raise ValueError("native backend supports only en-US") + if type(apply_normalize) is not bool: + raise ValueError("apply_normalize must be Boolean") + text = text.strip() + if not text or len(text.encode("utf-8")) > self._max_input: + raise ValueError("text is empty or exceeds max_input_bytes") + with self._condition: + if self._closed: + raise RuntimeError("text preparer is closed") + if not self.enabled or not apply_normalize: + return text + self._condition.wait_for(lambda: self._handles or self._closed) + if self._closed: + raise RuntimeError("text preparer is closed") + handle = self._handles.pop() + self._active += 1 + try: + return self._process(handle, text) + finally: + with self._condition: + self._handles.append(handle) + self._active -= 1 + self._condition.notify_all() + + def _process(self, handle, text): + request, _storage = _envelope({"text": text, "language_code": "en-US"}) + output = _Output(struct_size=C.sizeof(_Output)) + try: + self._check(self._api.process(handle, C.byref(request), C.byref(output))) + if ( + output.struct_size < C.sizeof(_Output) + or output.schema_version != 1 + or _bytes(output.request_id) != _bytes(request.request_id) + or _bytes(output.session_id) != _bytes(request.session_id) + or output.sequence != request.sequence + or output.has_sequence != request.has_sequence + or output.is_final != request.is_final + ): + raise RuntimeError("native normalizer changed envelope identity") + if not output.records or not 0 < output.record_count <= self._max_segments: + raise RuntimeError("native normalizer returned invalid segment count") + segments = [] + size = output.record_count - 1 + for index in range(output.record_count): + record = output.records[index] + if ( + record.struct_size < C.sizeof(_Record) + or record.kind != 1 + or _bytes(record.name) != f"normalized_text.{index}".encode() + ): + raise RuntimeError("native normalizer returned invalid segment order/type") + data = _bytes(record.data, self._max_output) + size += len(data) + segment = data.decode("utf-8") + if not segment.strip() or size > self._max_output: + raise RuntimeError("native normalizer returned empty or oversized text") + segments.append(segment) + return re.sub(r" +([,.?!।])", r"\1", " ".join(segments)) + finally: + self._api.release_output(handle, C.byref(output)) + + def close(self): + with self._condition: + self._closed = True + self._condition.notify_all() + self._condition.wait_for(lambda: self._active == 0) + handles, self._handles = self._handles, [] + for handle in handles: + self._api.destroy(handle) diff --git a/python/nemotron_cpp_bridge/protocol.py b/python/nemotron_cpp_bridge/protocol.py new file mode 100644 index 0000000..eaf8f9b --- /dev/null +++ b/python/nemotron_cpp_bridge/protocol.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import asyncio +import json +import math +import struct + +try: + import numpy as np +except ImportError: + np = None + +MAX_FRAME_BYTES = 64 * 1024 * 1024 +MAX_CONTROL_FRAME_BYTES = 8 * 1024 * 1024 + + +def control_frame(kind, **fields): + payload = b"J" + json.dumps( + {"kind": kind, "protocol_version": 1, **fields}, + ensure_ascii=False, + allow_nan=False, + separators=(",", ":"), + ).encode("utf-8") + return frame(payload) + + +def frame(payload): + if not 0 < len(payload) <= MAX_FRAME_BYTES: + raise ValueError("bridge frame is outside the allowed size") + return struct.pack(">I", len(payload)) + payload + + +def audio_frame(request_id, sequence, pcm): + if not pcm or len(pcm) % 2: + raise ValueError("invalid PCM16 extent") + return frame(b"A" + struct.pack(">QI", request_id, sequence) + pcm) + + +def float_audio_frame(request_id, sequence, pcm): + if not pcm or len(pcm) % 4: + raise ValueError("invalid PCM-F32LE extent") + return frame(b"F" + struct.pack(">QI", request_id, sequence) + pcm) + + +def _object(pairs): + result = {} + for key, value in pairs: + if key in result: + raise ValueError("duplicate JSON field") + result[key] = value + return result + + +async def read_control(reader): + length = struct.unpack(">I", await reader.readexactly(4))[0] + if not 1 < length <= min(MAX_FRAME_BYTES, MAX_CONTROL_FRAME_BYTES): + raise ValueError("invalid control frame length") + payload = await reader.readexactly(length) + if payload[:1] != b"J": + raise ValueError("client must send JSON control frames") + value = json.loads( + payload[1:], + object_pairs_hook=_object, + parse_constant=lambda _: (_ for _ in ()).throw(ValueError("nonfinite JSON value")), + ) + if not isinstance(value, dict) or value.get("protocol_version", 1) != 1: + raise ValueError("unsupported bridge protocol") + return value + + +def pcm16_scalar(payload): + if len(payload) % 4: + raise ValueError("invalid PCM-F32LE extent") + output = bytearray(len(payload) // 2) + for index, (sample,) in enumerate(struct.iter_unpack("= 1: + value = 32767 + else: + scale = 32768 if sample < 0 else 32767 + scaled = struct.unpack(" last or last > max(available): + raise ValueError("preparation CPUs must be within the process CPU allocation") + selected.update(range(first, last + 1)) + if not selected or not selected.issubset(available): + raise ValueError("preparation CPUs must be within the process CPU allocation") + return frozenset(selected) + + +def initialize_preparation_worker(cpus): + if cpus is not None: + os.sched_setaffinity(0, cpus) + + +@dataclass(eq=False) +class Request: + request_id: int + native_id: str + text: str + token_count: int = 0 + next_sequence: int = 0 + total_samples: int = 0 + submitted: bool = False + accepted: bool = False + canceled: bool = False + terminal: bool = False + scheduled: bool = False + preparation_task: asyncio.Task | None = None + preparation_started: bool = False + preparation_dispatched: bool = False + output: collections.deque = field(default_factory=collections.deque) + trace: dict = field(default_factory=dict) + + +class Session: + def __init__(self, bridge, reader, writer): + self.bridge = bridge + self.reader = reader + self.writer = writer + self.native_reader = None + self.native_writer = None + self.native_lock = asyncio.Lock() + self.requests = {} + self.native_requests = {} + self.retired = collections.deque(maxlen=4096) + self.ready = asyncio.Queue(maxsize=bridge.args.max_pending_requests + 1) + self.output_bytes = 0 + self.preparations = set() + self.running_preparations = set() + self.dispatched_preparations = set() + self.background = set() + self.slabs = SlabReader() + self.closed = False + + def spawn(self, coroutine): + task = asyncio.create_task(coroutine) + self.background.add(task) + task.add_done_callback(self.background.discard) + + def done(task): + if not task.cancelled() and task.exception() is not None: + self.writer.close() + if self.native_writer is not None: + self.native_writer.close() + + task.add_done_callback(done) + return task + + async def send_native(self, kind, **fields): + async with self.native_lock: + if self.native_writer is None or self.closed: + raise ConnectionError("native connection is unavailable") + async with asyncio.timeout(self.bridge.args.write_timeout): + await write_frame(self.native_writer, message(kind, **fields)) + + def enqueue(self, request, payload, *, terminal=False): + if self.closed or request.terminal: + return False + if ( + len(request.output) >= self.bridge.args.request_queue_chunks + or self.output_bytes + len(payload) > self.bridge.args.max_output_bytes + ): + return False + request.output.append((payload, terminal)) + self.output_bytes += len(payload) + request.terminal = terminal + if not request.scheduled: + request.scheduled = True + self.ready.put_nowait(request) + return True + + def discard_output(self, request): + while request.output: + self.output_bytes -= len(request.output.popleft()[0]) + + async def fail(self, request, code, detail, *, abort=True, aborted=False): + if request.terminal: + return + request.canceled = True + if ( + request.preparation_task is not None + and not request.preparation_started + and not request.preparation_dispatched + ): + request.preparation_task.cancel() + self.discard_output(request) + if aborted: + payload = control_frame("aborted", request_id=request.request_id) + else: + payload = control_frame( + "failed", + request_id=request.request_id, + code=code, + message=detail, + retryable=code in {"output_backpressure", "worker_stopping", "queue_full"}, + ) + if not self.enqueue(request, payload, terminal=True): + raise ConnectionError("client output budget exhausted") + if abort and request.submitted: + await self.send_native("abort", request_id=request.native_id) + + def retire(self, request): + self.mark(request, "retired") + if self.bridge.diagnostics is not None: + self.bridge.diagnostics.write( + json.dumps( + { + "request_id": request.request_id, + "native_id": request.native_id, + "events_ns": request.trace, + } + ) + + "\n" + ) + self.requests.pop(request.request_id, None) + self.native_requests.pop(request.native_id, None) + self.retired.append(request.native_id) + + def mark(self, request, event): + if self.bridge.diagnostics is not None and event not in request.trace: + request.trace[event] = time.monotonic_ns() + + async def write_output(self): + fair_batch = self.bridge.fair_batch + processed = 0 + while True: + if fair_batch and processed == fair_batch: + processed = 0 + await asyncio.sleep(0) + request = await self.ready.get() + processed += 1 + if not request.output: + request.scheduled = False + continue + payload, terminal = request.output.popleft() + self.output_bytes -= len(payload) + self.writer.write(payload) + async with asyncio.timeout(self.bridge.args.write_timeout): + await self.writer.drain() + if payload[4:5] in (b"A", b"F"): + self.mark(request, "first_audio_written") + request.scheduled = False + if terminal: + self.retire(request) + elif request.output: + request.scheduled = True + self.ready.put_nowait(request) + + def prepare_queued(self, request, value): + if request.canceled or self.closed: + return None + started = time.monotonic_ns() + text, tokens = self.bridge.prepare_text(value) + return started, time.monotonic_ns(), text, tokens + + async def prepare(self, request, value): + try: + if request.canceled or self.closed: + return + loop = asyncio.get_running_loop() + if value.get("kind") == "submit_prepared": + self.mark(request, "preparation_started") + text, tokens = value["text"], value["token_ids"] + self.mark(request, "preparation_finished") + elif self.bridge.continuous_preparation: + request.preparation_dispatched = True + self.dispatched_preparations.add(request.preparation_task) + try: + result = await loop.run_in_executor( + self.bridge.executor, self.prepare_queued, request, value + ) + finally: + self.dispatched_preparations.discard(request.preparation_task) + if result is None: + return + started, finished, text, tokens = result + if self.bridge.diagnostics is not None: + request.trace.setdefault("preparation_started", started) + request.trace.setdefault("preparation_finished", finished) + else: + async with self.bridge.preparation_slots: + if request.canceled or self.closed: + return + request.preparation_started = True + self.running_preparations.add(request.preparation_task) + self.mark(request, "preparation_started") + try: + text, tokens = await loop.run_in_executor( + self.bridge.executor, self.bridge.prepare_text, value + ) + finally: + self.running_preparations.discard(request.preparation_task) + self.mark(request, "preparation_finished") + if request.canceled or self.closed: + return + request.text = text + request.token_count = len(tokens) + request.submitted = True + await self.send_native( + "submit", + request_id=request.native_id, + session_id=None, + normalized_text=text, + token_ids=tokens, + speaker_id="english_sample", + max_tokens=512, + temperature=0.7, + top_k=80, + sample_rate_hz=22050, + seed=value.get("seed", 0), + ) + self.mark(request, "native_submitted") + except asyncio.CancelledError: + raise + except Exception as error: + await self.fail(request, "preparation_failed", str(error)) + + async def read_client(self): + fair_batch = self.bridge.fair_batch + processed = 0 + while True: + if fair_batch and processed == fair_batch: + processed = 0 + await asyncio.sleep(0) + value = await read_control(self.reader) + processed += 1 + kind = value.get("kind") + ident = value.get("request_id") + if type(ident) is not int or not 0 <= ident <= 0xFFFFFFFFFFFFFFFF: + raise ValueError("request_id must be a uint64 integer") + if kind in {"submit", "submit_prepared"}: + if ident in self.requests: + raise ValueError("request_id is already active") + if ( + len(self.requests) >= self.bridge.args.max_pending_requests + or len(self.preparations) >= self.bridge.args.max_pending_requests + ): + raise ValueError("client pending request limit exceeded") + request = Request(ident, uuid.uuid4().hex, "") + self.mark(request, "received") + self.requests[ident] = request + self.native_requests[request.native_id] = request + try: + self.bridge.validate_input(value) + except ValueError as error: + await self.fail(request, "invalid_request", str(error), abort=False) + continue + task = self.spawn(self.prepare(request, value)) + request.preparation_task = task + self.preparations.add(task) + task.add_done_callback(self.preparations.discard) + elif kind == "abort": + request = self.requests.get(ident) + if request is not None: + await self.fail(request, "aborted", "request aborted", aborted=True) + else: + raise ValueError("unsupported bridge client message") + + async def read_native(self): + fair_batch = self.bridge.fair_batch + processed = 0 + while True: + if fair_batch and processed == fair_batch: + processed = 0 + await asyncio.sleep(0) + value = await read_frame(self.native_reader) + processed += 1 + if value.get("protocol_version") != 1: + raise ValueError("native protocol version changed") + kind = value.get("kind") + ident = value.get("request_id") + request = self.native_requests.get(ident) + payload = None + if kind == "audio_slab": + if request is not None: + self.mark(request, "first_audio_received") + payload = self.slabs.copy(value) + await self.send_native( + "audio_ack", + arena=value["arena"], + slot=value["slot"], + generation=value["generation"], + ) + if request is None: + if ident in self.retired: + continue + raise ValueError("native output belongs to an unknown request") + if request.canceled or request.terminal: + continue + if kind == "accepted": + if request.accepted: + raise ValueError("native request was accepted twice") + request.accepted = True + if not self.enqueue( + request, + control_frame( + "accepted", + request_id=request.request_id, + processed_text=request.text, + token_count=request.token_count, + ), + ): + await self.fail(request, "output_backpressure", "client output queue is full") + elif kind == "audio_slab": + if not request.accepted or value["sequence"] != request.next_sequence: + raise ValueError("native audio sequence mismatch") + request.next_sequence += 1 + request.total_samples += value["samples"] + forwarded = ( + float_audio_frame(request.request_id, value["sequence"], payload) + if os.environ.get("NEMOTRON_BRIDGE_FORWARD_F32") == "1" + else audio_frame(request.request_id, value["sequence"], pcm16(payload)) + ) + if not self.enqueue(request, forwarded): + await self.fail(request, "output_backpressure", "client output queue is full") + elif kind == "stream_complete_v1": + if ( + type(value.get("next_sequence")) is not int + or type(value.get("chunk_count")) is not int + or type(value.get("total_samples")) is not int + or value["next_sequence"] != request.next_sequence + or value["chunk_count"] != request.next_sequence + or value["total_samples"] != request.total_samples + or request.next_sequence < 1 + or request.total_samples < 1 + ): + raise ValueError("native completion accounting mismatch") + if not self.enqueue( + request, + control_frame( + "complete", + request_id=request.request_id, + total_samples=request.total_samples, + chunk_count=request.next_sequence, + ), + terminal=True, + ): + await self.fail(request, "output_backpressure", "client output queue is full") + elif kind == "failed": + await self.fail( + request, + str(value.get("code", "backend_failed")), + str(value.get("message", "native request failed")), + abort=False, + ) + elif kind == "aborted": + await self.fail( + request, + "aborted", + "native request aborted", + abort=False, + aborted=True, + ) + else: + raise ValueError("unsupported native output message") + + async def run(self): + tasks = [] + try: + hello = await asyncio.wait_for(read_control(self.reader), 10) + if hello != {"kind": "hello", "protocol_version": 1}: + raise ValueError("invalid bridge handshake") + self.native_reader, self.native_writer = await asyncio.wait_for( + asyncio.open_unix_connection(str(self.bridge.args.native_worker_socket)), + self.bridge.args.connect_timeout, + ) + await self.send_native( + "client_hello", + expected_manifest_sha256=self.bridge.manifest_hash, + client_instance_id="nemo-speech-cpp-" + uuid.uuid4().hex, + ) + ready = await asyncio.wait_for(read_frame(self.native_reader), 10) + self.bridge.validate_ready(ready) + self.writer.write( + control_frame( + "ready", + model_name="nemotron-tts", + sample_rate=22050, + voices=["english_sample"], + batch_size=self.bridge.args.width, + mode=self.bridge.args.mode, + text_normalization_enabled=self.bridge.preparer.enabled, + **( + {"text_preparation": self.bridge.frontend_preparation} + if self.bridge.frontend_preparation is not None + else {} + ), + ) + ) + await asyncio.wait_for(self.writer.drain(), self.bridge.args.write_timeout) + tasks = [ + asyncio.create_task(fn()) + for fn in (self.read_client, self.read_native, self.write_output) + ] + done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + for task in done: + task.result() + except (asyncio.IncompleteReadError, ConnectionError, BrokenPipeError): + pass + except asyncio.CancelledError: + raise + except Exception: + LOG.exception("bridge connection failed") + finally: + self.closed = True + for task in (*tasks, *self.background): + if ( + task not in self.running_preparations + and task not in self.dispatched_preparations + ): + task.cancel() + await asyncio.gather(*tasks, *self.background, return_exceptions=True) + if self.native_writer is not None: + self.native_writer.close() + with contextlib.suppress(Exception): + await asyncio.wait_for(self.native_writer.wait_closed(), 5) + self.slabs.close() + self.writer.close() + with contextlib.suppress(Exception): + await asyncio.wait_for(self.writer.wait_closed(), 5) + + +class Bridge: + def __init__(self, args, *, tokenizer=None, preparer=None): + self.args = args + fair_batch = os.environ.get("NEMOTRON_BRIDGE_FAIR_BATCH", "0") + if re.fullmatch(r"[0-9]{1,4}", fair_batch) is None or int(fair_batch) > 1024: + raise ValueError("NEMOTRON_BRIDGE_FAIR_BATCH must be an integer from 0 through 1024") + self.fair_batch = int(fair_batch) + worker_cpus = preparation_cpus() + continuous = os.environ.get("NEMOTRON_BRIDGE_CONTINUOUS_PREPARATION", "0") + if continuous not in {"0", "1"}: + raise ValueError("NEMOTRON_BRIDGE_CONTINUOUS_PREPARATION must be 0 or 1") + self.continuous_preparation = continuous == "1" + frontend = os.environ.get("NEMOTRON_FRONTEND_PREPARATION", "0") + if frontend not in {"0", "1"}: + raise ValueError("NEMOTRON_FRONTEND_PREPARATION must be 0 or 1") + self.frontend_preparation = None + raw = args.manifest.read_bytes() + self.manifest_hash = hashlib.sha256(raw).hexdigest() + self.manifest = json.loads(raw) + if self.manifest["physical_width"] != args.width: + raise ValueError("bridge width differs from the native serving plan") + if self.manifest["host"].get("batching_mode", "fixed_cohort") != args.mode: + raise ValueError("bridge mode differs from the native serving plan") + if self.manifest["transport_compatibility"]["output_contract"] != ( + "pcm_f32_le_22050_mono_variable_streaming_v3" + ): + raise ValueError("bridge requires the native streaming-v3 output contract") + if tokenizer is None: + seal_path = args.model_seal or args.model / "model-seal.json" + seal = json.loads(seal_path.read_text()) + tokenizer_path = (args.model / "tokenizer.json").resolve(strict=True) + tokenizer_seal = seal["files"]["tokenizer.json"] + if ( + seal["model_identity"] != self.manifest["model"]["checkpoint_sha256"] + or not tokenizer_path.is_relative_to(args.model.resolve(strict=True)) + or tokenizer_path.stat().st_size != tokenizer_seal["bytes"] + or hashlib.sha256(tokenizer_path.read_bytes()).hexdigest() + != tokenizer_seal["sha256"] + ): + raise ValueError("tokenizer does not match the selected sealed model") + preparation_library = os.environ.get("NEMOTRON_NATIVE_PREPARATION_LIBRARY") + if frontend == "1" and not preparation_library: + raise ValueError("frontend preparation requires the native preparation library") + if preparation_library and preparer is None: + if frontend == "1": + from .native_preparation import FrontendPreparation + + preparer = FrontendPreparation( + preparation_library, tokenizer_path, args.normalization_config + ) + self.frontend_preparation = preparer.metadata + else: + from .native_preparation import NativePreparation + + preparer = NativePreparation( + preparation_library, tokenizer_path, args.normalization_config + ) + else: + from tokenizers import Tokenizer + + tokenizer = Tokenizer.from_file(str(tokenizer_path)) + if preparer is None: + from .normalization import TextPreparer + + preparer = TextPreparer(args.normalization_config) + self.tokenizer = tokenizer + self.preparer = preparer + self.executor = ThreadPoolExecutor( + max_workers=args.preparation_threads, + thread_name_prefix="native-text", + initializer=initialize_preparation_worker, + initargs=(worker_cpus,), + ) + self.preparation_slots = asyncio.Semaphore(args.preparation_threads) + self.session = None + self.tasks = set() + diagnostic_path = os.environ.get("NEMOTRON_BRIDGE_DIAGNOSTICS") + self.diagnostics = open(diagnostic_path, "x") if diagnostic_path else None + + def validate_input(self, value): + prepared = value.get("kind") == "submit_prepared" + if prepared != (self.frontend_preparation is not None): + raise ValueError("request preparation does not match the negotiated mode") + if prepared: + tokens = value.get("token_ids") + if ( + not isinstance(tokens, list) + or not 1 <= len(tokens) <= 512 + or any(type(token) is not int or not 0 <= token < 131075 for token in tokens) + or tokens[-1] != 131073 + ): + raise ValueError("invalid prepared token IDs") + self.validate_submit({key: item for key, item in value.items() if key != "token_ids"}) + else: + self.validate_submit(value) + + @staticmethod + def validate_submit(value): + allowed = { + "kind", + "protocol_version", + "request_id", + "text", + "voice", + "language", + "seed", + "apply_normalize", + } + if set(value) - allowed: + raise ValueError("unsupported submit fields") + text = value.get("text") + if not isinstance(text, str) or not text.strip() or len(text.encode("utf-8")) > 1_000_000: + raise ValueError("text must contain 1..1000000 UTF-8 bytes") + if value.get("voice", "english_sample") != "english_sample": + raise ValueError("native TTS supports only voice english_sample") + if value.get("language", "en-US") not in ( + "en", + "en-US", + "english", + "English", + "", + ): + raise ValueError("the native package supports English only") + seed = value.get("seed", 0) + if type(seed) is not int or not 0 <= seed <= 0xFFFFFFFFFFFFFFFF: + raise ValueError("seed must be a uint64 integer") + if type(value.get("apply_normalize", True)) is not bool: + raise ValueError("apply_normalize must be boolean") + + def prepare_text(self, value): + if hasattr(self.preparer, "prepare_tokens"): + text, tokens = self.preparer.prepare_tokens( + value["text"], value.get("apply_normalize", True) + ) + else: + text = self.preparer.prepare(value["text"], "en-US", value.get("apply_normalize", True)) + tokens = self.tokenizer.encode_batch([text], add_special_tokens=False)[0].ids + [131073] + if not isinstance(text, str) or not text.strip(): + raise ValueError("text preparation returned empty text") + if not 1 <= len(tokens) <= 512: + raise ValueError("native TTS supports at most 512 text tokens including EOS") + if any(type(token) is not int or not 0 <= token < 131075 for token in tokens): + raise ValueError("tokenizer returned an out-of-vocabulary token") + return text, tokens + + def validate_ready(self, value): + expected = self.manifest["transport_compatibility"] + if ( + value.get("kind") != "worker_ready" + or value.get("protocol_version") != 1 + or value.get("manifest_sha256") != self.manifest_hash + or value.get("model_sha256") != self.manifest["model"]["checkpoint_sha256"] + or value.get("output_contract") != expected["output_contract"] + or value.get("enabled_buckets") != sorted({1, self.args.width}) + or any( + value.get(key) != item + for key, item in expected["installed_transport_versions"].items() + ) + ): + raise ValueError("native worker identity differs from the selected package") + + async def accept(self, reader, writer): + if self.session is not None: + writer.close() + await writer.wait_closed() + return + self.session = Session(self, reader, writer) + task = asyncio.current_task() + self.tasks.add(task) + try: + await self.session.run() + finally: + self.session = None + self.tasks.discard(task) + + async def close(self): + for task in self.tasks: + task.cancel() + await asyncio.gather(*self.tasks, return_exceptions=True) + await asyncio.to_thread(self.executor.shutdown, wait=True, cancel_futures=True) + self.preparer.close() + if self.diagnostics is not None: + self.diagnostics.close() + + +def parse_args(argv=None): + parser = argparse.ArgumentParser(description="C++ frontend bridge for native Nemotron TTS") + for name in ("native-worker-socket", "manifest", "model", "listen"): + parser.add_argument("--" + name, type=Path, required=True) + parser.add_argument( + "--mode", + choices=("fixed_cohort", "fixed_request", "inflight"), + default="inflight", + ) + parser.add_argument("--width", type=int, required=True) + parser.add_argument("--normalization-config", type=Path) + parser.add_argument("--model-seal", type=Path) + parser.add_argument("--preparation-threads", type=int, default=16) + parser.add_argument("--max-pending-requests", type=int, default=4096) + parser.add_argument("--request-queue-chunks", type=int, default=32) + parser.add_argument("--max-output-bytes", type=int, default=64 * 1024 * 1024) + parser.add_argument("--write-timeout", type=float, default=10) + parser.add_argument("--connect-timeout", type=float, default=10) + args = parser.parse_args(argv) + if args.width not in [1 << exponent for exponent in range(13)]: + parser.error("width must be a supported power of two from 1 through 4096") + for key in ( + "preparation_threads", + "max_pending_requests", + "request_queue_chunks", + "max_output_bytes", + "write_timeout", + "connect_timeout", + ): + if getattr(args, key) <= 0: + parser.error(key + " must be positive") + return args + + +async def serve(args): + bridge = Bridge(args) + server = None + created = False + try: + if args.listen.exists() or args.listen.is_symlink(): + raise ValueError("refusing to replace an existing bridge socket") + args.listen.parent.mkdir(parents=True, exist_ok=True) + server = await asyncio.start_unix_server(bridge.accept, path=str(args.listen)) + created = True + args.listen.chmod(0o600) + stopping = asyncio.Event() + loop = asyncio.get_running_loop() + for sig in (signal.SIGINT, signal.SIGTERM): + loop.add_signal_handler(sig, stopping.set) + print( + json.dumps({"status": "BRIDGE_LISTENING", "socket": str(args.listen)}), + flush=True, + ) + await stopping.wait() + finally: + if server is not None: + server.close() + await server.wait_closed() + await bridge.close() + if created: + args.listen.unlink(missing_ok=True) + + +def main(argv=None): + logging.basicConfig(level=logging.INFO) + print( + json.dumps({"serialization": install_serialization(native_protocol)}), + flush=True, + ) + asyncio.run(serve(parse_args(argv))) + + +if __name__ == "__main__": + main() diff --git a/python/nemotron_cpp_bridge/slabs.py b/python/nemotron_cpp_bridge/slabs.py new file mode 100644 index 0000000..e53c612 --- /dev/null +++ b/python/nemotron_cpp_bridge/slabs.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import mmap +import os +import re +import stat + + +class SlabReader: + SLOT_BYTES = 8 * 1764 * 4 + + def __init__(self, directory="/dev/shm"): + self.directory = directory + self.arena = None + self.mapping = None + self.generations = {} + + def copy(self, value): + integer_fields = ( + "slab_version", + "slot", + "generation", + "arena_bytes", + "slot_bytes", + "offset", + "data_bytes", + "samples", + "sequence", + "sample_rate_hz", + "channels", + ) + if any(type(value.get(key)) is not int for key in integer_fields): + raise ValueError("slab metadata requires integer fields") + arena = value.get("arena") + if not isinstance(arena, str) or re.fullmatch(r"r20-pcm-[0-9a-f]{32}", arena) is None: + raise ValueError("invalid PCM arena name") + extent = value["arena_bytes"] + slot = value["slot"] + size = value["data_bytes"] + if ( + value["slab_version"] != 1 + or value["slot_bytes"] != self.SLOT_BYTES + or not 0 < extent <= 4096 * self.SLOT_BYTES + or extent % self.SLOT_BYTES + or not 0 <= slot < extent // self.SLOT_BYTES + or value["generation"] <= self.generations.get(slot, 0) + or value["offset"] != slot * self.SLOT_BYTES + or not 0 < size <= self.SLOT_BYTES + or size != value["samples"] * 4 + or value["offset"] + size > extent + or value["sample_rate_hz"] != 22050 + or value["channels"] != 1 + or value.get("encoding") != "pcm_f32_le" + or value.get("is_final") is not False + or not 0 <= value["sequence"] <= 0xFFFFFFFF + ): + raise ValueError("invalid PCM slab bounds or audio contract") + if self.mapping is None: + fd = os.open( + os.path.join(self.directory, arena), os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW + ) + try: + info = os.fstat(fd) + if not stat.S_ISREG(info.st_mode) or info.st_size != extent: + raise ValueError("PCM arena file does not match its descriptor") + self.mapping = mmap.mmap(fd, extent, access=mmap.ACCESS_READ) + self.arena = arena + finally: + os.close(fd) + if arena != self.arena or len(self.mapping) != extent: + raise ValueError("PCM arena identity changed on an active connection") + self.generations[slot] = value["generation"] + return self.mapping[value["offset"] : value["offset"] + size] + + def close(self): + if self.mapping is not None: + self.mapping.close() + self.mapping = None diff --git a/python/nemotron_native_backend/__init__.py b/python/nemotron_native_backend/__init__.py new file mode 100644 index 0000000..0cb67e3 --- /dev/null +++ b/python/nemotron_native_backend/__init__.py @@ -0,0 +1,3 @@ +from .models import AudioResult, SynthesisRequest + +__all__ = ["AudioResult", "SynthesisRequest"] diff --git a/python/nemotron_native_backend/backend.py b/python/nemotron_native_backend/backend.py new file mode 100644 index 0000000..5ab5fcc --- /dev/null +++ b/python/nemotron_native_backend/backend.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +from .models import AudioResult, SynthesisRequest + + +class BackendError(RuntimeError): + def __init__(self, code: str, message: str, *, retryable: bool = False) -> None: + super().__init__(message) + self.code = code + self.retryable = retryable + + +@dataclass(frozen=True, slots=True) +class BackendIdentity: + cuda_device_name: str + cuda_compute_capability: tuple[int, int] + enabled_buckets: tuple[int, ...] + real_weights: bool + execution_contract: str = "unspecified" + + +class SynthesisBackend(ABC): + @property + @abstractmethod + def identity(self) -> BackendIdentity: + raise NotImplementedError + + @abstractmethod + async def synthesize_batch(self, requests: Sequence[SynthesisRequest]) -> list[AudioResult]: + raise NotImplementedError + + def consume_execution_evidence( + self, requests: Sequence[SynthesisRequest], results: Sequence[AudioResult] + ) -> dict[str, Any]: + raise BackendError("backend_evidence", "backend did not provide execution evidence") + + async def shutdown(self) -> None: + return None + + +def exact16_execution_evidence( + identity: BackendIdentity, + requests: Sequence[SynthesisRequest], + results: Sequence[AudioResult], +) -> dict[str, Any]: + return { + "real_weights": identity.real_weights, + "execution_contract": identity.execution_contract, + "request_ids": tuple(request.request_id for request in requests), + "result_ids": tuple(result.request_id for result in results), + "predictor_plan_slices": (len(requests),), + "stage1_plan_slices": (len(requests),), + "decode_steps": 15, + "predictor_plan_invocations": 15, + "predictor_passes_per_lane": 240, + "padding_count": 0, + "predictor_fallback_count": 0, + "stage1_fallback_count": 0, + "code_h2d_copies": 0, + "code_d2h_copies": 0, + } + + +def validate_exact16_execution_evidence( + identity: BackendIdentity, + requests: Sequence[SynthesisRequest], + results: Sequence[AudioResult], + evidence: Mapping[str, Any], +) -> None: + expected = exact16_execution_evidence(identity, requests, results) + if dict(evidence) != expected: + raise BackendError("backend_evidence", "per-invocation exact16 evidence drift") + + +def backend_identity_dict(identity: BackendIdentity) -> dict[str, Any]: + return { + "cuda_device_name": identity.cuda_device_name, + "cuda_compute_capability": list(identity.cuda_compute_capability), + "enabled_buckets": list(identity.enabled_buckets), + "real_weights": identity.real_weights, + "execution_contract": identity.execution_contract, + } diff --git a/python/nemotron_native_backend/batcher.py b/python/nemotron_native_backend/batcher.py new file mode 100644 index 0000000..45b7762 --- /dev/null +++ b/python/nemotron_native_backend/batcher.py @@ -0,0 +1,300 @@ +from __future__ import annotations + +import asyncio +from collections import defaultdict, deque +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +from .backend import BackendError, SynthesisBackend, validate_exact16_execution_evidence +from .models import AudioResult, SynthesisRequest +from .receipts import ReceiptWriter + + +class QueueFull(BackendError): + def __init__(self) -> None: + super().__init__("queue_full", "worker request queue is full", retryable=True) + + +class RequestAborted(BackendError): + def __init__(self) -> None: + super().__init__("aborted", "request was aborted", retryable=False) + + +@dataclass(slots=True) +class _Pending: + request: SynthesisRequest + future: asyncio.Future[AudioResult] + state: str = "queued" + aborted: bool = False + forced_error: BaseException | None = None + + +def exact_partitions(count: int, enabled_buckets: Sequence[int]) -> tuple[list[int], int]: + + if count < 0: + raise ValueError("count must be nonnegative") + remainder = count + result: list[int] = [] + for bucket in sorted({int(value) for value in enabled_buckets}, reverse=True): + if bucket < 1: + raise ValueError("bucket sizes must be positive") + while remainder >= bucket: + result.append(bucket) + remainder -= bucket + return result, remainder + + +class ExactBatcher: + def __init__( + self, + backend: SynthesisBackend, + *, + batch_window_ms: float, + cohort_timeout_ms: float = 1000.0, + expected_cohort_size: int | None = None, + max_queue_depth: int, + max_inflight_batches: int = 1, + receipt_writer: ReceiptWriter | None = None, + ) -> None: + self.backend = backend + self.batch_window_seconds = float(batch_window_ms) / 1000.0 + self.cohort_timeout_seconds = float(cohort_timeout_ms) / 1000.0 + self.expected_cohort_size = expected_cohort_size + if ( + expected_cohort_size is not None + and expected_cohort_size not in backend.identity.enabled_buckets + ): + raise ValueError("expected cohort size has no enabled exact backend plan") + self.max_queue_depth = int(max_queue_depth) + self.max_inflight_batches = int(max_inflight_batches) + if self.max_inflight_batches not in (1, 2): + raise ValueError("max_inflight_batches must be one or two") + self.receipt_writer = receipt_writer + self._queue: deque[_Pending] = deque() + self._pending: dict[str, _Pending] = {} + self._wake = asyncio.Event() + self._runner: asyncio.Task[None] | None = None + self._dispatch_tasks: set[asyncio.Task[None]] = set() + self._dispatch_slots = asyncio.Semaphore(self.max_inflight_batches) + self._closed = False + self._active_batches = 0 + + def start(self) -> None: + if self._runner is not None: + raise RuntimeError("batcher already started") + self._runner = asyncio.create_task(self._run(), name="nemotron-exact-batcher") + + def submit(self, request: SynthesisRequest) -> asyncio.Future[AudioResult]: + if self._closed: + raise BackendError("worker_stopping", "worker is stopping", retryable=True) + if request.request_id in self._pending: + raise BackendError("duplicate_request_id", "request ID is already pending") + if len(self._pending) >= self.max_queue_depth: + raise QueueFull() + loop = asyncio.get_running_loop() + pending = _Pending(request=request, future=loop.create_future()) + self._pending[request.request_id] = pending + self._queue.append(pending) + self._wake.set() + return pending.future + + def abort(self, request_id: str) -> bool: + pending = self._pending.get(request_id) + if pending is None: + return False + pending.aborted = True + if pending.state == "queued" and not pending.future.done(): + pending.future.set_exception(RequestAborted()) + self._pending.pop(request_id, None) + self._wake.set() + return True + + @property + def queue_depth(self) -> int: + return sum(1 for item in self._queue if not item.aborted) + + @property + def active_batches(self) -> int: + return self._active_batches + + @property + def dispatch_task_count(self) -> int: + return len(self._dispatch_tasks) + + async def _run(self) -> None: + try: + while not self._closed: + await self._wake.wait() + self._wake.clear() + deadline = asyncio.get_running_loop().time() + self.cohort_timeout_seconds + if self.batch_window_seconds: + await asyncio.sleep(self.batch_window_seconds) + while self.expected_cohort_size is not None: + counts: dict[tuple[Any, ...], int] = defaultdict(int) + for item in self._queue: + if not item.aborted: + counts[item.request.compatibility_key] += 1 + if max(counts.values(), default=0) >= self.expected_cohort_size: + break + remaining = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + break + self._wake.clear() + try: + await asyncio.wait_for(self._wake.wait(), timeout=remaining) + except TimeoutError: + break + snapshot: list[_Pending] = [] + while self._queue: + item = self._queue.popleft() + if item.aborted: + continue + snapshot.append(item) + groups: dict[tuple[Any, ...], list[_Pending]] = defaultdict(list) + for item in snapshot: + groups[item.request.compatibility_key].append(item) + for compatible in groups.values(): + self._launch_compatible(compatible) + if self._queue: + self._wake.set() + except asyncio.CancelledError: + raise + except Exception as error: + for item in list(self._pending.values()): + if not item.future.done(): + item.future.set_exception(error) + self._pending.clear() + raise + + def _launch_compatible(self, items: list[_Pending]) -> None: + partitions, remainder = exact_partitions(len(items), self.backend.identity.enabled_buckets) + cursor = 0 + for batch in partitions: + cohort = items[cursor : cursor + batch] + cursor += batch + for item in cohort: + item.state = "scheduled" + task = asyncio.create_task( + self._dispatch(cohort), + name=f"nemotron-exact-dispatch-b{batch}", + ) + self._dispatch_tasks.add(task) + task.add_done_callback(self._dispatch_done) + if remainder: + error = BackendError( + "unsupported_exact_batch", + f"{remainder} request(s) cannot be represented by enabled exact plans " + f"{list(self.backend.identity.enabled_buckets)}", + retryable=True, + ) + for item in items[cursor:]: + self._finish_error(item, error) + + def _dispatch_done(self, task: asyncio.Task[None]) -> None: + self._dispatch_tasks.discard(task) + if task.cancelled(): + return + error = task.exception() + if error is not None and self._runner is not None and not self._runner.done(): + self._runner.cancel(error) + + async def _dispatch(self, cohort: list[_Pending]) -> None: + async with self._dispatch_slots: + for item in cohort: + item.state = "active" + self._active_batches += 1 + try: + if all(item.aborted or item.forced_error is not None for item in cohort): + for item in cohort: + terminal = self._terminal_error(item) + if terminal is None: + raise BackendError( + "batcher_state", "terminal cohort has no terminal error" + ) + self._finish_error(item, terminal) + return + requests = [item.request for item in cohort] + results = await self.backend.synthesize_batch(requests) + evidence = self.backend.consume_execution_evidence(requests, results) + validate_exact16_execution_evidence( + self.backend.identity, requests, results, evidence + ) + if len(results) != len(cohort): + raise BackendError( + "backend_contract", + f"backend returned {len(results)} results for B{len(cohort)}", + ) + by_id = {result.request_id: result for result in results} + if len(by_id) != len(results) or set(by_id) != { + item.request.request_id for item in cohort + }: + raise BackendError("backend_contract", "backend result IDs differ from cohort") + for result in results: + result.validate() + if self.receipt_writer is not None: + await self.receipt_writer.write_success(requests, evidence) + for item in cohort: + terminal = self._terminal_error(item) + if terminal is not None: + self._finish_error(item, terminal) + else: + self._finish_result(item, by_id[item.request.request_id]) + except asyncio.CancelledError as error: + for item in cohort: + terminal = self._terminal_error(item, default=error) + self._finish_error(item, terminal or error) + raise + except Exception as error: + for item in cohort: + terminal = self._terminal_error(item, default=error) + if terminal is None: + raise BackendError( + "batcher_state", "failed cohort has no terminal error" + ) from error + self._finish_error(item, terminal) + finally: + self._active_batches -= 1 + + @staticmethod + def _terminal_error( + item: _Pending, *, default: BaseException | None = None + ) -> BaseException | None: + if item.forced_error is not None: + return item.forced_error + if item.aborted: + return RequestAborted() + return default + + def _finish_result(self, item: _Pending, result: AudioResult) -> None: + self._pending.pop(item.request.request_id, None) + item.state = "done" + if not item.future.done(): + item.future.set_result(result) + + def _finish_error(self, item: _Pending, error: BaseException) -> None: + self._pending.pop(item.request.request_id, None) + item.state = "done" + if not item.future.done(): + item.future.set_exception(error) + + async def shutdown(self) -> None: + self._closed = True + self._wake.set() + if self._runner is not None: + self._runner.cancel() + try: + await self._runner + except asyncio.CancelledError: + pass + stopping = BackendError("worker_stopping", "worker stopped", retryable=True) + for item in list(self._pending.values()): + item.forced_error = stopping + if item.state == "queued": + self._finish_error(item, stopping) + if self._dispatch_tasks: + await asyncio.gather(*tuple(self._dispatch_tasks), return_exceptions=True) + for item in list(self._pending.values()): + self._finish_error(item, stopping) + self._pending.clear() + await self.backend.shutdown() diff --git a/python/nemotron_native_backend/bundle.py b/python/nemotron_native_backend/bundle.py new file mode 100644 index 0000000..6ed66a8 --- /dev/null +++ b/python/nemotron_native_backend/bundle.py @@ -0,0 +1,197 @@ +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path, PurePosixPath +from types import MappingProxyType +from typing import Any + + +class BundleError(ValueError): + pass + + +@dataclass(frozen=True) +class Artifact: + name: str + relative_path: str + sha256: str + + +@dataclass(frozen=True) +class NativeBundlePlan: + identity: str + physical_width: int + artifacts: tuple[Artifact, ...] + contract: Mapping[str, Any] + manifest_sha256: str + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise BundleError(message) + + +def _freeze(value: Any) -> Any: + if isinstance(value, dict): + return MappingProxyType({key: _freeze(item) for key, item in value.items()}) + if isinstance(value, list): + return tuple(_freeze(item) for item in value) + return value + + +def _relative_path(value: Any) -> str: + _require(isinstance(value, str) and bool(value), "Artifact path must be nonempty") + path = PurePosixPath(value) + _require( + not path.is_absolute() and ".." not in path.parts, "Artifacts require bundle-relative paths" + ) + _require( + "\\" not in value and "\x00" not in value and path.as_posix() == value, + "Artifact path must be canonical POSIX syntax", + ) + return value + + +def _sha256(value: Any) -> str: + _require( + isinstance(value, str) + and len(value) == 64 + and all(character in "0123456789abcdef" for character in value), + "Expected lowercase SHA-256", + ) + return value + + +def load_plan(path: Path) -> NativeBundlePlan: + payload = Path(path).read_bytes() + value = json.loads(payload) + _require(isinstance(value, dict), "Bundle must be an object") + _require(value.get("schema") == "nemotron.native-package-plan.r22.v1", "Unknown schema") + _require( + value.get("status") == "PROPOSED_UNQUALIFIED" and value.get("enabled") is False, + "This preparatory loader cannot enable serving", + ) + width = value["physical_width"] + _require(type(width) is int and width > 0, "Physical width must be positive") + _require(value["neural_cohort_partitioning"] is False, "Physical cohorts cannot split") + artifact_rows = value["artifacts"] + artifacts = tuple( + Artifact(name, _relative_path(row["path"]), _sha256(row["sha256"])) + for name, row in sorted(artifact_rows.items()) + ) + names = {item.name for item in artifacts} + _require( + len({item.relative_path for item in artifacts}) == len(artifacts), + "Artifact destinations must be unique", + ) + kernels = value["kernels"] + _require(set(kernels) == {"prefill", "stage0", "codec"}, "Three neural bindings required") + for region, kernel in kernels.items(): + _require(kernel["artifact"] in names, "Missing kernel artifact: " + region) + _require(kernel["physical_width"] == width, "Kernel physical widths must agree") + _require( + isinstance(kernel["entrypoint"], str) and kernel["entrypoint"].isidentifier(), + "Invalid kernel entrypoint", + ) + generation = value["generation"] + _require( + generation["cap"] == 511 + and generation["startup_skip"] is True + and generation["chunk_frames"] == [3, 5, 8], + "Unsupported generation contract", + ) + _require(generation["contraction"] is False, "This imported wide plan has no contraction") + publication = value["publication"] + _require( + publication["continuation_proof_required"] is True + and publication["sample_rate"] == 22050 + and publication["neural_output"] == "pcm_f32_le" + and publication["http_output"] == "pcm_s16_le", + "Changed publication contract", + ) + state = value["state"] + _require( + state["mode"] == 1 + and state["format"] == "canonical-FP32-private-BF16" + and state["recurrence_arithmetic"] == "FP32", + "Unsupported imported state format", + ) + _require( + state["scales_artifact"] in names and state["map_artifact"] in names, + "State bindings require explicit scale and map artifacts", + ) + execution = value["execution"] + _require( + execution["mode"] == "serial" and execution["grid_policy"] == "overlap-only", + "Imported plan is qualified only for serial execution", + ) + for key in ("stage0_grid", "full_stage0_grid", "codec_grid"): + _require(type(execution[key]) is int and execution[key] > 0, "Invalid resource count") + _require( + execution["stage0_grid"] == execution["full_stage0_grid"], + "Serial producer allocation differs from compiled participant count", + ) + _require( + kernels["stage0"]["barrier_participants"] == execution["stage0_grid"], + "Producer barrier participants differ from launch", + ) + _sha256(value["model"]["checkpoint_sha256"]) + host = value["host"] + _require( + host["tokens"] == "pinned" and host["serialization"] == "msgpack-c-1.2.2", + "Unrecognized imported host path", + ) + _require( + host["pcm_conversion"] == "exact-avx2-with-scalar-fallback", + "SIMD changes require a separately qualified plan", + ) + _require( + value["transport_compatibility"]["identity_migration_enabled"] is False, + "Legacy ready identity must remain compatible until qualification", + ) + return NativeBundlePlan( + value["id"], width, artifacts, _freeze(value), hashlib.sha256(payload).hexdigest() + ) + + +def resolve_artifacts(plan: NativeBundlePlan, artifact_root: Path) -> Mapping[str, Path]: + root = Path(artifact_root).resolve(strict=True) + resolved = {} + for artifact in plan.artifacts: + path = (root / artifact.relative_path).resolve(strict=True) + _require(path.is_relative_to(root) and path.is_file(), "Artifact escapes bundle root") + digest = hashlib.sha256() + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1 << 20), b""): + digest.update(block) + _require(digest.hexdigest() == artifact.sha256, "Artifact hash mismatch: " + artifact.name) + resolved[artifact.name] = path + return MappingProxyType(resolved) + + +def create_engine( + plan: NativeBundlePlan, + artifact_root: Path, + model_root: Path, + *, + source_pins: Mapping[str, str] | None = None, + model_identity: str | None = None, + qualification_only: bool = False, +) -> Any: + if qualification_only is not True or source_pins is None or model_identity is None: + raise BundleError( + "Production serving remains disabled; explicit qualification bindings are required" + ) + from .execution.composition import create_for_qualification + + return create_for_qualification( + plan, + artifact_root, + model_root, + source_pins=source_pins, + model_identity=model_identity, + qualification_only=True, + ) diff --git a/python/nemotron_native_backend/execution/__init__.py b/python/nemotron_native_backend/execution/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/nemotron_native_backend/execution/chunk_boundary.py b/python/nemotron_native_backend/execution/chunk_boundary.py new file mode 100644 index 0000000..cb00282 --- /dev/null +++ b/python/nemotron_native_backend/execution/chunk_boundary.py @@ -0,0 +1,9 @@ +def pack_and_enqueue(window, binding, start_event, stream): + window.pack() + start_event.record(stream) + return binding.enqueue() + + +def finish_and_observe(codec, ticket, window, done): + result = codec.finish(ticket, copy_pcm=True) + return (result, window.lengths.cpu().tolist(), window.stops.cpu().tolist(), done.cpu().tolist()) diff --git a/python/nemotron_native_backend/execution/chunk_submission.py b/python/nemotron_native_backend/execution/chunk_submission.py new file mode 100644 index 0000000..1a095d7 --- /dev/null +++ b/python/nemotron_native_backend/execution/chunk_submission.py @@ -0,0 +1,44 @@ +def install(e): + t = e.t + c = e.core + n = c.neural + remaining = t.empty_like(e.counts) + + def submit(F, counts, bank): + if e._codec_pending is not None: + raise RuntimeError("codec buffer/scratch lease still outstanding") + (packed, lengths, index, offsets, mask, _, _, banks) = e.bindings[F] + (ready, a, b) = e._event_pool[len(e._chunk_events)] + sizes = [ + min(F, available - used) + for (available, used) in zip(counts, e.consumed_host, strict=False) + ] + with t.cuda.stream(n.stream): + t.sub(e.counts, e.consumed, out=remaining) + remaining.clamp_(min=0, max=F) + lengths.copy_(remaining) + t.add(e.consumed.view(e.batch, 1), offsets, out=index) + index.clamp_(max=e.steps - 1) + t.gather(e.queue, 1, index.unsqueeze(-1).expand(e.batch, F, 16), out=packed) + t.lt(offsets, lengths.view(e.batch, 1), out=mask) + packed.mul_(mask.unsqueeze(-1)) + e.consumed.add_(lengths) + ready.record(n.stream) + e.consumed_host = [ + used + size for (used, size) in zip(e.consumed_host, sizes, strict=False) + ] + e.codec_stream.wait_event(ready) + with t.cuda.stream(e.codec_stream): + a.record(e.codec_stream) + banks[bank].event = b + ticket = banks[bank].enqueue() + pending = dict(ticket=ticket, sizes=sizes, events=(a, b)) + e._codec_pending = pending + e._chunk_events.append((a, b)) + return pending + + original = e._submit_chunk + if hasattr(e, "_r8_pcm_owners"): + raise RuntimeError("chunk binding must precede pinned transfer binding") + e._submit_chunk = submit + e._r9_chunk_owners = (remaining, original) diff --git a/python/nemotron_native_backend/execution/composition.py b/python/nemotron_native_backend/execution/composition.py new file mode 100644 index 0000000..025f6c0 --- /dev/null +++ b/python/nemotron_native_backend/execution/composition.py @@ -0,0 +1,287 @@ +from __future__ import annotations + +import copy +import ctypes as C +import hashlib +import json +from collections.abc import Mapping +from pathlib import Path + +from ..bundle import BundleError, NativeBundlePlan, resolve_artifacts +from . import ( + chunk_submission, + core, + helper_graphs, + pcm_transfer, + prefix_cache, + state_binding, + token_staging, +) +from .engine import NativeCohortEngine +from .lane_control import update_first_eos +from .native_assets import ModelAssets +from .native_inputs import requests_from_texts +from .source_build import is_source_build, verify_source_build + + +class PackageSupport: + ModelAssets = ModelAssets + Cohort = core.Cohort + requests_from_texts = staticmethod(requests_from_texts) + update_first_eos = staticmethod(update_first_eos) + + def __init__(self, plan, artifacts, source_pins): + self.plan = plan + self.artifacts = artifacts + self.source_pins = dict(source_pins) + self.source_sha256 = hashlib.sha256(Path(core.__file__).read_bytes()).hexdigest() + + def verify_sources(self): + package = Path(__file__).resolve().parent.parent + if not self.source_pins: + raise BundleError("Frozen package source pins are required") + actual = {str(path.relative_to(package)) for path in package.rglob("*.py")} + if set(self.source_pins) != actual: + raise BundleError("Package source inventory differs from the qualification seal") + for relative, expected in self.source_pins.items(): + target = (package / relative).resolve(strict=True) + if ( + not target.is_relative_to(package) + or hashlib.sha256(target.read_bytes()).hexdigest() != expected + ): + raise BundleError("Package source differs from qualification seal: " + relative) + + def select(self, batch): + if batch != self.plan.physical_width: + raise BundleError("Requested physical width differs from the kernel binding") + paths = { + name: self.artifacts[key] + for name, key in (("prefill", "prefill"), ("stage0", "stage0"), ("stage1", "codec")) + } + return paths, dict( + physical_batch=batch, + variant=self.plan.identity, + neural_cohort_partitioning=False, + bundle_manifest_sha256=self.plan.manifest_sha256, + ) + + +def verify_compiled_contract(plan, artifacts): + value = plan.contract + built = verify_source_build(plan, artifacts) if is_source_build(value) else None + binding = ( + built["compiled_binding"] + if built is not None + else json.loads(artifacts["compiled_binding"].read_text()) + ) + if ( + binding["batch"] != plan.physical_width + or binding["physical_batch"] != plan.physical_width + or binding["neural_cohort_partitioning"] is not False + or binding["state_mode"] != value["state"]["mode"] + or binding["state_format"] != value["state"]["format"] + or binding["ctas"] != value["execution"]["stage0_grid"] + or binding["prefill_workspace_sequences"] + != value["kernels"]["prefill"]["workspace_sequences"] + ): + raise BundleError("Compiled binding does not satisfy the proposed program") + expected_options = { + "binding_module": "state_r9.binding", + "batched_handoff": True, + "pinned_pcm": True, + "early_proof": True, + "engine_binding": "integration_r9.binding", + "engine_bindings": ["control_r9.chunk", "control_r9.prefix"], + "transport_binding": "publication.wide_delivery", + } + if value["kernels"]["prefill"].get("validation_kind") == "retained-device-dedup": + expected_options.pop("transport_binding") + if binding["runtime_options"] != expected_options: + raise BundleError("Compiled binding requires unported construction hooks") + if built is None: + for region in ("prefill", "stage0", "codec"): + receipt = json.loads(artifacts[region + "_build_receipt"].read_text()) + if ( + receipt["library_sha256"] + != hashlib.sha256(artifacts[region].read_bytes()).hexdigest() + ): + raise BundleError("Library and immutable build receipt disagree: " + region) + recipe = json.loads(artifacts["state_recipe"].read_text()) + if ( + recipe["mode"] != value["state"]["mode"] + or recipe["reordering"] != value["state"]["map_kind"] + ): + raise BundleError("State recipe and compiled state mode disagree") + for field, key in (("scales", "state_scales"), ("map", "state_map")): + if recipe[field]["sha256"] != hashlib.sha256(artifacts[key].read_bytes()).hexdigest(): + raise BundleError("State artifact differs from its recipe: " + field) + return binding + + +def bind_state_and_startup(neural, plan, artifacts): + function = neural.lib.mk2_set_native_predictor_active_mask + function.argtypes = [C.c_void_p, C.c_void_p] + function.restype = C.c_int + neural._check( + function(neural.stage0_handle, neural.inputs.audio_valid.data_ptr()), "startup mask" + ) + neural._startup_skip_mask_owner = neural.inputs.audio_valid + state_bytes = neural.lib.mk2_r9_state_bytes + state_bytes.argtypes = [] + state_bytes.restype = C.c_int + if state_bytes() != 2: + raise BundleError("Compiled recurrent state is not BF16 private storage") + state_binding.install( + neural, + state_mode=plan.contract["state"]["mode"], + scales_path=artifacts["state_scales"], + map_path=artifacts["state_map"], + recipe_identity=artifacts["state_recipe"], + ) + if plan.contract["kernels"]["prefill"].get("validation_kind") == "retained-device-dedup": + if neural.batch != 512: + raise BundleError("Legacy deduplicated prefill is qualified only at B512") + return + original = neural.prefill_lib.mk2_native_prefill_run + validate = neural.prefill_lib.mk2_native_prefill_validate + validate.argtypes = [C.c_void_p, C.c_void_p] + validate.restype = C.c_int + + def checked_prefill(handle, stream): + status = original(handle, stream) + return status if status else validate(handle, stream) + + neural.prefill_lib.mk2_native_prefill_run = checked_prefill + neural._r14_prefill_validation = dict( + entry="mk2_native_prefill_validate", + scope="cache preparation and every miss before consumers or snapshots", + workspace_sequences=1, + ) + + +def bind_receipt_contract(engine, plan, binding): + original = engine.stream + admission = dict( + initial_singleton_readiness=plan.contract["publication"]["initial_singleton_readiness"], + post_readiness_width=plan.physical_width, + incomplete_deadline_ms=plan.contract["publication"]["incomplete_deadline_ms"], + ) + + def stream(*args, **kwargs): + iterator = original(*args, **kwargs) + try: + for event in iterator: + if event["type"] == "receipt": + receipt = event["receipt"] + receipt["r9_configuration"] = dict( + id=binding["id"], + components=binding["r9_components"], + state_format=plan.contract["state"]["format"], + transport_format="pcm_f32_le", + ) + receipt["immutable_prefix_cache"] = copy.deepcopy(engine._r9_prefix_cache[3]) + receipt["r21_admission"] = dict(admission) + receipt["native_package"] = dict( + id=plan.identity, + manifest_sha256=plan.manifest_sha256, + qualification_only=True, + ) + yield event + finally: + iterator.close() + + engine.stream = stream + + +class ConfiguredEngine(NativeCohortEngine): + def __init__( + self, + plan, + artifacts, + artifact_root, + model_root, + source_pins, + binding, + *, + research_schedule=(3, 5), + research_steps=None, + ): + self.native_plan = plan + self.native_artifacts = artifacts + self.native_compiled_binding = binding + support = PackageSupport(plan, artifacts, source_pins) + super().__init__( + plan.physical_width, + model_root=model_root, + steps=research_steps + if research_steps is not None + else plan.contract["generation"]["cap"], + retain_codes=False, + support=support, + artifact_root=artifact_root, + qualification_mode=True, + execution=dict(plan.contract["execution"]), + active_lanes={"mode": "disabled"}, + research_schedule=research_schedule, + ) + try: + expected = plan.contract["execution"] + if ( + self.residency["stage0"]["batch"] != plan.physical_width + or self.residency["stage0"]["grid"] != expected["full_stage0_grid"] + or self.residency["codec_blocks"] != expected["codec_grid"] + ): + raise BundleError( + "Actual launch does not match the physical width and resource program" + ) + geometry = plan.contract["kernels"]["stage0"] + actual = self.residency["stage0"] + for declared, observed in ( + ("row_tile", "batch_pad"), + ("k_tile", "kchunk"), + ("pipeline_stages", "stages"), + ): + if geometry[declared] != actual[observed]: + raise BundleError( + "Declared producer geometry differs from compiled ABI: " + declared + ) + token_staging.install(self, mode=plan.contract["host"]["tokens"]) + except BaseException: + self.close() + raise + + def _allocate_queues(self): + bind_state_and_startup(self.core.neural, self.native_plan, self.native_artifacts) + self.startup_predictor_skip = True + self.helper_graphs = helper_graphs.install(self.core.neural) + super()._allocate_queues() + self.r8_options = dict(self.native_compiled_binding["runtime_options"]) + chunk_submission.install(self) + prefix_cache.install(self) + self._r9_prefix_cache[3]["key"]["state"] = self.native_plan.contract["state"]["format"] + bind_receipt_contract(self, self.native_plan, self.native_compiled_binding) + pcm_transfer.install(self) + + +def create_for_qualification( + plan: NativeBundlePlan, + artifact_root: Path, + model_root: Path, + *, + source_pins: Mapping[str, str], + model_identity: str, + qualification_only: bool, +): + if qualification_only is not True or plan.contract["enabled"] is not False: + raise BundleError( + "The native port is available only through an explicit qualification call" + ) + if model_identity != plan.contract["model"]["checkpoint_sha256"]: + raise BundleError("Caller model identity differs from the frozen proposal") + if plan.physical_width not in (128, 256, 512, 1024, 2048, 4096): + raise BundleError("Unported wide construction contract") + artifacts = resolve_artifacts(plan, artifact_root) + binding = verify_compiled_contract(plan, artifacts) + return ConfiguredEngine( + plan, artifacts, Path(artifact_root), Path(model_root), source_pins, binding + ) diff --git a/python/nemotron_native_backend/execution/controller_rules.py b/python/nemotron_native_backend/execution/controller_rules.py new file mode 100644 index 0000000..d4ffc33 --- /dev/null +++ b/python/nemotron_native_backend/execution/controller_rules.py @@ -0,0 +1,121 @@ +import operator + + +def effective_seed(seed, epoch): + if type(seed) is not int or not 0 <= seed <= 9223372036854775807: + raise ValueError("seed must be an unsigned 63-bit integer") + if type(epoch) is not int or epoch < 1: + raise ValueError("epoch must be a positive integer") + return seed + epoch - 1 & 9223372036854775807 + + +def generation_deficit(counts, consumed, done, target, executed, limit): + if not len(counts) == len(consumed) == len(done): + raise ValueError("queue metadata widths differ") + if target not in (3, 5, 8) or not 0 <= executed <= limit: + raise ValueError("invalid generation/chunk bounds") + if True in done: + pending = [ + n - used + for (n, used, terminal) in zip(counts, consumed, done, strict=False) + if not terminal + ] + need = max(0, target - min(pending)) if pending else 0 + else: + need = max(0, target - min(map(operator.sub, counts, consumed))) + return min(need, limit - executed) + + +def continuation_lanes(counts, consumed, held): + return [ + lane + for (lane, (count, used, pending)) in enumerate(zip(counts, consumed, held, strict=False)) + if pending and count > used + ] + + +def publication_probe_steps(needed, held, done): + return ( + min(needed, 1) + if any(pending and (not terminal) for (pending, terminal) in zip(held, done, strict=False)) + else needed + ) + + +def validate_execution(row): + mode = row.get("mode", "serial") + policy = row.get("grid_policy", "static") + if mode not in ("serial", "overlap", "after-first-publication"): + raise ValueError("unsupported overlap mode") + if policy not in ("static", "overlap-only"): + raise ValueError("unsupported grid policy") + out = dict( + mode=mode, + grid_policy=policy, + stage0_grid=row.get("stage0_grid", 0), + codec_grid=row.get("codec_grid", 0), + ) + for key in ("stage0_grid", "codec_grid"): + if type(out[key]) is not int or out[key] < 0: + raise ValueError("grid must be a nonnegative integer") + if mode != "serial" and (not out["stage0_grid"] or not out["codec_grid"]): + raise ValueError("overlap requires explicit bounded grids") + if policy == "overlap-only": + full = row.get("full_stage0_grid") + if type(full) is not int or full < out["stage0_grid"] or full < 1: + raise ValueError("overlap-only requires a full grid at least as large as overlap grid") + out["full_stage0_grid"] = full + return out + + +def should_prefetch( + mode, counts, consumed, done, held, sizes, first_published, next_frames, executed, limit +): + if mode == "serial" or (mode == "after-first-publication" and (not all(first_published))): + return False + future_held = [pending or size > 0 for (pending, size) in zip(held, sizes, strict=False)] + for lane in continuation_lanes(counts, consumed, future_held): + future_held[lane] = False + needed = generation_deficit(counts, consumed, done, next_frames, executed, limit) + return publication_probe_steps(needed, future_held, done) == 1 + + +def select_retained_rows(done, target): + if not done or type(target) is not int or (not 0 < target < len(done)): + raise ValueError("invalid width contraction") + if any(type(v) is not bool for v in done): + raise ValueError("done must contain host booleans") + live = [i for (i, v) in enumerate(done) if not v] + if not 0 < len(live) <= target: + return None + return live + [i for (i, v) in enumerate(done) if v][: target - len(live)] + + +def validate_active_lanes(row, batch, execution): + if row is None or row.get("mode") == "disabled": + return None + if row.get("mode") != "fixed-subset" or row.get("target_batch") != 16 or batch not in (32, 64): + raise ValueError("contraction supports B32/B64 to B16 only") + if row.get("after_chunk0") is not True: + raise ValueError("contraction must preserve the chunk0 path") + grid = row.get("narrow_grid") + if type(grid) is not int or grid <= 0: + raise ValueError("narrow grid must be a positive integer") + policy = row.get("grid_policy", "static") + if policy not in ("static", "follow-stage0"): + raise ValueError("unsupported narrow grid policy") + if policy == "follow-stage0" and ( + execution.get("grid_policy") != "overlap-only" or grid != execution.get("stage0_grid") + ): + raise ValueError( + "dynamic narrow grid must follow an explicit bounded Stage0 execution grid" + ) + for name in ("stage0", "transfer"): + artifact = row.get(name, {}) + if ( + set(artifact) != {"path", "sha256"} + or not isinstance(artifact["path"], str) + or len(artifact["sha256"]) != 64 + ): + raise ValueError("auxiliary library requires exact path and SHA256") + return dict(row) diff --git a/python/nemotron_native_backend/execution/engine.py b/python/nemotron_native_backend/execution/engine.py new file mode 100644 index 0000000..4c20b85 --- /dev/null +++ b/python/nemotron_native_backend/execution/engine.py @@ -0,0 +1,1220 @@ +from __future__ import annotations + +import ctypes as C +import dataclasses +import hashlib +import json +import operator +import os +import threading +import time +from pathlib import Path + +SAMPLE_RATE = 22050 +SAMPLES_PER_FRAME = 1764 +WIDTHS = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096) +MAX_STEPS = 511 +_CONSTRUCTION_LOCK = threading.RLock() + + +def effective_seed(seed, epoch): + if type(seed) is not int or not 0 <= seed <= 9223372036854775807: + raise ValueError("seed must be an unsigned 63-bit integer") + if type(epoch) is not int or epoch < 1: + raise ValueError("epoch must be a positive integer") + return seed + epoch - 1 & 9223372036854775807 + + +def generation_deficit(counts, consumed, done, target, executed, limit): + if not len(counts) == len(consumed) == len(done): + raise ValueError("queue metadata widths differ") + if type(target) is not int or not 1 <= target <= 8 or not 0 <= executed <= limit: + raise ValueError("invalid generation/chunk bounds") + if True in done: + pending = [ + n - used + for (n, used, terminal) in zip(counts, consumed, done, strict=False) + if not terminal + ] + need = max(0, target - min(pending)) if pending else 0 + else: + need = max(0, target - min(map(operator.sub, counts, consumed))) + return min(need, limit - executed) + + +def continuation_lanes(counts, consumed, held): + return [ + lane + for (lane, (count, used, pending)) in enumerate(zip(counts, consumed, held, strict=False)) + if pending and count > used + ] + + +def publication_probe_steps(needed, held, done): + return ( + min(needed, 1) + if any(pending and (not terminal) for (pending, terminal) in zip(held, done, strict=False)) + else needed + ) + + +def validate_execution(row): + mode = row.get("mode", "serial") + policy = row.get("grid_policy", "static") + if mode not in ("serial", "overlap", "after-first-publication"): + raise ValueError("unsupported overlap mode") + if policy not in ("static", "overlap-only"): + raise ValueError("unsupported grid policy") + out = dict( + mode=mode, + grid_policy=policy, + stage0_grid=row.get("stage0_grid", 0), + codec_grid=row.get("codec_grid", 0), + ) + for key in ("stage0_grid", "codec_grid"): + if type(out[key]) is not int or out[key] < 0: + raise ValueError("grid must be a nonnegative integer") + if mode != "serial" and (not out["stage0_grid"] or not out["codec_grid"]): + raise ValueError("overlap requires explicit bounded grids") + if policy == "overlap-only": + full = row.get("full_stage0_grid") + if type(full) is not int or full < out["stage0_grid"] or full < 1: + raise ValueError("overlap-only requires a full grid at least as large as overlap grid") + out["full_stage0_grid"] = full + return out + + +def should_prefetch( + mode, counts, consumed, done, held, sizes, first_published, next_frames, executed, limit +): + if mode == "serial" or (mode == "after-first-publication" and (not all(first_published))): + return False + future_held = [pending or size > 0 for (pending, size) in zip(held, sizes, strict=False)] + for lane in continuation_lanes(counts, consumed, future_held): + future_held[lane] = False + needed = generation_deficit(counts, consumed, done, next_frames, executed, limit) + return publication_probe_steps(needed, future_held, done) == 1 + + +class Stage0BuildInfo(C.Structure): + _fields_ = [ + (name, C.c_int) + for name in ( + "version", + "batch", + "batch_pad", + "stages", + "kchunk", + "direct_input", + "b1_local_barriers", + "skip_hidden_snapshot", + "threads", + "dynamic_smem", + "static_smem", + "registers", + "device", + "sms", + "blocks_per_sm", + "grid", + ) + ] + + +def inspect_residency(core, execution): + n = core.neural + n.lib.mk2_build_info.argtypes = [C.POINTER(Stage0BuildInfo)] + n.lib.mk2_build_info.restype = C.c_int + info = Stage0BuildInfo() + n._check(n.lib.mk2_build_info(C.byref(info)), "overlap Stage0 build information") + actual_codec = core.codec.lib.native_stage1_resident_blocks(core.codec.handle) + props = core.t.cuda.get_device_properties(0) + if info.sms != props.multi_processor_count: + raise RuntimeError("stage0/runtime device SM count mismatch") + expected_stage0 = ( + execution.get("full_stage0_grid") + if execution["grid_policy"] == "overlap-only" + else execution["stage0_grid"] + ) + if expected_stage0 and info.grid != expected_stage0: + raise RuntimeError("requested Stage0 grid was clamped or ignored") + if execution["codec_grid"] and actual_codec != execution["codec_grid"]: + raise RuntimeError("requested codec grid was ignored") + if execution["mode"] != "serial": + if info.blocks_per_sm != 1: + raise RuntimeError("experiment currently qualifies one Stage0 CTA per SM only") + if execution["stage0_grid"] + actual_codec > info.sms: + raise RuntimeError("persistent grids exceed distinct-SM capacity") + if execution["grid_policy"] == "overlap-only": + setter = n.lib.mk2_set_native_stage0_grid + setter.argtypes = [C.c_void_p, C.c_int] + setter.restype = C.c_int + return dict( + stage0={name: getattr(info, name) for (name, _) in info._fields_}, + codec_blocks=actual_codec, + device_name=props.name, + distinct_sm_capacity_checked=execution["stage0_grid"] + actual_codec <= info.sms, + progress_guarantee="finite independently resident grids; no inter-kernel device " + "wait; actual overlap " + "requires GPU evidence", + ) + + +def select_retained_rows(done, target): + if not done or type(target) is not int or (not 0 < target < len(done)): + raise ValueError("invalid width contraction") + if any(type(v) is not bool for v in done): + raise ValueError("done must contain host booleans") + live = [i for (i, v) in enumerate(done) if not v] + if not 0 < len(live) <= target: + return None + return live + [i for (i, v) in enumerate(done) if v][: target - len(live)] + + +def validate_active_lanes(row, batch, execution): + if row is None or row.get("mode") == "disabled": + return None + if row.get("mode") != "fixed-subset" or row.get("target_batch") != 16 or batch not in (32, 64): + raise ValueError("contraction supports B32/B64 to B16 only") + if row.get("after_chunk0") is not True: + raise ValueError("contraction must preserve the chunk0 path") + grid = row.get("narrow_grid") + if type(grid) is not int or grid <= 0: + raise ValueError("narrow grid must be a positive integer") + policy = row.get("grid_policy", "static") + if policy not in ("static", "follow-stage0"): + raise ValueError("unsupported narrow grid policy") + if policy == "follow-stage0" and ( + execution.get("grid_policy") != "overlap-only" or grid != execution.get("stage0_grid") + ): + raise ValueError( + "dynamic narrow grid must follow an explicit bounded Stage0 execution grid" + ) + for name in ("stage0", "transfer"): + artifact = row.get(name, {}) + if ( + set(artifact) != {"path", "sha256"} + or not isinstance(artifact["path"], str) + or len(artifact["sha256"]) != 64 + ): + raise ValueError("auxiliary library requires exact path and SHA256") + return dict(row) + + +class ActiveLaneContraction: + def __init__(self, neural, config, root, codec_grid, info_type, execution): + from .split_native_runtime import bind_stage0_library + + self.n = neural + self.t = neural.torch + self.config = config + self.handle = None + self._original_step = neural.step + self._hook = None + self.enabled = False + self.selected_rows = None + self.transition = None + self.widths = [] + self.grids = [] + self.full_width = neural.batch + self.target = config["target_batch"] + self.grid = config["narrow_grid"] + self.grid_policy = config.get("grid_policy", "static") + self.full_grid = ( + execution.get("full_stage0_grid", self.grid) + if self.grid_policy == "follow-stage0" + else self.grid + ) + self.current_grid = self.full_grid + self.artifacts = {} + for key in ("stage0", "transfer"): + row = config[key] + path = Path(row["path"]) + if not path.is_absolute(): + path = Path(root) / path + path = path.resolve(strict=True) + got = hashlib.sha256(path.read_bytes()).hexdigest() + if got != row["sha256"]: + raise ValueError(f"active-lane {key} library hash mismatch") + self.artifacts[key] = {"path": str(path), "sha256": got} + self.lib = bind_stage0_library(self.artifacts["stage0"]["path"]) + self.transfer = C.CDLL(self.artifacts["transfer"]["path"]) + binding_type = type(neural.bindings) + self.transfer.native_lanes_bindings_size.argtypes = [] + self.transfer.native_lanes_bindings_size.restype = C.c_size_t + if self.transfer.native_lanes_bindings_size() != C.sizeof(binding_type): + raise RuntimeError("lane-transfer ABI layout mismatch") + for name in ("native_lanes_gather", "native_lanes_scatter"): + fn = getattr(self.transfer, name) + fn.argtypes = [ + C.POINTER(binding_type), + C.POINTER(binding_type), + C.c_void_p, + C.c_int, + C.c_int, + C.c_void_p, + ] + fn.restype = C.c_int + self.lib.mk2_build_info.argtypes = [C.POINTER(info_type)] + self.lib.mk2_build_info.restype = C.c_int + old_grid = os.environ.get("NEMOTRON_STAGE0_MEGAKERNEL_GRID") + os.environ["NEMOTRON_STAGE0_MEGAKERNEL_GRID"] = str(self.full_grid) + try: + self.handle = self.lib.mk2_create_stage0( + C.byref(neural.pw), C.byref(neural.bw), self.target + ) + neural._require_handle(self.handle, "contracted Stage0 create", self.lib) + info = info_type() + neural._check( + self.lib.mk2_build_info(C.byref(info)), "contracted Stage0 resources", self.lib + ) + if ( + self.lib.mk2_batch(self.handle) != self.target + or info.batch != self.target + or info.grid != self.full_grid + ): + raise RuntimeError("contracted Stage0 width or requested grid mismatch") + if info.blocks_per_sm != 1 or ( + execution["mode"] != "serial" and self.grid + codec_grid > info.sms + ): + raise RuntimeError( + "contracted persistent grids exceed verified distinct-SM capacity" + ) + self.residency = {name: getattr(info, name) for (name, _) in info._fields_} + if self.grid_policy == "follow-stage0": + self.lib.mk2_set_native_stage0_grid.argtypes = [C.c_void_p, C.c_int] + self.lib.mk2_set_native_stage0_grid.restype = C.c_int + except BaseException: + if self.handle: + self.lib.mk2_destroy(self.handle) + self.handle = None + raise + finally: + if old_grid is None: + os.environ.pop("NEMOTRON_STAGE0_MEGAKERNEL_GRID", None) + else: + os.environ["NEMOTRON_STAGE0_MEGAKERNEL_GRID"] = old_grid + try: + t = self.t + width = self.target + blocks = neural.bindings.kv_max_blocks_per_seq + + def empty(shape, dtype): + return t.empty(shape, dtype=dtype, device="cuda:0") + + self.rows = empty((width,), t.int32) + self.embedding = empty((width, 1536), t.bfloat16) + self.noise = empty((width, 16, 80), t.float32) + self.codes = empty((width, 16), t.int64) + self.phones = empty((width,), t.int32) + self.seq = empty((width,), t.int32) + self.slots = empty((width,), t.int32) + self.tables = [empty((width, 2, blocks), t.int32) for _ in range(3)] + self.bindings = binding_type.from_buffer_copy(neural.bindings) + for field, tensor in ( + ("embedding", self.embedding), + ("gumbel", self.noise), + ("codes_out", self.codes), + ("phoneme_out", self.phones), + ("seq_len", self.seq), + ("slot_map", self.slots), + ): + setattr(self.bindings, field, tensor.data_ptr()) + for i, table in enumerate(self.tables): + self.bindings.kv_block_table[i] = table.data_ptr() + neural.stream.wait_stream(t.cuda.current_stream(0)) + self._hook = self.step + neural.step = self._hook + except BaseException: + self.close() + raise + + def reset(self): + self.enabled = False + self.selected_rows = None + self.transition = None + self.widths = [] + self.grids = [] + + def set_requested_grid(self, grid): + if self.grid_policy != "follow-stage0": + return + if grid not in (self.grid, self.full_grid): + raise ValueError("producer requested an unqualified narrow launch grid") + if grid != self.current_grid: + self.n._check( + self.lib.mk2_set_native_stage0_grid(self.handle, grid), + "set contracted launch grid", + self.lib, + ) + self.current_grid = grid + + def activate(self, rows, generation, *, reason): + if self.enabled or len(rows) != self.target or len(set(rows)) != self.target: + raise ValueError("activation requires a fresh unique fixed subset") + if any(type(i) is not int or not 0 <= i < self.full_width for i in rows): + raise ValueError("invalid original logical row") + with self.t.cuda.stream(self.n.stream): + self.rows.copy_(self.t.tensor(rows, dtype=self.t.int32)) + self.selected_rows = list(rows) + self.enabled = True + self.transition = dict( + generation=generation, original_logical_rows=list(rows), reason=reason + ) + + def observe(self, done, generation, chunk0_seen): + if self.enabled or not chunk0_seen: + return + rows = select_retained_rows(done, self.target) + if rows is not None: + self.activate( + rows, + generation, + reason="observed EOS after chunk0; retain live rows plus distinct dead padding", + ) + + def step(self): + n = self.n + if not self.enabled: + result = self._original_step() + self.widths.append(self.full_width) + self.grids.append(None) + return result + if ( + n.closed + or not n.started + or n.inputs.generation >= min(n.cache["max_generations"], len(n.step_events)) + ): + raise RuntimeError("contracted step exceeds prepared diagnostic capacity") + with self.t.cuda.stream(n.stream): + n.inputs.prepare_decode() + n.noise.generation(n.batch) + args = ( + C.byref(n.bindings), + C.byref(self.bindings), + self.rows.data_ptr(), + self.full_width, + self.target, + n.stream.cuda_stream, + ) + rc = self.transfer.native_lanes_gather(*args) + if rc: + raise RuntimeError(f"lane gather CUDA error {rc}") + (a, b) = n.step_events[n.inputs.generation] + a.record(n.stream) + n._check( + self.lib.mk2_run_native_stage0( + self.handle, C.byref(self.bindings), n.stream.cuda_stream + ), + "contracted full Stage0", + self.lib, + ) + b.record(n.stream) + rc = self.transfer.native_lanes_scatter(*args) + if rc: + raise RuntimeError(f"lane scatter CUDA error {rc}") + n.inputs.commit_outputs(n.codes, n.phonemes) + n.history.append((a, b)) + self.widths.append(self.target) + self.grids.append(self.current_grid) + return (n.codes, n.phonemes, n.inputs.stop) + + def receipt(self): + return dict( + mode="fixed-subset", + batch=self.full_width, + target_batch=self.target, + after_chunk0=True, + narrow_grid=self.grid, + grid_policy=self.grid_policy, + full_narrow_grid=self.full_grid, + transition=self.transition, + stage0_widths=self.widths, + stage0_grids=self.grids, + contracted_generations=self.widths.count(self.target), + artifacts=self.artifacts, + residency=self.residency, + rng="unchanged full64 draws; gather original logical rows", + inactive_outputs="dropped dead rows codes=0 and phonemes=0; live and padding " + "rows scattered " + "before unchanged original-width helper graph", + state="original conv/SSM/KV pool addresses; gathered original slot-map and KV-table " + "rows; canonical live seq lengths scattered", + scope="two auxiliary transfer kernels per contracted generation; one full native " + "Stage0 neural kernel; unchanged request membership and original-width codec", + ) + + def close(self): + if self.handle: + self.n.stream.synchronize() + self.lib.mk2_destroy(self.handle) + self.handle = None + if self._hook is not None: + if self.n.step is not self._hook: + raise RuntimeError("contracted step hook ownership changed") + self.n.step = self._original_step + self._hook = None + + +class NativeCohortEngine: + def __init__( + self, + batch, + model_root, + steps=MAX_STEPS, + retain_codes=False, + *, + support, + artifact_root, + execution=None, + stage0_library=None, + codec_library=None, + prefill_library=None, + qualification_mode=False, + serial_match_mode="overlap", + active_lanes=None, + research_schedule=(3, 5), + ): + if batch not in WIDTHS or not 6 <= steps <= MAX_STEPS: + raise ValueError("requires B1/B16/B32/B64 and 6..511 generation steps") + if type(research_schedule) is not tuple or research_schedule not in ( + (3, 5), + (2, 2, 2, 4), + (1, 2, 2, 3), + (2, 1, 5), + (2, 2, 4), + (2, 3, 4, 5), + ): + raise ValueError("unsupported research chunk schedule") + if research_schedule != (3, 5) and qualification_mode is not True: + raise ValueError("changed packet contract is qualification-only") + self.research_schedule = research_schedule + self.batch = batch + self.steps = steps + self.retain_codes = bool(retain_codes) + self._lock = threading.Lock() + self._closed = False + self.last_receipt = None + self.core = None + if (execution is not None or stage0_library or codec_library or prefill_library) and ( + not qualification_mode + ): + raise ValueError("explicit unqualified overrides require qualification_mode=True") + self.execution = validate_execution(execution if execution is not None else {}) + self.active_lanes_config = validate_active_lanes(active_lanes, batch, self.execution) + if active_lanes is not None and (not qualification_mode): + raise ValueError("explicit active-lane override requires qualification_mode=True") + self.contraction = None + self.qualification_mode = bool(qualification_mode) + if serial_match_mode not in ("overlap", "after-first-publication"): + raise ValueError("invalid serial grid reference mode") + if serial_match_mode != "overlap" and (not qualification_mode): + raise ValueError("serial grid reference override is qualification-only") + self._serial_grid_reference_mode = serial_match_mode + self._codec_pending = None + self._deferred_serial_grid = False + self._overlap_records = [] + self.s = support + self.s.verify_sources() + (self.paths, self.selection) = self.s.select(batch) + self.startup_predictor_skip = False + self.helper_graphs = None + for name, override in ( + ("stage0", stage0_library), + ("stage1", codec_library), + ("prefill", prefill_library), + ): + if override: + self.paths[name] = Path(override).resolve(strict=True) + self._actual_library_records = { + name: {"path": str(path), "sha256": hashlib.sha256(path.read_bytes()).hexdigest()} + for (name, path) in self.paths.items() + } + self._support_source_sha256 = self.s.source_sha256 + self._engine_source_sha256 = hashlib.sha256(Path(__file__).read_bytes()).hexdigest() + self.assets = self.s.ModelAssets(Path(model_root)) + try: + with _CONSTRUCTION_LOCK: + old_grid = os.environ.get("NEMOTRON_STAGE0_MEGAKERNEL_GRID") + initial_grid = ( + self.execution.get("full_stage0_grid") + if self.execution["grid_policy"] == "overlap-only" + else self.execution["stage0_grid"] + ) + os.environ["NEMOTRON_STAGE0_MEGAKERNEL_GRID"] = str(initial_grid) + try: + self.core = self.s.Cohort( + self.assets, + batch, + self.paths, + "Warming the first request path.", + steps, + codec_grid=self.execution["codec_grid"], + ) + self.residency = inspect_residency(self.core, self.execution) + finally: + if old_grid is None: + os.environ.pop("NEMOTRON_STAGE0_MEGAKERNEL_GRID", None) + else: + os.environ["NEMOTRON_STAGE0_MEGAKERNEL_GRID"] = old_grid + self.t = self.core.t + self._allocate_queues() + if self.active_lanes_config: + self.contraction = ActiveLaneContraction( + self.core.neural, + self.active_lanes_config, + artifact_root, + self.residency["codec_blocks"], + Stage0BuildInfo, + self.execution, + ) + except BaseException: + if self.contraction is not None: + self.contraction.close() + if self.core is not None: + self.core.close() + raise + + def _allocate_queues(self): + (t, c, B) = (self.t, self.core, self.batch) + self.queue = t.zeros((B, self.steps, 16), dtype=t.int64, device="cuda:0") + self.counts = t.zeros(B, dtype=t.int64, device="cuda:0") + self.consumed = t.zeros_like(self.counts) + self.consumed_host = [0] * B + self.element_valid = t.empty((B, 16), dtype=t.bool, device="cuda:0") + self.element_tmp = t.empty_like(self.element_valid) + self.valid = t.empty(B, dtype=t.bool, device="cuda:0") + self.nonzero = t.empty_like(self.valid) + self.active = t.empty_like(self.valid) + self.valid_i64 = t.empty_like(self.counts) + self.write_index = t.empty_like(self.counts) + self._write_index3 = self.write_index.view(B, 1, 1).expand(B, 1, 16) + self._step_codes3 = c.neural.codes.view(B, 1, 16) + self._done_event = t.cuda.Event() + self._done_event.record() + self._observe_done = t.empty(B, dtype=t.bool, pin_memory=True) + self._observe_counts = t.empty(B, dtype=t.int64, pin_memory=True) + self._observe_event = t.cuda.Event() + self._observe_event.record() + self.codec_stream = t.cuda.Stream(device=0) + self.codec_stream.wait_stream(t.cuda.current_stream(0)) + self._origin = t.cuda.Event(enable_timing=True) + self._origin.record() + self._event_pool = [ + tuple(t.cuda.Event(enable_timing=True) for _ in range(3)) for _ in range(self.steps) + ] + for triple in self._event_pool: + for event in triple: + event.record() + self._chunk_events = [] + self.bindings = {} + for F in sorted(set((*self.research_schedule, 8))): + packed = t.empty((B, F, 16), dtype=t.int64, device="cuda:0") + lengths = t.empty(B, dtype=t.int32, device="cuda:0") + index = t.empty((B, F), dtype=t.int64, device="cuda:0") + offsets = t.arange(F, dtype=t.int64, device="cuda:0").view(1, F) + mask = t.empty((B, F), dtype=t.bool, device="cuda:0") + (a, b) = (t.cuda.Event(enable_timing=True), t.cuda.Event(enable_timing=True)) + a.record() + b.record() + pcm = c.pcm[: B * F * SAMPLES_PER_FRAME].view(B, F * SAMPLES_PER_FRAME) + banks = [ + c.codec.bind_buffers( + packed, + lengths, + c.slots, + c.banks[k], + c.banks[1 - k], + pcm, + c.samples, + c.status, + stream=self.codec_stream, + completion_event=b, + ) + for k in (0, 1) + ] + self.bindings[F] = (packed, lengths, index, offsets, mask, a, b, banks) + c.neural.stream.wait_stream(t.cuda.current_stream(0)) + t.cuda.synchronize(0) + + def _requests(self, texts, token_ids, speaker): + if speaker != "english_sample": + raise ValueError("selected native assets support english_sample only") + if not isinstance(texts, (list, tuple)) or len(texts) != self.batch: + raise ValueError("texts must contain exactly the configured batch width") + if any(not isinstance(text, str) or not text.strip() for text in texts): + raise ValueError("texts must be nonempty normalized strings") + if token_ids is None: + return self.s.requests_from_texts(self.assets, list(texts)) + from .. import r31 + from .native_inputs import Request, validate_requests + + if len(token_ids) != self.batch: + raise ValueError("token_ids width differs from texts") + context = self.core.requests[0].context_ids + requests = [ + Request(text, tuple(ids), context) + for (text, ids) in zip(texts, token_ids, strict=False) + ] + if r31.enabled("staged_admission"): + if ( + len(context) != 3 + or any(type(x) is not int or not 0 <= x < 131075 for x in context) + or any( + not 1 <= len(r.target_ids) <= 512 or r.target_ids[-1] != 131073 + for r in requests + ) + ): + raise ValueError("invalid native request context/target/EOS contract") + return requests + validate_requests(requests) + return requests + + def _reset(self, requests, seed): + (t, c, n) = (self.t, self.core, self.core.neural) + tokens = t.zeros((self.batch, 512), dtype=t.int64) + lengths = t.tensor([len(r.target_ids) for r in requests], dtype=t.int64) + for lane, request in enumerate(requests): + tokens[lane, : len(request.target_ids)] = t.tensor(request.target_ids, dtype=t.int64) + c.requests = requests + n.inputs.requests = tuple(requests) + if self.contraction is not None: + self.contraction.reset() + c.reset(seed) + with t.cuda.stream(n.stream): + n.inputs.text_tokens.copy_(tokens) + n.inputs.text_lengths.copy_(lengths) + n.inputs.max_text_offsets.copy_(lengths - 1) + self.counts.zero_() + self.consumed.zero_() + self.consumed_host = [0] * self.batch + self._chunk_events = [] + self._overlap_records = [] + self._deferred_serial_grid = False + self._origin.record(n.stream) + + def _append_step(self, generation): + (t, c, n) = (self.t, self.core, self.core.neural) + restore_grid = self._deferred_serial_grid + if restore_grid: + self._set_stage0_grid(overlapped=True) + self._deferred_serial_grid = False + (codes, phones, stop) = n.step() + if restore_grid: + self._set_stage0_grid(overlapped=False) + c.raw_cols[generation].copy_(codes) + c.stop_cols[generation].copy_(stop) + c.phone_cols[generation].copy_(phones) + t.ge(codes, 0, out=self.element_valid) + t.lt(codes, 1024, out=self.element_tmp) + t.logical_and(self.element_valid, self.element_tmp, out=self.element_valid) + t.all(self.element_valid, dim=1, out=self.valid) + t.ne(codes, 0, out=self.element_tmp) + t.any(self.element_tmp, dim=1, out=self.nonzero) + t.logical_and(self.valid, self.nonzero, out=self.valid) + t.logical_not(c.done, out=self.active) + t.logical_and(self.valid, self.active, out=self.valid) + t.clamp(self.counts, max=self.steps - 1, out=self.write_index) + self.queue.scatter_(1, self._write_index3, self._step_codes3) + self.valid_i64.copy_(self.valid) + self.counts.add_(self.valid_i64) + t.logical_or(c.done, stop, out=c.done) + + def _observe(self, first_eos, start, end): + (c, n) = (self.core, self.core.neural) + if c.stops.dtype is not self.t.bool or c.done.dtype is not self.t.bool: + raise RuntimeError("EOS observations must be boolean") + self._done_event.record(n.stream) + self._done_event.synchronize() + self._observe_done.copy_(c.done, non_blocking=True) + self._observe_counts.copy_(self.counts, non_blocking=True) + self._observe_event.record() + self._observe_event.synchronize() + done = self._observe_done.tolist() + if True in done: + rows = c.stops[:, start:end].cpu().tolist() + first_eos = self.s.update_first_eos(first_eos, rows, start) + else: + first_eos = list(first_eos) + if done != [g is not None for g in first_eos]: + raise RuntimeError("native EOS carry mismatch") + if self.contraction is not None: + self.contraction.observe(done, end, bool(self._chunk_events)) + return (first_eos, done, self._observe_counts.tolist()) + + def _set_stage0_grid(self, *, overlapped): + if self.execution["grid_policy"] != "overlap-only": + return + grid = self.execution["stage0_grid"] if overlapped else self.execution["full_stage0_grid"] + n = self.core.neural + n._check( + n.lib.mk2_set_native_stage0_grid(n.stage0_handle, grid), "set finite Stage0 launch grid" + ) + if self.contraction is not None: + self.contraction.set_requested_grid(grid) + + def _submit_chunk(self, F, counts, bank): + if self._codec_pending is not None: + raise RuntimeError("codec buffer/scratch lease still outstanding") + (t, n) = (self.t, self.core.neural) + (packed, lengths, index, offsets, mask, _, _, banks) = self.bindings[F] + (ready, a, b) = self._event_pool[len(self._chunk_events)] + sizes = [ + min(F, available - used) + for (available, used) in zip(counts, self.consumed_host, strict=False) + ] + with t.cuda.stream(n.stream): + lengths.copy_(t.tensor(sizes, dtype=t.int32)) + t.add(self.consumed.view(self.batch, 1), offsets, out=index) + index.clamp_(max=self.steps - 1) + t.gather(self.queue, 1, index.unsqueeze(-1).expand(self.batch, F, 16), out=packed) + t.lt(offsets, lengths.view(self.batch, 1), out=mask) + packed.mul_(mask.unsqueeze(-1)) + self.consumed.add_(lengths) + ready.record(n.stream) + self.consumed_host = [ + used + size for (used, size) in zip(self.consumed_host, sizes, strict=False) + ] + self.codec_stream.wait_event(ready) + with t.cuda.stream(self.codec_stream): + a.record(self.codec_stream) + banks[bank].event = b + ticket = banks[bank].enqueue() + pending = dict(ticket=ticket, sizes=sizes, events=(a, b)) + self._codec_pending = pending + self._chunk_events.append((a, b)) + return pending + + def _finish_chunk(self, pending): + if pending is not self._codec_pending: + raise RuntimeError("wrong codec ticket") + result = self.core.codec.finish(pending["ticket"], copy_pcm=True) + self._codec_pending = None + samples = result["samples"].cpu().tolist() + if samples != [size * SAMPLES_PER_FRAME for size in pending["sizes"]]: + raise RuntimeError("native codec sample counts differ from valid-frame queue") + (a, b) = pending["events"] + return (result, pending["sizes"], samples, a.elapsed_time(b)) + + def _overlap_receipt(self): + rows = [] + for row in self._overlap_records: + (a, b) = row["codec_events"] + (c0, c1) = (self._origin.elapsed_time(a), self._origin.elapsed_time(b)) + generation = row["stage0_generation"] + item = dict( + chunk_index=row["chunk_index"], + stage0_generation=generation, + codec_start_ms=c0, + codec_end_ms=c1, + event_interval_overlap_ms=0.0, + ) + if generation is not None: + (x, y) = self.core.neural.step_events[generation] + (s0, s1) = (self._origin.elapsed_time(x), self._origin.elapsed_time(y)) + item.update( + stage0_start_ms=s0, + stage0_end_ms=s1, + event_interval_overlap_ms=max(0.0, min(c1, s1) - max(c0, s0)), + ) + rows.append(item) + return dict( + rows=rows, + prefetched_generations=sum(r["stage0_generation"] is not None for r in rows), + event_interval_overlap_ms=sum(r["event_interval_overlap_ms"] for r in rows), + caveat="CUDA event intervals include admission wait; use an Nsight timeline to prove " + "simultaneous kernel residency", + ) + + def stream(self, texts, *, token_ids=None, speaker="english_sample", seed=0, epoch=1): + seed_value = effective_seed(seed, epoch) + requests = self._requests(texts, token_ids, speaker) + if not self._lock.acquire(blocking=False): + raise RuntimeError("native engine already has a live cohort") + try: + if self._closed: + raise RuntimeError("native engine is closed") + yield from self._stream(requests, speaker, seed, epoch, seed_value) + self._r31_predraw(effective_seed(seed, epoch + 1)) + finally: + try: + if self.core is not None: + self.core.neural.stream.synchronize() + if self._codec_pending is not None: + self.core.codec.finish(self._codec_pending["ticket"], copy_pcm=False) + self._codec_pending = None + self.codec_stream.synchronize() + finally: + self._lock.release() + + def _r31_predraw(self, seed): + from .. import r31 + + if not r31.enabled("rng_predraw") or self.core is None: + return + n = self.core.neural + if not hasattr(n.noise, "predraw"): + return + with self.t.cuda.stream(n.stream): + n.noise.predraw(seed, 8) + + def _stream(self, requests, speaker, base_seed, epoch, seed): + import numpy as np + + from .. import r31 + + (t, c, n) = (self.t, self.core, self.core.neural) + first_eos = [None] * self.batch + finished = [False] * self.batch + held = [False] * self.batch + first_published = [False] * self.batch + total_samples = [0] * self.batch + deliveries = [[] for _ in range(self.batch)] + pcm_hashes = [hashlib.sha256() for _ in range(self.batch)] + fast_first_chunk = r31.enabled("fast_first_chunk") + pending_hashes = [] + + def flush_hashes(): + for lane, payload in pending_hashes: + pcm_hashes[lane].update(payload) + pending_hashes.clear() + + lane_chunk_indices = [0] * self.batch + chunks = [] + continuation_observations = [] + start_wall = time.perf_counter() + self._reset(requests, seed) + n.start() + with t.cuda.stream(n.stream): + for generation in range(5): + (codes, phones, stop) = n.step() + c.raw_cols[generation].copy_(codes) + c.stop_cols[generation].copy_(stop) + c.phone_cols[generation].copy_(phones) + t.any(c.stops[:, :5], dim=1, out=c.done) + executed = 5 + (first_eos, done, counts) = self._observe(first_eos, 0, executed) + bank = 0 + window_index = 0 + prefetched_generation = None + + def finish_ready(): + for lane in range(self.batch): + terminal = done[lane] or executed == self.steps + drained = counts[lane] == self.consumed_host[lane] + if not finished[lane] and terminal and drained: + finished[lane] = True + first_published[lane] = True + held[lane] = False + yield { + "type": "finish", + "lane": lane, + "reason": "eos" if done[lane] else "length_limit", + "samples": total_samples[lane], + "sample_rate": SAMPLE_RATE, + } + + while True: + F = ( + self.research_schedule[window_index] + if window_index < len(self.research_schedule) + else 8 + ) + while True: + release = continuation_lanes(counts, self.consumed_host, held) + if release: + for lane in release: + held[lane] = False + first_published[lane] = True + continuation_observations.append( + { + "executed_generations": executed, + "lanes": release, + "elapsed_ms": (time.perf_counter() - start_wall) * 1000, + } + ) + yield {"type": "continuation", "lanes": release} + yield from finish_ready() + needed = generation_deficit( + counts, self.consumed_host, done, F, executed, self.steps + ) + if not needed: + break + needed = publication_probe_steps(needed, held, done) + before = executed + if prefetched_generation is not None: + if prefetched_generation != before or needed != 1: + raise RuntimeError("prefetch diverged from serial observation policy") + executed += 1 + prefetched_generation = None + else: + with t.cuda.stream(n.stream): + for generation in range(before, before + needed): + self._append_step(generation) + executed += 1 + flush_hashes() + (first_eos, done, counts) = self._observe(first_eos, before, executed) + pending = [ + count - used for (count, used) in zip(counts, self.consumed_host, strict=False) + ] + if any(pending): + submission_generations = executed + pending_codec = self._submit_chunk(F, counts, bank) + sizes = pending_codec["sizes"] + next_frames = ( + self.research_schedule[window_index + 1] + if window_index + 1 < len(self.research_schedule) + else 8 + ) + prefetch = should_prefetch( + self.execution["mode"], + counts, + self.consumed_host, + done, + held, + sizes, + first_published, + next_frames, + executed, + self.steps, + ) + if ( + self.execution["mode"] == "serial" + and self.execution["grid_policy"] == "overlap-only" + ): + self._deferred_serial_grid = should_prefetch( + self._serial_grid_reference_mode, + counts, + self.consumed_host, + done, + held, + sizes, + first_published, + next_frames, + executed, + self.steps, + ) + if prefetch: + self._set_stage0_grid(overlapped=True) + if prefetched_generation is not None: + raise RuntimeError("only one future generation may be outstanding") + with t.cuda.stream(n.stream): + self._append_step(executed) + prefetched_generation = executed + self._set_stage0_grid(overlapped=False) + early_proof = ( + bool(getattr(self, "r8_options", {}).get("early_proof")) + and window_index == 0 + and (not prefetch) + and should_prefetch( + "overlap", + counts, + self.consumed_host, + done, + held, + sizes, + first_published, + next_frames, + executed, + self.steps, + ) + ) + if early_proof: + if prefetched_generation is not None: + raise RuntimeError("proof generation already submitted") + n.stream.wait_event(pending_codec["events"][1]) + with t.cuda.stream(n.stream): + self._append_step(executed) + prefetched_generation = executed + (result, sizes, samples, codec_ms) = self._finish_chunk(pending_codec) + self._overlap_records.append( + { + "chunk_index": window_index, + "stage0_generation": submission_generations if prefetch else None, + "codec_events": pending_codec["events"], + } + ) + data = result["pcm"].numpy() + finite = result.get("finite") + if not (np.isfinite(data).all() if finite is None else finite): + raise RuntimeError("native codec produced nonfinite PCM") + delivery_ms = (time.perf_counter() - start_wall) * 1000 + chunks.append( + { + "index": window_index, + "frames": F, + "executed_generations": submission_generations, + "prefetched_generation": submission_generations if prefetch else None, + "true_lengths": sizes, + "samples": samples, + "delivery_ms": delivery_ms, + "stage1_event_ms": codec_ms, + "status": result["status"], + } + ) + last_chunk_lane = max((i for (i, value) in enumerate(samples) if value), default=-1) + for lane, sample_count in enumerate(samples): + if sample_count: + pcm = ( + data[lane, :sample_count] + .astype( + "= MAX_GENERATIONS: + raise RuntimeError( + "prepare requires uncommitted-free state and remaining generation capacity" + ) + if inputs.generation == 0: + self.validate_bindings() + self.prepare_graph.replay() + inputs.prepared = True + self.prepare_replays += 1 + return inputs.embedding + + def commit_outputs(self, codes, phonemes): + self._check_stream() + inputs = self.inputs + if not inputs.prepared: + raise RuntimeError("commit requires prepared inputs") + if ( + codes.dtype != self.torch.int64 + or tuple(codes.shape) != (inputs.batch, 16) + or codes.device != inputs.device + or (codes.data_ptr() != self._codes_pointer) + or (tuple(codes.stride()) != self._signatures["neural.codes"][2]) + ): + raise ValueError("codes must use the captured stable output pointer and ABI") + if ( + phonemes.dtype != self.torch.int32 + or tuple(phonemes.shape) != (inputs.batch,) + or phonemes.device != inputs.device + or (phonemes.data_ptr() != self._phonemes_pointer) + or (tuple(phonemes.stride()) != self._signatures["neural.phonemes"][2]) + ): + raise ValueError("phonemes must use the captured stable output pointer and ABI") + self.commit_graph.replay() + inputs.prepared = False + inputs.generation += 1 + self.commit_replays += 1 + return inputs.stop + + def receipt(self): + return dict( + schema="native.helper-graphs.v1", + installed=self.installed, + closed=self.closed, + batch=self.inputs.batch, + source_sha256=_digest(__file__), + input_source=str(self.input_source), + input_source_sha256=INPUT_SOURCE_SHA, + prepare_replays=self.prepare_replays, + commit_replays=self.commit_replays, + captured_regions=["ReusableInputs.prepare_decode", "ReusableInputs.commit_outputs"], + neural_kernels_captured=False, + rng_captured=False, + warmup_pairs_rolled_back=3, + graph_lifetime="Retired after stream synchronization and before original neural.close " + "releases buffers", + storage_contract="Owned tensor addresses/shapes/strides remain fixed; full " + "validation at each " + "generation0; output-pointer validation on every commit", + ) + + def _retire_graphs(self): + self.stream.synchronize() + for name in ("commit_graph", "prepare_graph"): + graph = getattr(self, name, None) + if graph is not None: + graph.reset() + setattr(self, name, None) + + def close(self): + if self.closed: + return + if self.installed: + for obj, name, hook in ( + (self.inputs, "prepare_decode", self.prepare_hook), + (self.inputs, "commit_outputs", self.commit_hook), + (self.neural, "close", self.close_hook), + ): + if obj.__dict__.get(name) is not hook: + raise RuntimeError("helper hook ownership changed before uninstall: " + name) + self._retire_graphs() + if self.installed: + for obj, name in ( + (self.inputs, "prepare_decode"), + (self.inputs, "commit_outputs"), + (self.neural, "close"), + (self.neural, "_helper_graphs"), + ): + old = self._old_slots[id(obj), name] + if old is _ABSENT: + delattr(obj, name) + else: + setattr(obj, name, old) + self.installed = False + self.closed = True + self._bound.clear() + + def _close_runner(self): + original = self.original_close + self.close() + return original() + + +def install(neural): + existing = getattr(neural, "_helper_graphs", None) + if existing is not None: + if not isinstance(existing, HelperGraphs) or existing.closed: + raise RuntimeError("unexpected helper graph owner") + existing.validate_bindings() + return existing + return HelperGraphs(neural) + + +def uninstall(neural): + adapter = getattr(neural, "_helper_graphs", None) + if adapter is not None: + adapter.close() diff --git a/python/nemotron_native_backend/execution/inflight_engine.py b/python/nemotron_native_backend/execution/inflight_engine.py new file mode 100644 index 0000000..468b8cc --- /dev/null +++ b/python/nemotron_native_backend/execution/inflight_engine.py @@ -0,0 +1,1367 @@ +from __future__ import annotations + +import ctypes as C +import hashlib +import math +import threading +import time +from collections import deque +from dataclasses import dataclass +from pathlib import Path + +from .engine import SAMPLE_RATE, SAMPLES_PER_FRAME, Stage0BuildInfo, inspect_residency +from .inflight_policy import DEFAULT_POLICY, InflightPolicy +from .native_assets import WIDTHS +from .native_inputs import Request, validate_requests +from .request_noise import RequestNoise +from .reusable_inputs import ReusableInputs + + +@dataclass +class Slot: + epoch: int + request_id: str + seed: int + admitted: float + generation: int = 0 + produced: int = 0 + consumed: int = 0 + chunk_index: int = 0 + samples: int = 0 + reason: str | None = None + finished: bool = False + first_audio: float | None = None + + def target(self, schedule): + return schedule[self.chunk_index] if self.chunk_index < len(schedule) else 8 + + +def ready_groups(slots, progress_slots, schedule, *, first_packet_f2=False): + groups = {} + for index in progress_slots: + slot = slots[index] + if slot is None or slot.finished: + continue + pending = slot.produced - slot.consumed + target = slot.target(schedule) + if pending and (pending >= target or slot.reason is not None): + frames = 3 if target <= 3 else 5 if target <= 5 else 8 + if first_packet_f2 and target <= 2: + frames = 2 + groups.setdefault(frames, []).append(index) + return sorted(groups.items(), key=lambda item: min(slots[i].admitted for i in item[1])) + + +def reset_ranges(slots, batch=32): + if type(batch) is not int or batch not in WIDTHS: + raise ValueError("unsupported physical reset width") + slots = list(slots) + if any(type(slot) is not int or not 0 <= slot < batch for slot in slots): + raise ValueError("reset slots must be valid physical indices") + if len(set(slots)) != len(slots): + raise ValueError("reset slots must be unique") + selected = sum(1 << slot for slot in slots) + ranges = [] + + def visit(start, width): + mask = ((1 << width) - 1) << start + if selected & mask == mask: + ranges.append((start, start + width)) + elif selected & mask: + visit(start, width // 2) + visit(start + width // 2, width // 2) + + visit(0, batch) + return ranges + + +def codec_row_order(selected, batch): + if type(batch) is not int or batch not in WIDTHS: + raise ValueError("unsupported physical codec width") + selected = tuple(selected) + if any(type(lane) is not int or not 0 <= lane < batch for lane in selected): + raise ValueError("codec rows must be valid physical slots") + active = set(selected) + if len(active) != len(selected): + raise ValueError("codec physical slots must be unique") + order = tuple(sorted(active)) + tuple(lane for lane in range(batch) if lane not in active) + inverse = [0] * batch + for row, lane in enumerate(order): + inverse[lane] = row + return order, tuple(inverse) + + +def startup_steps(slots, progress, limit): + if not progress or any( + (slot := slots[index]) is None + or slot.finished + or slot.reason is not None + or slot.generation >= min(5, limit) + for index in progress + ): + return 1 + return max(1, min(min(5, limit) - slots[index].generation for index in progress)) + + +def startup_progress(slots, progress, now, *, defer=False, align_phases=False): + startup = [index for index in progress if slots[index].first_audio is None] + if defer or not startup: + return progress + peers = [slots[index] for index in progress if slots[index].first_audio is not None] + if any( + slot.first_audio is None or slot.first_audio + slot.samples / SAMPLE_RATE - now < 0.040 + for slot in peers + ): + return progress + if align_phases: + startup_phase = [index for index in startup if slots[index].generation < 5] + if startup_phase: + return startup_phase + return startup + + +class InflightEngine: + @classmethod + def from_engine(cls, engine, *, stage0_library, policy=None): + return cls(engine, stage0_library=stage0_library, policy=policy) + + def __init__(self, engine, *, stage0_library, policy=None): + if engine.batch not in WIDTHS or engine.steps > 511: + raise ValueError("in-flight execution requires a supported physical width") + if engine._closed or engine._lock.locked(): + raise RuntimeError("cannot adopt a closed engine or an active cohort") + if tuple(engine.research_schedule) != (2, 3, 4, 5): + raise ValueError("in-flight execution requires the selected 2/3/4/5 startup schedule") + if engine._codec_pending is not None: + raise RuntimeError("cannot adopt an engine with outstanding codec work") + self.engine = engine + self.policy = InflightPolicy.parse(policy) + self.batch = engine.batch + self.t = t = engine.t + self.core = c = engine.core + self.neural = n = c.neural + self.schedule = tuple(engine.research_schedule) + self.steps = engine.steps + self.slots = [None] * self.batch + self.epochs = [0] * self.batch + self.bank = 0 + self.closed = False + self.poisoned = False + self.tick_count = 0 + self.producer_launches = 0 + self.codec_launches = 0 + self.live_refills = 0 + self.observation_pending = False + self.admission_requested = threading.Event() + self.admission_yields = 0 + self.codec_ready_order = deque() + self.startup_priority_deferred = False + self.startup_priority_launches = 0 + self.startup_priority_steps = 0 + self.startup_priority_started = None + if self.policy.playback_reserve_guard: + self.playback_producer_envelopes = {} + self.playback_codec_envelopes = {} + self.playback_producer_pending = None + self.playback_codec_pending = None + self.playback_reserve_last = None + self.trace = deque(maxlen=4096) + self.diagnostic_events = deque(maxlen=8192) + self.diagnostic_dropped = 0 + n.stream.synchronize() + engine.codec_stream.synchronize() + from . import helper_graphs + + if engine.contraction is not None: + engine.contraction.close() + engine.contraction = None + engine.active_lanes_config = None + c.reset(0) + n.start() + n.stream.synchronize() + self._capture_prefix() + snapshot = getattr(engine, "_fp32_prefix_snapshot", None) + if snapshot is not None: + snapshot.detach() + helper_graphs.uninstall(n) + engine.helper_graphs = None + cache = getattr(engine, "_r9_prefix_cache", None) + if cache is not None: + n.prefill_lib.mk2_native_prefill_run = cache[2] + engine._r9_prefix_cache = None + del cache + if self.batch >= 128: + t.cuda.empty_cache() + self._replace_stage0(stage0_library) + with t.cuda.stream(n.stream): + self.progress = t.zeros(self.batch, dtype=t.uint8, device=n.inputs.device) + self.progress_bool = t.zeros(self.batch, dtype=t.bool, device=n.inputs.device) + self.progress_long = t.zeros(self.batch, dtype=t.int64, device=n.inputs.device) + self.old_phoneme_ended = t.zeros_like(n.inputs.phoneme_ended) + self.noise = RequestNoise(n, self.progress) + n.bindings.gumbel = self.noise.buffer.data_ptr() + self.progress_host = t.zeros(self.batch, dtype=t.uint8, pin_memory=True) + self.progress_host_array = self.progress_host.numpy() + self.meta = t.empty((2, self.batch), dtype=t.int64, device=n.inputs.device) + self.meta_host = t.empty((2, self.batch), dtype=t.int64, pin_memory=True) + self.meta_done = t.cuda.Event() + self.ready = t.cuda.Event() + self.tokens_host = t.zeros((self.batch, 512), dtype=t.int64, pin_memory=True) + self.tokens_host_array = self.tokens_host.numpy() + self.phoneme_long = t.zeros_like(n.inputs.last_phoneme) + engine.counts.zero_() + engine.consumed.zero_() + c.done.zero_() + n.codes.zero_() + n.phonemes.zero_() + setter = n.lib.mk2_set_native_progress_mask + setter.argtypes = [C.c_void_p, C.c_void_p] + setter.restype = C.c_int + n._check(setter(n.stage0_handle, self.progress.data_ptr()), "in-flight progress mask") + self._capture_helpers() + if not hasattr(engine, "submit_pcm_transfer"): + from . import pcm_transfer + + pcm_transfer.install(engine) + self._allocate_codec() + self._capture_admissions() + n.stream.synchronize() + + def _view(self, pointer, shape, dtype): + t = self.t + item_size = t.empty((), dtype=dtype).element_size() + count = 1 + for extent in shape: + count *= extent + + class Alias: + __cuda_array_interface__ = { + "shape": (count * item_size,), + "typestr": "|u1", + "data": (int(pointer), False), + "version": 3, + "strides": None, + } + + return t.as_tensor(Alias(), device=self.neural.inputs.device).view(dtype).view(shape) + + def _capture_prefix(self): + t, n = self.t, self.neural + blocks = n.info.blocks_per_sequence + p = n.inputs.prompt_positions + self.seq_len = self._view(n.bindings.seq_len, (self.batch,), t.int32) + self.slot_map = self._view(n.bindings.slot_map, (self.batch,), t.int32) + if not t.equal(self.slot_map, t.arange(self.batch, dtype=t.int32, device=n.inputs.device)): + raise ValueError("in-flight prefix requires identity physical slots") + expected_table = t.arange( + self.batch * 2 * blocks, dtype=t.int32, device=n.inputs.device + ).view(self.batch, 2, blocks) + for i in range(3): + table = self._view(n.bindings.kv_block_table[i], (self.batch, 2, blocks), t.int32) + if not t.equal(table, expected_table): + raise ValueError("in-flight prefix requires stable row-owned KV block tables") + self.conv = [ + self._view( + n.bindings.conv_state[i], (self.batch, n.bindings.conv_slot_stride[i]), t.bfloat16 + ) + for i in range(14) + ] + self.ssm = [] + if getattr(n, "_r8_initial_state", None) is None: + if n.cache["ssm_dtype"] != "FP32": + raise ValueError("uncanonical in-flight state must use FP32 storage") + self.ssm = [ + self._view( + n.bindings.ssm_state[i], (self.batch, n.bindings.ssm_slot_stride[i]), t.float32 + ) + for i in range(14) + ] + self.kv_prefix = [] + for i in range(3): + pool = self._view( + n.bindings.kv_pool_base[i], (self.batch, 2, blocks, 4, 32, 128), t.bfloat16 + ) + for block in range((p + 31) // 32): + self.kv_prefix.append(pool[:, :, block, :, : min(32, p - block * 32), :]) + with t.cuda.stream(n.stream): + self.conv_prefix = [value[0].clone() for value in self.conv] + self.ssm_prefix = [value[0].clone() for value in self.ssm] + self.kv_initial = [value[0].clone() for value in self.kv_prefix] + n.stream.synchronize() + + def _replace_stage0(self, path): + from .split_native_runtime import bind_stage0_library + + n = self.neural + path = Path(path).resolve(strict=True) + lib = bind_stage0_library(path) + version = lib.mk2_native_inflight_abi_version + version.argtypes = [] + version.restype = C.c_int + if version() != 1: + raise ValueError("unsupported in-flight Stage0 ABI") + lib.mk2_build_info.argtypes = [C.POINTER(Stage0BuildInfo)] + lib.mk2_build_info.restype = C.c_int + info = Stage0BuildInfo() + n._check(lib.mk2_build_info(C.byref(info)), "in-flight resource contract", lib) + if info.batch != self.batch or info.blocks_per_sm < 1: + raise ValueError("in-flight kernel changed the supported physical width or residency") + if self.engine.execution["mode"] == "overlap" and info.blocks_per_sm != 1: + raise ValueError("in-flight overlap requires one resident producer block per SM") + n.lib.mk2_destroy(n.stage0_handle) + n.stage0_handle = None + handle = lib.mk2_create_stage0(C.byref(n.pw), C.byref(n.bw), self.batch) + n._require_handle(handle, "in-flight Stage0 allocation", lib) + try: + initial = getattr(n, "_r8_initial_state", None) + if initial is not None: + setter = lib.mk2_r8_bind_initial_state + setter.argtypes = [ + C.c_void_p, + C.POINTER(C.c_void_p), + C.c_void_p, + C.c_void_p, + C.c_int, + ] + setter.restype = C.c_int + n._check( + setter( + handle, + initial["initial"], + initial["ready"], + n.bindings.seq_len, + n.inputs.prompt_positions, + ), + "in-flight canonical state", + lib, + ) + state = getattr(n, "_r9_state", None) + if state is not None: + scales = (C.c_void_p * 14)(*[state["scales"][i].data_ptr() for i in range(14)]) + maps = (C.c_void_p * 14)(*[state["maps"][i].data_ptr() for i in range(14)]) + setter = lib.mk2_r9_state_bind + setter.argtypes = [C.c_void_p, C.POINTER(C.c_void_p), C.POINTER(C.c_void_p)] + setter.restype = C.c_int + n._check(setter(handle, scales, maps), "in-flight state format binding", lib) + setter = lib.mk2_set_native_predictor_active_mask + setter.argtypes = [C.c_void_p, C.c_void_p] + setter.restype = C.c_int + n._check(setter(handle, n.inputs.audio_valid.data_ptr()), "in-flight startup mask", lib) + lib.mk2_set_native_stage0_grid.argtypes = [C.c_void_p, C.c_int] + lib.mk2_set_native_stage0_grid.restype = C.c_int + n._check( + lib.mk2_set_native_stage0_grid( + handle, + self.engine.execution.get( + "full_stage0_grid", self.engine.execution["stage0_grid"] + ), + ), + "in-flight full producer grid", + lib, + ) + except BaseException: + lib.mk2_destroy(handle) + raise + n.lib, n.stage0_handle = lib, handle + self.startup_entry = None + if hasattr(n, "_r8_startup_stage0"): + self.startup_entry = lib.mk2_run_native_startup_stage0 + self.startup_entry.argtypes = lib.mk2_run_native_stage0.argtypes + self.startup_entry.restype = lib.mk2_run_native_stage0.restype + n.library = path + n.library_sha256 = hashlib.sha256(path.read_bytes()).hexdigest() + self.engine.paths = {**self.engine.paths, "stage0": path} + self.engine._actual_library_records = { + **self.engine._actual_library_records, + "stage0": {"path": str(path), "sha256": n.library_sha256}, + } + self.engine.residency = inspect_residency(self.core, self.engine.execution) + + def _prepare_device(self): + t, i = self.t, self.neural.inputs + self.progress_bool.copy_(self.progress) + self.progress_long.copy_(self.progress) + self.old_phoneme_ended.copy_(i.phoneme_ended) + i.prepared = False + i.generation = 0 + ReusableInputs.prepare_decode(i) + t.where(self.progress_bool, i.phoneme_ended, self.old_phoneme_ended, out=i.phoneme_ended) + t.logical_and(i.audio_valid, self.progress_bool, out=i.audio_valid) + i.embedding.mul_(self.progress_bool[:, None]) + + def _commit_device(self): + t, n, e, c = self.t, self.neural, self.engine, self.core + i = n.inputs + t.eq(n.codes, 1025, out=i.code_eos) + t.any(i.code_eos, dim=1, out=i.any_eos) + t.logical_and(i.any_eos, i.audio_valid, out=i.stop) + t.where(self.progress_bool[:, None], n.codes, i.last_codes, out=i.last_codes) + self.phoneme_long.copy_(n.phonemes) + t.where(self.progress_bool, self.phoneme_long, i.last_phoneme, out=i.last_phoneme) + t.ge(n.codes, 0, out=e.element_valid) + t.lt(n.codes, 1024, out=e.element_tmp) + t.logical_and(e.element_valid, e.element_tmp, out=e.element_valid) + t.all(e.element_valid, dim=1, out=e.valid) + t.ne(n.codes, 0, out=e.element_tmp) + t.any(e.element_tmp, dim=1, out=e.nonzero) + t.logical_and(e.valid, e.nonzero, out=e.valid) + t.logical_and(e.valid, i.audio_valid, out=e.valid) + t.logical_not(c.done, out=e.active) + t.logical_and(e.valid, e.active, out=e.valid) + t.clamp(e.counts, max=self.steps - 1, out=e.write_index) + e.queue.scatter_(1, e._write_index3, e._step_codes3) + e.valid_i64.copy_(e.valid) + e.counts.add_(e.valid_i64) + t.logical_or(c.done, i.stop, out=c.done) + i.offsets.add_(self.progress_long) + self.meta[0].copy_(e.counts) + self.meta[1].copy_(c.done) + + def _capture_helpers(self): + t, n = self.t, self.neural + names = ( + "offsets", + "last_codes", + "last_phoneme", + "phoneme_ended", + "audio_valid", + "stop", + "embedding", + "combined_model", + ) + with t.cuda.stream(n.stream): + snapshots = {name: getattr(n.inputs, name).clone() for name in names} + for _ in range(2): + self._prepare_device() + self.noise.generation() + self._commit_device() + n.stream.synchronize() + self.prepare_graph = t.cuda.CUDAGraph() + with t.cuda.graph(self.prepare_graph, stream=n.stream): + self._prepare_device() + self.noise.generation() + self.commit_graph = t.cuda.CUDAGraph() + with t.cuda.graph(self.commit_graph, stream=n.stream): + self._commit_device() + for name, value in snapshots.items(): + getattr(n.inputs, name).copy_(value) + self.engine.counts.zero_() + self.core.done.zero_() + n.inputs.prepared = False + n.stream.synchronize() + + def _allocate_codec(self): + t, c, stream = self.t, self.core, self.engine.codec_stream + self.codec_buffers = {} + self.codec_packing = {} + self.codec_ticket_rows = None + self.codec_packing_launches = 0 + frames_supported = ( + (2, 3, 5, 8) + if getattr(self, "policy", DEFAULT_POLICY).codec_first_packet_f2 + else (3, 5, 8) + ) + for frames in frames_supported: + packed = t.zeros((self.batch, frames, 16), dtype=t.int64, device="cuda:0") + lengths = t.zeros(self.batch, dtype=t.int32, device="cuda:0") + index = t.zeros((self.batch, frames), dtype=t.int64, device="cuda:0") + offsets = t.arange(frames, dtype=t.int64, device="cuda:0").view(1, frames) + mask = t.empty((self.batch, frames), dtype=t.bool, device="cuda:0") + codec_slots = c.slots + if getattr(self, "policy", DEFAULT_POLICY).codec_local_packing: + codec_slots = t.arange(self.batch, dtype=t.int32, device="cuda:0") + self.codec_packing[frames] = { + "slots": codec_slots, + "order": t.arange(self.batch, dtype=t.int64, device="cuda:0"), + "physical_codes": t.empty_like(packed), + "physical_lengths": t.empty_like(lengths), + "staged_order": tuple(range(self.batch)), + } + pcm = c.pcm[: self.batch * frames * SAMPLES_PER_FRAME].view( + self.batch, frames * SAMPLES_PER_FRAME + ) + event = t.cuda.Event() + bindings = [ + c.codec.bind_buffers( + packed, + lengths, + codec_slots, + c.banks[bank], + c.banks[1 - bank], + pcm, + c.samples, + c.status, + stream=stream, + completion_event=event, + ) + for bank in (0, 1) + ] + self.codec_buffers[frames] = (packed, lengths, index, offsets, mask, bindings) + self.neural.stream.wait_stream(t.cuda.current_stream(0)) + stream.wait_stream(t.cuda.current_stream(0)) + + def _check_open(self): + if self.closed or self.poisoned: + raise RuntimeError("in-flight engine is closed or failed") + + def _diagnostic(self, kind, **details): + if getattr(self, "policy", DEFAULT_POLICY).diagnostics: + row = { + "type": kind, + "physical_batch": self.batch, + "host_monotonic_ns": time.monotonic_ns(), + **details, + } + if "slots" in details: + row["owners"] = [ + { + "slot": index, + "epoch": self.slots[index].epoch, + "request_id": self.slots[index].request_id, + } + for index in details["slots"] + if self.slots[index] is not None + ] + if len(self.diagnostic_events) == self.diagnostic_events.maxlen: + self.diagnostic_dropped += 1 + self.diagnostic_events.append(row) + + def drain_diagnostics(self): + rows = list(self.diagnostic_events) + self.diagnostic_events.clear() + if self.diagnostic_dropped: + rows.append({"type": "diagnostic_overflow", "events": self.diagnostic_dropped}) + self.diagnostic_dropped = 0 + return rows + + def _slot(self, index, epoch): + self._check_open() + if type(index) is not int or not 0 <= index < self.batch: + raise ValueError("invalid physical slot") + slot = self.slots[index] + if slot is None or slot.epoch != epoch: + raise ValueError("stale or absent request slot") + return slot + + def _reset_range(self, start, end): + n, c = self.neural, self.core + rows = slice(start, end) + for value, initial in zip(self.conv, self.conv_prefix, strict=True): + value[rows].copy_(initial) + for value, initial in zip(self.ssm, self.ssm_prefix, strict=True): + value[rows].copy_(initial) + for value, initial in zip(self.kv_prefix, self.kv_initial, strict=True): + value[rows].copy_(initial) + self.seq_len[rows].fill_(n.inputs.prompt_positions) + self.slot_map[rows].copy_(self.slot_ids[rows]) + for bank in c.banks: + for value in bank: + value[rows].zero_() + i = n.inputs + i.text_tokens[rows].copy_(self.tokens_host[rows], non_blocking=True) + i.text_lengths[rows].copy_(self.admission_host[0, rows], non_blocking=True) + i.max_text_offsets[rows].copy_(self.admission_host[1, rows], non_blocking=True) + self.noise.seeds[rows].copy_(self.admission_host[2, rows], non_blocking=True) + for name in ( + "offsets", + "last_codes", + "last_phoneme", + "phoneme_ended", + "audio_valid", + "stop", + ): + getattr(i, name)[rows].zero_() + n.codes[rows].zero_() + n.phonemes[rows].zero_() + self.engine.counts[rows].zero_() + self.engine.consumed[rows].zero_() + c.done[rows].zero_() + + def _reset_slot(self, slot): + self._reset_range(slot, slot + 1) + + def _capture_admissions(self): + t, stream = self.t, self.neural.stream + self.admission_host = t.zeros((3, self.batch), dtype=t.int64, pin_memory=True) + self.admission_host_array = self.admission_host.numpy() + self.admission_host[0].fill_(1) + self.admission_graphs = {} + self.admission_events = {} + self.admission_done = [None] * self.batch + self.admission_pending = [False] * self.batch + with t.cuda.stream(stream): + self.slot_ids = t.arange(self.batch, dtype=t.int32, device=self.neural.inputs.device) + self._reset_range(0, self.batch) + stream.synchronize() + if self.batch > 32: + from .inflight_reset import IndexedReset + + self.indexed_reset = IndexedReset(self) + return + for width in (1 << shift for shift in range(self.batch.bit_length())): + for start in range(0, self.batch, width): + key = (start, start + width) + graph = t.cuda.CUDAGraph() + with t.cuda.graph(graph, stream=stream): + self._reset_range(*key) + self.admission_graphs[key] = graph + self.admission_events[key] = t.cuda.Event() + + def _validate_admission(self, slot, epoch, *, text, token_ids, seed, request_id): + self._check_open() + if type(slot) is not int or not 0 <= slot < self.batch or self.slots[slot] is not None: + raise ValueError("admission requires an unused physical slot") + if type(epoch) is not int or epoch <= self.epochs[slot]: + raise ValueError("slot epoch must increase on every admission") + if ( + not isinstance(request_id, str) + or not request_id + or not isinstance(text, str) + or not text + ): + raise ValueError("request text and identity must be nonempty") + if type(seed) is not int or not 0 <= seed < 1 << 64: + raise ValueError("seed must be an unsigned 64-bit integer") + request = Request(text, tuple(token_ids), tuple(self.core.requests[0].context_ids)) + validate_requests([request]) + + def admit(self, slot, epoch, *, text, token_ids, seed, request_id): + self.admit_many( + [ + dict( + slot=slot, + epoch=epoch, + text=text, + token_ids=token_ids, + seed=seed, + request_id=request_id, + ) + ] + ) + + def admit_many(self, records): + self._check_open() + self._observe() + records = list(records) + if len({record["slot"] for record in records}) != len(records): + raise ValueError("admission slots must be unique") + for record in records: + self._validate_admission(**record) + if not records: + return + ranges = reset_ranges((record["slot"] for record in records), self.batch) + indexed = getattr(self, "indexed_reset", None) + live = any(value is not None and value.generation > 0 for value in self.slots) + try: + self._diagnostic("reset_wait_start", slots=[record["slot"] for record in records]) + if indexed is not None: + indexed.wait() + pending = { + self.admission_done[record["slot"]] + for record in records + if self.admission_pending[record["slot"]] + } + for event in pending: + event.synchronize() + self._diagnostic("reset_wait_end") + staging_rows = ( + [record["slot"] for record in records] + if indexed is None + else list(range(len(records))) + ) + lengths = [len(record["token_ids"]) for record in records] + seeds = [ + record["seed"] if record["seed"] < 1 << 63 else record["seed"] - (1 << 64) + for record in records + ] + self.admission_host_array[0, staging_rows] = lengths + self.admission_host_array[1, staging_rows] = [length - 1 for length in lengths] + self.admission_host_array[2, staging_rows] = seeds + for position, record in enumerate(records): + slot = record["slot"] + self.admission_pending[slot] = False + row = self.tokens_host_array[staging_rows[position]] + row.fill(0) + row[: len(record["token_ids"])] = record["token_ids"] + with self.t.cuda.stream(self.neural.stream): + if indexed is not None: + indexed.submit(records) + else: + for key in ranges: + self.admission_graphs[key].replay() + event = self.admission_events[key] + event.record(self.neural.stream) + for slot in range(*key): + self.admission_done[slot] = event + self.admission_pending[slot] = True + for record in records: + slot = record["slot"] + self.slots[slot] = Slot( + record["epoch"], record["request_id"], record["seed"], time.perf_counter() + ) + self.epochs[slot] = record["epoch"] + self.live_refills += int(live) + self.trace.append( + { + "type": "admit", + "slot": slot, + "epoch": record["epoch"], + "request_id": record["request_id"], + "host_monotonic_ns": time.monotonic_ns(), + } + ) + self._diagnostic("reset_enqueued", slots=[record["slot"] for record in records]) + except BaseException: + self.poisoned = True + raise + + def _playback_queued_rows(self): + pending = self.playback_codec_pending + actual = getattr(self.core.codec, "_ticket", None) + packed = getattr(self, "codec_ticket_rows", None) + if pending is None: + return actual is None and packed is None, set() + ticket, _, frames, owners = pending + if actual is not ticket or ticket.owner is not self.core.codec: + return False, set() + selected = tuple(row[0] for row in owners) + if len(set(selected)) != len(selected): + return False, set() + for lane, epoch, request_id, size in owners: + if not 0 <= lane < self.batch or not 0 < size <= frames: + return False, set() + slot = self.slots[lane] + if slot is None or (slot.epoch, slot.request_id) != (epoch, request_id): + return False, set() + if packed is not None: + if packed[0] is not ticket or packed[3] != selected: + return False, set() + if len(packed[2]) != self.batch or any( + packed[2][lane] != size for lane, _, _, size in owners + ): + return False, set() + return True, set(selected) + + def _playback_reserve_choice(self, progress, selected, repeats, now): + if selected == progress and repeats == 1: + self.playback_reserve_last = {"reason": "already_shared_single_step"} + return selected, repeats + published = [ + index + for index, slot in enumerate(self.slots) + if slot is not None and not slot.finished and slot.first_audio is not None + ] + if not published: + self.playback_reserve_last = {"reason": "no_published_owners"} + return selected, repeats + valid, queued = self._playback_queued_rows() + if valid: + published = [ + index + for index in published + if self.slots[index].reason is None + or self.slots[index].produced != self.slots[index].consumed + or index in queued + ] + if not published: + self.playback_reserve_last = {"reason": "no_pending_published_owners"} + return selected, repeats + ready = [ + index + for index in published + if index not in queued + and self.slots[index].produced > self.slots[index].consumed + and ( + self.slots[index].reason is not None + or self.slots[index].produced - self.slots[index].consumed + >= self.slots[index].target(self.schedule) + ) + ] + capped = bool(queued and ready) + if capped: + repeats = 1 + kind = ( + "startup" + if self.startup_entry is not None + and all(self.slots[index].generation < 5 for index in selected) + else "full" + ) + spans = [ + value + for (entry, count), value in self.playback_producer_envelopes.items() + if entry == kind and count >= repeats + ] + full = self.playback_producer_envelopes.get(("full", 1)) + codecs = self.playback_codec_envelopes + calibrated = all(frames in codecs for frames in self.codec_buffers) + reason = "reserve_clear" + if not valid: + reason = "queued_owner_unknown" + elif self.observation_pending: + reason = "producer_metadata_pending" + elif not spans or not calibrated: + reason = "uncalibrated_service" + elif any( + not math.isfinite(value) or value < 0 + for value in (*self.playback_producer_envelopes.values(), *codecs.values()) + ): + reason = "invalid_service_envelope" + evaluations = [] + if reason == "reserve_clear": + quantum = max(spans) + intervening = max(self.playback_producer_envelopes.values()) + codec = max(codecs[frames] for frames in self.codec_buffers) + for index in published: + slot = self.slots[index] + pending = slot.produced - slot.consumed + missing = ( + 0 + if index in queued or slot.reason is not None + else max(0, slot.target(self.schedule) - pending) + ) + if pending < 0 or (missing and full is None): + reason = "uncalibrated_service" if pending >= 0 else "invalid_owner_state" + break + class_turns = len(self.codec_buffers) * ( + 2 if self.policy.codec_deadline_priority else 1 + ) + turns = 1 if index in queued else class_turns + bool(queued) + reserve = ( + quantum + + missing * (full or 0.0) + + (missing + turns) * codec + + (missing + max(0, turns - 1)) * intervening + + 0.040 + ) + lead = slot.first_audio + slot.samples / SAMPLE_RATE - now + evaluations.append((index, slot.epoch, slot.request_id, missing, lead, reserve)) + if not math.isfinite(lead) or not math.isfinite(reserve) or lead < reserve: + reason = "playback_reserve" + self.playback_reserve_last = { + "reason": reason, + "published": published, + "queued": sorted(queued), + "ready": ready, + "prefetch_cap": capped, + "owners": evaluations, + } + if reason != "reserve_clear": + return progress, 1 + return selected, repeats + + def _playback_observed_producer(self): + pending = self.playback_producer_pending + if pending is not None: + started, kind, repeats = pending + key = (kind, repeats) + elapsed = time.perf_counter() - started + if math.isfinite(elapsed) and elapsed >= 0: + self.playback_producer_envelopes[key] = max( + elapsed, self.playback_producer_envelopes.get(key, 0.0) + ) + self.playback_producer_pending = None + + def _playback_finished_codec(self, ticket): + pending = self.playback_codec_pending + if pending is not None and pending[0] is ticket: + _, started, frames, owners = pending + if all( + self.slots[lane] is not None + and (self.slots[lane].epoch, self.slots[lane].request_id) == (epoch, request_id) + for lane, epoch, request_id, _ in owners + ): + elapsed = time.perf_counter() - started + if math.isfinite(elapsed) and elapsed >= 0: + self.playback_codec_envelopes[frames] = max( + elapsed, + self.playback_codec_envelopes.get(frames, 0.0), + ) + self.playback_codec_pending = None + + def _step(self, progress, *, repeats=1): + if self.observation_pending: + raise RuntimeError("previous producer metadata must be observed before mask reuse") + if type(repeats) is not int or not 1 <= repeats <= startup_steps( + self.slots, progress, self.steps + ): + raise ValueError("producer repetition would cross a request control boundary") + if getattr(self, "policy", DEFAULT_POLICY).playback_reserve_guard: + kind = ( + "startup" + if self.startup_entry is not None + and all(self.slots[index].generation < 5 for index in progress) + else "full" + ) + self.playback_producer_pending = (time.perf_counter(), kind, repeats) + t, n = self.t, self.neural + self.progress_host_array.fill(0) + self.progress_host_array[progress] = 1 + with t.cuda.stream(n.stream): + self.progress.copy_(self.progress_host, non_blocking=True) + for _ in range(repeats): + self.prepare_graph.replay() + launch = n.lib.mk2_run_native_stage0 + if self.startup_entry is not None and all( + self.slots[index].generation < 5 for index in progress + ): + launch = self.startup_entry + n._check( + launch(n.stage0_handle, C.byref(n.bindings), n.stream.cuda_stream), + "in-flight Stage0", + ) + self.commit_graph.replay() + for index in progress: + self.slots[index].generation += 1 + self.producer_launches += 1 + self._diagnostic("producer_enqueued", slots=list(progress), repeats=repeats) + self.trace.append( + { + "type": "producer", + "tick": self.tick_count, + "slots": list(progress), + "startup_batch_steps": repeats, + "host_monotonic_ns": time.monotonic_ns(), + } + ) + self.meta_host.copy_(self.meta, non_blocking=True) + self.meta_done.record(n.stream) + self.observation_pending = True + + def _observe(self): + if not self.observation_pending: + return + try: + self._diagnostic("producer_wait_start") + self.meta_done.synchronize() + self._diagnostic("producer_wait_end") + except BaseException: + self.poisoned = True + raise + counts, done = self.meta_host.tolist() + for index, slot in enumerate(self.slots): + if slot is None or slot.finished: + continue + slot.produced = counts[index] + if done[index]: + slot.reason = "eos" + elif slot.generation >= self.steps: + slot.reason = "length_limit" + self.observation_pending = False + if getattr(self, "policy", DEFAULT_POLICY).playback_reserve_guard: + self._playback_observed_producer() + + def _progress(self, requested): + return [ + index + for index in requested + if (slot := self.slots[index]) is not None + and not slot.finished + and slot.reason is None + and slot.produced - slot.consumed < slot.target(self.schedule) + ] + + def _advance_producer(self, progress): + now = time.perf_counter() + policy = getattr(self, "policy", DEFAULT_POLICY) + started = getattr(self, "startup_priority_started", None) + defer = getattr(self, "startup_priority_deferred", False) or ( + started is not None and now - started >= policy.startup_budget_ms / 1000 + ) + selected = startup_progress( + self.slots, + progress, + now, + defer=defer or policy.producer_priority == "shared", + align_phases=policy.startup_phase_alignment, + ) + priority = selected != progress + repeats = startup_steps(self.slots, selected, self.steps) + if policy.partial_startup_single_step and any(slot is None for slot in self.slots): + repeats = 1 + if policy.playback_reserve_guard: + if priority: + repeats = min(repeats, policy.startup_max_steps - self.startup_priority_steps) + selected, repeats = self._playback_reserve_choice(progress, selected, repeats, now) + priority = selected != progress + if priority: + if started is None: + self.startup_priority_started = now + remaining = policy.startup_max_steps - getattr(self, "startup_priority_steps", 0) + repeats = min(repeats, remaining) + self.startup_priority_steps = policy.startup_max_steps - remaining + repeats + self.startup_priority_deferred = self.startup_priority_steps >= policy.startup_max_steps + self.startup_priority_launches += repeats + else: + self.startup_priority_deferred = False + self.startup_priority_steps = 0 + self.startup_priority_started = None + if repeats == 1: + self._step(selected) + else: + self._step(selected, repeats=repeats) + + def _submit_codec(self, frames, selected): + reserve_started = ( + time.perf_counter() + if getattr(self, "policy", DEFAULT_POLICY).playback_reserve_guard + else None + ) + t, n, e = self.t, self.neural, self.engine + if getattr(self, "codec_ticket_rows", None) is not None: + raise RuntimeError("codec row map cannot be reused before ticket completion") + packed, lengths, index, offsets, mask, bindings = self.codec_buffers[frames] + packing = getattr(self, "codec_packing", {}).get(frames) + packed_transfer = getattr(self, "policy", DEFAULT_POLICY).codec_packed_transfer + if packed_transfer and packing is None: + raise RuntimeError("packed codec transfer requires an owned row map") + order, inverse = ( + codec_row_order(selected, self.batch) if packing is not None else (None, None) + ) + reordered = packing is not None and order != tuple(range(self.batch)) + physical_codes = packing["physical_codes"] if reordered else packed + physical_lengths = packing["physical_lengths"] if reordered else lengths + sizes = [0] * self.batch + for lane in selected: + slot = self.slots[lane] + sizes[lane] = min(slot.target(self.schedule), slot.produced - slot.consumed) + with t.cuda.stream(n.stream): + if packing is not None and order != packing["staged_order"]: + packing["order"].copy_(t.tensor(order, dtype=t.int64)) + packing["slots"].copy_(packing["order"]) + packing["staged_order"] = order + physical_lengths.copy_(t.tensor(sizes, dtype=t.int32)) + t.add(e.consumed.view(self.batch, 1), offsets, out=index) + index.clamp_(max=self.steps - 1) + t.gather( + e.queue, 1, index.unsqueeze(-1).expand(self.batch, frames, 16), out=physical_codes + ) + t.lt(offsets, physical_lengths.view(self.batch, 1), out=mask) + physical_codes.mul_(mask.unsqueeze(-1)) + e.consumed.add_(physical_lengths) + if reordered: + t.index_select(physical_codes, 0, packing["order"], out=packed) + t.index_select(physical_lengths, 0, packing["order"], out=lengths) + self.ready.record(n.stream) + for lane in selected: + self.slots[lane].consumed += sizes[lane] + e.codec_stream.wait_event(self.ready) + with t.cuda.stream(e.codec_stream): + ticket = bindings[self.bank].enqueue() + if packing is not None: + self.codec_ticket_rows = (ticket, inverse, tuple(sizes), tuple(selected)) + self.codec_packing_launches += int(reordered) + if packed_transfer: + e.submit_pcm_transfer(ticket, row_count=len(selected)) + else: + e.submit_pcm_transfer(ticket) + self._diagnostic("codec_enqueued", slots=selected, frames=frames) + self.codec_launches += 1 + self.trace.append( + { + "type": "codec", + "tick": self.tick_count, + "frames": frames, + "slots": selected, + "codec_reordered": reordered, + "host_monotonic_ns": time.monotonic_ns(), + } + ) + if reserve_started is not None: + self.playback_codec_pending = ( + ticket, + reserve_started, + frames, + tuple( + (lane, self.slots[lane].epoch, self.slots[lane].request_id, sizes[lane]) + for lane in selected + ), + ) + return ticket, sizes + + def _finish_codec(self, ticket, sizes, selected): + import numpy as np + + owner = getattr(self, "codec_ticket_rows", None) + inverse = range(self.batch) + if owner is not None: + if owner[0] is not ticket or owner[2:] != (tuple(sizes), tuple(selected)): + raise RuntimeError("codec row map ticket ownership mismatch") + inverse = owner[1] + self._diagnostic("codec_wait_start", slots=selected) + result = self.core.codec.finish(ticket, copy_pcm=True) + self._diagnostic("codec_wait_end", slots=selected) + pcm = result["pcm"].numpy() + finite = result.get("finite") + if finite is None: + finite = bool(np.isfinite(pcm).all()) + if not finite: + raise RuntimeError("in-flight codec produced nonfinite PCM") + samples = result["samples"].cpu().tolist() + if owner is not None: + samples = [samples[row] for row in inverse] + if samples != [size * SAMPLES_PER_FRAME for size in sizes]: + raise RuntimeError("in-flight codec sample-count mismatch") + if getattr(self, "policy", DEFAULT_POLICY).codec_packed_transfer: + if pcm.ndim != 2 or any( + not 0 <= inverse[lane] < pcm.shape[0] or samples[lane] > pcm.shape[1] + for lane in selected + ): + raise RuntimeError("packed PCM prefix does not cover selected codec rows") + self.bank = 1 - self.bank + events = [] + for lane in selected: + slot = self.slots[lane] + sample_count = samples[lane] + payload = pcm[inverse[lane], :sample_count].astype(" 6 or destination.element_size() not in (1, 2, 4, 8): + raise ValueError("unsupported indexed reset layout") + source_stride = 0 + if source is not None: + if not source.is_contiguous() or source.element_size() != destination.element_size(): + raise ValueError("indexed reset source must have contiguous matching elements") + if source.numel() != math.prod(shape): + if tuple(source.shape) != tuple(destination.shape): + raise ValueError("indexed reset source must be one row or a full physical tensor") + source_stride = source.stride(0) + return ResetLayout( + destination.stride(0), + source_stride, + math.prod(shape), + len(shape), + destination.element_size(), + int(compact), + (C.c_uint64 * 6)(*shape), + (C.c_uint64 * 6)(*strides), + ) + + +class IndexedReset: + def __init__(self, owner): + self.owner = owner + t, n, b = owner.t, owner.neural, owner.batch + self.pending = False + self.done = t.cuda.Event() + self.ids_host = t.empty(b, dtype=t.int32, pin_memory=True) + self.count_host = t.empty(1, dtype=t.int32, pin_memory=True) + self.ids_host_array = self.ids_host.numpy() + self.count_host_array = self.count_host.numpy() + self.launch = n.lib.mk2_inflight_reset_rows + self.launch.argtypes = [ + C.c_void_p, + C.c_void_p, + C.c_void_p, + C.c_void_p, + C.c_int, + C.c_uint64, + C.POINTER(ResetLayout), + C.c_void_p, + ] + self.launch.restype = C.c_int + with t.cuda.stream(n.stream): + self.ids = t.zeros(b, dtype=t.int32, device=n.inputs.device) + self.count = t.zeros(1, dtype=t.int32, device=n.inputs.device) + self.tokens = t.empty_like(owner.tokens_host, device=n.inputs.device) + self.metadata = t.empty_like(owner.admission_host, device=n.inputs.device) + self.operations = [] + for targets, sources in ( + (owner.conv, owner.conv_prefix), + (owner.ssm, owner.ssm_prefix), + (owner.kv_prefix, owner.kv_initial), + ): + for target, source in zip(targets, sources, strict=True): + self.add(target, source) + self.add(owner.seq_len, fill=n.inputs.prompt_positions) + self.add(owner.slot_map, owner.slot_ids) + for bank in owner.core.banks: + for value in bank: + self.add(value) + inputs = n.inputs + self.add(inputs.text_tokens, self.tokens, compact=True) + self.add(inputs.text_lengths, self.metadata[0], compact=True) + self.add(inputs.max_text_offsets, self.metadata[1], compact=True) + self.add(owner.noise.seeds, self.metadata[2], compact=True) + for name in ( + "offsets", + "last_codes", + "last_phoneme", + "phoneme_ended", + "audio_valid", + "stop", + ): + self.add(getattr(inputs, name)) + for value in ( + n.codes, + n.phonemes, + owner.engine.counts, + owner.engine.consumed, + owner.core.done, + ): + self.add(value) + self.run() + n.stream.synchronize() + self.graph = t.cuda.CUDAGraph() + with t.cuda.graph(self.graph, stream=n.stream): + self.run() + n.stream.synchronize() + + def add(self, target, source=None, *, fill=0, compact=False): + self.operations.append( + (target, source, fill, reset_layout(target, source, compact=compact)) + ) + + def run(self): + n = self.owner.neural + for target, source, fill, layout in self.operations: + n._check( + self.launch( + target.data_ptr(), + None if source is None else source.data_ptr(), + self.ids.data_ptr(), + self.count.data_ptr(), + self.owner.batch, + fill, + C.byref(layout), + n.stream.cuda_stream, + ), + "indexed request-state reset", + ) + + def wait(self): + if self.pending: + self.done.synchronize() + self.pending = False + + def submit(self, records): + if self.pending: + raise RuntimeError("indexed reset staging is still in use") + count = len(records) + if not 0 < count <= self.owner.batch: + raise ValueError("indexed reset requires nonempty bounded admission") + self.ids_host_array[:count] = [record["slot"] for record in records] + self.count_host_array[0] = count + o = self.owner + with o.t.cuda.stream(o.neural.stream): + self.ids[:count].copy_(self.ids_host[:count], non_blocking=True) + self.count.copy_(self.count_host, non_blocking=True) + self.tokens[:count].copy_(o.tokens_host[:count], non_blocking=True) + for index in range(3): + self.metadata[index, :count].copy_( + o.admission_host[index, :count], non_blocking=True + ) + self.graph.replay() + self.done.record(o.neural.stream) + self.pending = True + + def close(self): + self.wait() + self.graph.reset() diff --git a/python/nemotron_native_backend/execution/lane_control.py b/python/nemotron_native_backend/execution/lane_control.py new file mode 100644 index 0000000..4ce134b --- /dev/null +++ b/python/nemotron_native_backend/execution/lane_control.py @@ -0,0 +1,37 @@ +def update_first_eos(first_eos, stop_rows, start): + if len(first_eos) != len(stop_rows) or start < 0: + raise ValueError("invalid EOS observation geometry") + updated = list(first_eos) + for lane, row in enumerate(stop_rows): + if any(type(x) is not bool for x in row): + raise ValueError("EOS observations must be boolean") + if updated[lane] is None: + updated[lane] = next((start + j for (j, stop) in enumerate(row) if stop), None) + return updated + + +def completion_summary(first_eos, executed): + if ( + executed < 0 + or not first_eos + or any(g is not None and (not 0 <= g < executed) for g in first_eos) + ): + raise ValueError("EOS generation lies outside executed work") + rows = [ + { + "lane": i, + "first_eos_generation": g, + "completed": g is not None, + "post_eos_computed_lane_steps": 0 if g is None else executed - g - 1, + } + for (i, g) in enumerate(first_eos) + ] + all_done = all(x["completed"] for x in rows) + return { + "all_lanes_completed": all_done, + "completed_lanes": sum(x["completed"] for x in rows), + "cohort_completion_generation": max(first_eos) if all_done else None, + "request_completion": rows, + "executed_lane_steps": len(rows) * executed, + "discarded_post_eos_lane_steps": sum(x["post_eos_computed_lane_steps"] for x in rows), + } diff --git a/python/nemotron_native_backend/execution/lane_windows.py b/python/nemotron_native_backend/execution/lane_windows.py new file mode 100644 index 0000000..815358d --- /dev/null +++ b/python/nemotron_native_backend/execution/lane_windows.py @@ -0,0 +1,100 @@ +from __future__ import annotations + + +class LaneWindow: + def __init__(self, batch, frames, *, done, device="cuda:0", raw=None, stops=None): + import torch + + if ( + batch not in (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096) + or not 1 <= frames <= 16 + ): + raise ValueError("invalid window geometry") + self.torch = torch + self.batch = batch + self.frames = frames + (B, F) = (batch, frames) + device = torch.device(device) + if ( + done.dtype != torch.bool + or tuple(done.shape) != (B,) + or done.device != device + or (not done.is_contiguous()) + ): + raise ValueError("done must be contiguous bool[B] on the window device") + self.done = done + + def empty(shape, dtype): + return torch.empty(shape, dtype=dtype, device=device) + + self.raw = empty((B, F, 16), torch.int64) if raw is None else raw + self.stops = empty((B, F), torch.bool) if stops is None else stops + if ( + tuple(self.raw.shape) != (B, F, 16) + or tuple(self.stops.shape) != (B, F) + or self.raw.dtype != torch.int64 + or (self.stops.dtype != torch.bool) + or (self.raw.device != device) + or (self.stops.device != device) + ): + raise ValueError("borrowed history shape/dtype/device mismatch") + self.packed = empty((B, F, 16), torch.int64) + self.lengths = empty((B,), torch.int32) + self.element_ok = empty((B, F, 16), torch.bool) + self.element_tmp = empty((B, F, 16), torch.bool) + self.valid = empty((B, F), torch.bool) + self.nonzero = empty((B, F), torch.bool) + self.stop_prefix = empty((B, F), torch.int32) + self.stop_prefix_before = empty((B, F), torch.int32) + self.before_stop = empty((B, F), torch.bool) + self.was_active = empty((B,), torch.bool) + self.window_eos = empty((B,), torch.bool) + self.rank = empty((B, F), torch.int64) + self.matches = empty((B, F, F), torch.bool) + self.has_match = empty((B, F), torch.bool) + self.source_index = empty((B, F), torch.int64) + self.targets = torch.arange(1, F + 1, dtype=torch.int64, device=device).view(1, F, 1) + self._rank3 = self.rank.view(B, 1, F) + self._valid3 = self.valid.view(B, 1, F) + self._active2 = self.was_active.view(B, 1) + self._gather_index = self.source_index.unsqueeze(-1).expand(B, F, 16) + self._has_match3 = self.has_match.unsqueeze(-1) + self.raw_columns = tuple(self.raw[:, i, :] for i in range(F)) + self.stop_columns = tuple(self.stops[:, i] for i in range(F)) + self.fill = 0 + + def append(self, codes, stop): + if self.fill >= self.frames: + raise RuntimeError("window capacity exceeded") + self.raw_columns[self.fill].copy_(codes) + self.stop_columns[self.fill].copy_(stop) + self.fill += 1 + + def pack(self): + if self.fill != self.frames: + raise RuntimeError("window must be full before packing") + t = self.torch + t.ge(self.raw, 0, out=self.element_ok) + t.lt(self.raw, 1024, out=self.element_tmp) + t.logical_and(self.element_ok, self.element_tmp, out=self.element_ok) + t.all(self.element_ok, dim=2, out=self.valid) + t.ne(self.raw, 0, out=self.element_tmp) + t.any(self.element_tmp, dim=2, out=self.nonzero) + t.logical_and(self.valid, self.nonzero, out=self.valid) + t.logical_not(self.done, out=self.was_active) + t.cumsum(self.stops, dim=1, dtype=t.int32, out=self.stop_prefix) + self.stop_prefix_before[:, :1].zero_() + self.stop_prefix_before[:, 1:].copy_(self.stop_prefix[:, :-1]) + t.eq(self.stop_prefix_before, 0, out=self.before_stop) + t.logical_and(self.before_stop, self._active2, out=self.before_stop) + t.logical_and(self.valid, self.before_stop, out=self.valid) + t.cumsum(self.valid, dim=1, dtype=t.int64, out=self.rank) + t.eq(self._rank3, self.targets, out=self.matches) + t.logical_and(self.matches, self._valid3, out=self.matches) + t.max(self.matches, dim=2, out=(self.has_match, self.source_index)) + t.gather(self.raw, 1, self._gather_index, out=self.packed) + self.packed.mul_(self._has_match3) + t.sum(self.valid, dim=1, dtype=t.int32, out=self.lengths) + t.any(self.stops, dim=1, out=self.window_eos) + t.logical_or(self.done, self.window_eos, out=self.done) + return (self.packed, self.lengths) diff --git a/python/nemotron_native_backend/execution/narrow_prefix_cache.py b/python/nemotron_native_backend/execution/narrow_prefix_cache.py new file mode 100644 index 0000000..95910e9 --- /dev/null +++ b/python/nemotron_native_backend/execution/narrow_prefix_cache.py @@ -0,0 +1,109 @@ +import hashlib +import time + + +def install(engine): + n = engine.core.neural + t = engine.t + initial = getattr(n, "_r8_initial_state", None) + if initial is None or n.started or engine.batch != 32 or hasattr(engine, "_b32_prefix_owner"): + raise ValueError("canonical prefix requires a fresh B32 canonical owner") + private_bytes = engine._b32_prefill_state.get("private_state_bytes") + if private_bytes not in (14 * 32 * 196608 * 2, 14 * 32 * 196608 * 4): + raise ValueError("unrecognized private recurrent format") + + def view(pointer, count): + class Alias: + def __init__(self): + self.__cuda_array_interface__ = { + "shape": (count,), + "typestr": "|u1", + "data": (int(pointer), False), + "version": 3, + "strides": None, + } + + return t.as_tensor(Alias(), device="cuda:0") + + b = n.batch + p = n.inputs.prompt_positions + blocks = n.info.blocks_per_sequence + original = n.prefill_lib.mk2_native_prefill_run + version = n.inputs.prompt._version + context = tuple(tuple(request.context_ids) for request in n.inputs.requests) + identity = ( + n.prefill_handle, + n.stage0_handle, + tuple(int(initial["initial"][i]) for i in range(14)), + initial["ready"].value, + tuple(int(n.bindings.ssm_state[i]) for i in range(14)), + ) + views = [n.final_hidden, view(n.bindings.slot_map, b * 4), view(n.bindings.seq_len, b * 4)] + for layer in range(14): + views.append(view(n.bindings.conv_state[layer], b * n.bindings.conv_slot_stride[layer] * 2)) + views.append(view(initial["initial"][layer], 196608 * 4)) + for layer in range(3): + views.append(view(n.bindings.kv_block_table[layer], b * 2 * blocks * 4)) + pool = view(n.bindings.kv_pool_base[layer], b * 2 * blocks * 4 * 32 * 128 * 2).reshape( + b, 2, blocks, 4, 32, 256 + ) + for block in range((p + 31) // 32): + views.append(pool[:, :, block, :, : min(32, p - block * 32), :]) + ready = view(initial["ready"].value, 4) + views.append(ready) + started = time.monotonic() + with t.cuda.stream(n.stream): + n._check( + original(n.prefill_handle, n.stream.cuda_stream), + "canonical prefix preparation", + n.prefill_lib, + ) + snapshots = [value.clone() for value in views] + n.stream.synchronize() + if ready.view(t.int32).item() != 1: + raise ValueError("prefix did not establish a canonical initial state") + receipt = { + "kind": "canonical-FP32-prefix", + "private_recurrent_storage_restored": False, + "bytes": sum(value.numel() * value.element_size() for value in snapshots), + "preparation_seconds": time.monotonic() - started, + "hits": 0, + "misses": 0, + "prefill_sha256": n.prefill_library_sha256, + "stage0_sha256": n.library_sha256, + "prompt_sha256": hashlib.sha256( + n.inputs.prompt.cpu().view(t.uint8).numpy().tobytes() + ).hexdigest(), + } + + def run(handle, stream): + current = ( + n.prefill_handle, + n.stage0_handle, + tuple(int(initial["initial"][i]) for i in range(14)), + initial["ready"].value, + tuple(int(n.bindings.ssm_state[i]) for i in range(14)), + ) + if handle != n.prefill_handle or stream != n.stream.cuda_stream or current != identity: + raise RuntimeError("canonical prefix owner or stream changed") + if ( + n.inputs.prompt._version != version + or tuple(tuple(request.context_ids) for request in n.inputs.requests) != context + ): + receipt["misses"] += 1 + return original(handle, stream) + with t.cuda.stream(n.stream): + for destination, snapshot in zip(views, snapshots, strict=True): + destination.copy_(snapshot) + receipt["hits"] += 1 + return 0 + + n.prefill_lib.mk2_native_prefill_run = run + engine._b32_prefix_owner = (views, snapshots, original, run) + engine._b32_canonical_cache = receipt + + engine._b32_canonical_cache["ownership_contract"] = ( + "exclusive fixed-allocation engine ownership; supported prompt/context reset " + "invalidates; checkpoint, precision and state-format changes require a fresh " + "engine/cache" + ) diff --git a/python/nemotron_native_backend/execution/narrow_state_binding.py b/python/nemotron_native_backend/execution/narrow_state_binding.py new file mode 100644 index 0000000..8943202 --- /dev/null +++ b/python/nemotron_native_backend/execution/narrow_state_binding.py @@ -0,0 +1,70 @@ +import copy +import ctypes as C +import hashlib +from pathlib import Path + +from .engine import Stage0BuildInfo + + +def install(engine, narrow_library, narrow_sha256, artifact_root): + path = Path(narrow_library) + if hashlib.sha256(path.read_bytes()).hexdigest() != narrow_sha256: + raise ValueError("BF16 narrow library changed") + n = engine.core.neural + probe = n.lib.mk2_b32_private_state_bytes + probe.argtypes = [] + probe.restype = C.c_int + if probe() != 2: + raise ValueError("BF16 full producer format mismatch") + initial = (C.c_void_p * 14)() + ready = C.c_void_p() + get = n.prefill_lib.mk2_r8_initial_state + get.argtypes = [C.c_void_p, C.POINTER(C.c_void_p), C.POINTER(C.c_void_p)] + get.restype = C.c_int + n._check(get(n.prefill_handle, initial, C.byref(ready)), "B32 BF16 canonical source") + for i in range(14): + private = int(n.bindings.ssm_state[i]) + if private <= int(initial[i]) < private + n.batch * 196608 * 2: + raise ValueError("canonical state aliases BF16 private state") + setter = n.lib.mk2_r8_bind_initial_state + setter.argtypes = [C.c_void_p, C.POINTER(C.c_void_p), C.c_void_p, C.c_void_p, C.c_int] + setter.restype = C.c_int + n._check( + setter(n.stage0_handle, initial, ready, n.bindings.seq_len, n.inputs.prompt_positions), + "B32 BF16 canonical binding", + ) + n._r8_initial_state = { + "initial": initial, + "ready": ready, + "bytes": 14 * 196608 * 4, + "prefill_owner": n.prefill_handle, + } + engine._b32_prefill_state = {"canonical_bytes": 14 * 196608 * 4} + old = engine.contraction + if old is None or old.enabled: + raise ValueError("fresh retained contraction owner required") + config = copy.deepcopy(engine.active_lanes_config) + config["stage0"] = {"path": str(path.resolve()), "sha256": narrow_sha256} + owner_type = type(old) + old.close() + engine.contraction = owner_type( + n, + config, + artifact_root, + engine.residency["codec_blocks"], + Stage0BuildInfo, + engine.execution, + ) + engine.active_lanes_config = config + probe = engine.contraction.lib.mk2_b32_private_state_bytes + probe.argtypes = [] + probe.restype = C.c_int + if probe() != 2: + raise ValueError("BF16 narrow producer format mismatch") + engine._b32_prefill_state.update( + variant="n1", + private_state_bytes=14 * 32 * 196608 * 2, + prefill_fp32_workspace_bytes=14 * 32 * 196608 * 4, + narrow_library=str(path), + narrow_sha256=narrow_sha256, + ) diff --git a/python/nemotron_native_backend/execution/native_assets.py b/python/nemotron_native_backend/execution/native_assets.py new file mode 100644 index 0000000..21dcafa --- /dev/null +++ b/python/nemotron_native_backend/execution/native_assets.py @@ -0,0 +1,477 @@ +from __future__ import annotations + +import ctypes as C +import hashlib +import json +import math +import struct +from dataclasses import dataclass +from pathlib import Path + +WIDTHS = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096) +PATTERN = "MEMEM*EMEMEM*EMEMEMEM*EMEMEMEME" +HIDDEN = 1536 +EXPECTED_PREFIXES = { + "audio_embeddings": 16, + "decoder": 885, + "local_transformer": 19, + "local_transformer_out_projections": 32, + "phoneme_embeddings": 1, + "phoneme_final_proj": 2, + "text_embedding": 1, +} + + +def sha256(path: Path) -> str: + h = hashlib.sha256() + with path.open("rb") as f: + for block in iter(lambda: f.read(4 << 20), b""): + h.update(block) + return h.hexdigest() + + +def read_header(path: Path) -> tuple[dict, int, str]: + with path.open("rb") as f: + prefix = f.read(8) + if len(prefix) != 8: + raise ValueError(f"truncated safetensors: {path}") + length = struct.unpack("= 0 for v in shape): + raise ValueError(f"invalid tensor shape: {name}") + (start, end) = spec["data_offsets"] + if not 0 <= start <= end <= size or end - start != math.prod(shape) * sizes[spec["dtype"]]: + raise ValueError(f"invalid tensor extent: {name}") + spans.append((start, end)) + spans.sort() + if spans and ( + spans[0][0] != 0 + or spans[-1][1] != size + or any(a[1] != b[0] for (a, b) in zip(spans, spans[1:], strict=False)) + ): + raise ValueError("safetensors has a gap or overlapping payloads") + return (header, length + 8, hashlib.sha256(prefix + payload).hexdigest()) + + +@dataclass(frozen=True) +class WeightSpec: + family: str + field: str + index: int | None + sources: tuple[str, ...] + source_shapes: tuple[tuple[int, ...], ...] + shape: tuple[int, ...] + dtype: str + transform: str = "checkpoint_to_bf16" + + +def weight_specs() -> list[WeightSpec]: + specs = [] + + def add( + family, + field, + index, + source, + shape, + dtype="bfloat16", + transform="checkpoint_to_bf16", + source_shape=None, + ): + specs.append( + WeightSpec( + family, + field, + index, + (source,), + (tuple(source_shape or shape),), + tuple(shape), + dtype, + transform, + ) + ) + + for i in range(3): + p = f"local_transformer.layers.{i}." + for field, suffix, shape, src_shape in ( + ("qkv", "self_attention.qkv_net.weight", (4608, 1536), None), + ("o", "self_attention.o_net.weight", (1536, 1536), None), + ("ff_in", "pos_ff.proj.conv.weight", (6144, 1536), (6144, 1536, 1)), + ("ff_out", "pos_ff.o_net.conv.weight", (1536, 6144), (1536, 6144, 1)), + ): + add("predictor", field, i, p + suffix, shape, source_shape=src_shape) + for field, suffix in (("norm_self", "norm_self.weight"), ("norm_ff", "norm_pos_ff.weight")): + add("predictor", field, i, p + suffix, (1536,), "float32", "checkpoint_to_fp32") + add("predictor", "pos_emb", None, "local_transformer.position_embeddings.weight", (18, 1536)) + for i in range(16): + add("predictor", "audio_emb", i, f"audio_embeddings.{i}.weight", (1032, 1536)) + add("predictor", "head_w", i, f"local_transformer_out_projections.{i}.weight", (1032, 1536)) + add("predictor", "head_b", i, f"local_transformer_out_projections.{i}.bias", (1032,)) + counts = {"M": 0, "E": 0, "*": 0} + for layer, kind in enumerate(PATTERN): + p = f"decoder.layers.{layer}." + add("backbone", "norm", layer, p + "norm.weight", (1536,)) + p += "mixer." + i = counts[kind] + counts[kind] += 1 + if kind == "M": + for field, suffix, shape, src_shape in ( + ("mb_in_proj", "in_proj.weight", (5184, 1536), None), + ("mb_out_proj", "out_proj.weight", (1536, 1536), None), + ("mb_conv_w", "conv1d.weight", (3584, 4), (3584, 1, 4)), + ("mb_conv_b", "conv1d.bias", (3584,), None), + ("mb_gnorm", "norm.weight", (1536,), None), + ): + add("backbone", field, i, p + suffix, shape, source_shape=src_shape) + add("backbone", "mb_A", i, p + "A_log", (64,), "float32", "negative_exp_fp32") + for field, suffix in (("mb_D", "D"), ("mb_dt_bias", "dt_bias")): + add("backbone", field, i, p + suffix, (64,), "float32", "checkpoint_to_fp32") + elif kind == "*": + specs.append( + WeightSpec( + "backbone", + "at_qkv", + i, + tuple(p + x + "_proj.weight" for x in "qkv"), + ((1536, 1536), (512, 1536), (512, 1536)), + (2560, 1536), + "bfloat16", + "concat_qkv_checkpoint_to_bf16", + ) + ) + add("backbone", "at_o", i, p + "o_proj.weight", (1536, 1536)) + else: + for field, suffix, shape in ( + ("moe_exp_up", "up_proj.weight", (768, 1536)), + ("moe_exp_down", "down_proj.weight", (1536, 768)), + ): + specs.append( + WeightSpec( + "backbone", + field, + i, + tuple(p + f"experts.{e}." + suffix for e in range(24)), + (shape,) * 24, + (24, *shape), + "bfloat16", + "stack_numeric_expert_order_checkpoint_to_bf16", + ) + ) + for field, suffix, shape in ( + ("moe_sh_up", "shared_experts.up_proj.weight", (2048, 1536)), + ("moe_sh_down", "shared_experts.down_proj.weight", (1536, 2048)), + ("moe_gate", "gate.weight", (24, 1536)), + ("moe_gate_bias", "gate.e_score_correction_bias", (24,)), + ): + add("backbone", field, i, p + suffix, shape) + add("backbone", "norm_f", None, "decoder.norm_f.weight", (1536,)) + add("backbone", "phoneme_w", None, "phoneme_final_proj.weight", (2051, 1536)) + add("backbone", "phoneme_b", None, "phoneme_final_proj.bias", (2051,)) + add( + "input", + "text_embedding", + None, + "text_embedding.weight", + (131075, 1536), + "bfloat16", + "checkpoint_to_bf16", + ) + add( + "input", + "phoneme_embedding", + None, + "phoneme_embeddings.0.weight", + (2051, 1536), + "bfloat16", + "checkpoint_to_bf16", + ) + add("unused", "backbone_token_embedding", None, "decoder.embeddings.weight", (2, 1536)) + return specs + + +class ModelAssets: + def __init__(self, root: Path): + self.root = Path(root).resolve() + self.config = json.loads((self.root / "config.json").read_text()) + expected = { + "hidden_size": 1536, + "num_hidden_layers": 31, + "num_attention_heads": 12, + "torch_dtype": "bfloat16", + "num_key_value_heads": 4, + "head_dim": 128, + "mamba_num_heads": 64, + "mamba_head_dim": 24, + "ssm_state_size": 128, + "n_groups": 8, + "conv_kernel": 4, + "n_routed_experts": 24, + "num_experts_per_tok": 4, + "moe_intermediate_size": 768, + "moe_shared_expert_intermediate_size": 2048, + "hybrid_override_pattern": PATTERN, + "text_eos_id": 131073, + "text_vocab_size": 131075, + "num_task_embeddings": 0, + "streaming_speech_delay": 5, + "streaming_phonemes_delay": 3, + "phoneme_bos_id": 2048, + "phoneme_eos_id": 2049, + "forced_audio_bos_id": 1024, + "forced_audio_eos_id": 1025, + "local_transformer_n_layers": 3, + "local_transformer_n_heads": 12, + "local_transformer_hidden_dim": 1536, + "audio_embedding_dim": 1536, + "embedding_dim": 1536, + "num_audio_codebooks": 8, + "frame_stacking_factor": 2, + "codebook_size": 1024, + "phoneme_vocab_size": 2051, + "routed_scaling_factor": 2.5, + "norm_topk_prob": True, + "n_group": 1, + "topk_group": 1, + "layer_norm_epsilon": 1e-05, + } + drift = { + k: (self.config.get(k), v) for (k, v) in expected.items() if self.config.get(k) != v + } + if drift: + raise ValueError(f"unsupported model configuration: {drift}") + index = json.loads((self.root / "model.safetensors.index.json").read_text()) + self.weight_map = index["weight_map"] + from collections import Counter + + if dict(Counter(k.split(".")[0] for k in self.weight_map)) != EXPECTED_PREFIXES: + raise ValueError("checkpoint prefix inventory drift") + self.headers = {} + self.shards = {} + for name in sorted(set(self.weight_map.values())): + p = self.root / name + if Path(name).name != name or not p.is_file(): + raise ValueError("checkpoint shard must be a local regular filename") + (header, offset, digest) = read_header(p) + self.shards[name] = { + "path": str(p), + "bytes": p.stat().st_size, + "safetensors_header_bytes": offset, + "safetensors_header_sha256": digest, + } + for key, value in header.items(): + if key != "__metadata__": + if key in self.headers or self.weight_map.get(key) != name: + raise ValueError("checkpoint index/header disagreement") + self.headers[key] = value + if set(self.headers) != set(self.weight_map): + raise ValueError("checkpoint index/header tensor set drift") + self.specs = weight_specs() + used = set() + for spec in self.specs: + for key, shape in zip(spec.sources, spec.source_shapes, strict=False): + actual = self.headers.get(key) + if actual is None or tuple(actual["shape"]) != shape or actual["dtype"] != "BF16": + raise ValueError(f"checkpoint tensor shape/dtype drift: {key}") + used.add(key) + if used != set(self.headers): + raise ValueError(f"unmapped checkpoint tensors: {sorted(set(self.headers) - used)}") + + def inventory(self, *, hash_weights=False): + from dataclasses import asdict + + shards = [dict(v) for v in self.shards.values()] + if hash_weights: + for shard in shards: + shard["sha256"] = sha256(Path(shard["path"])) + small = {} + for name in ( + "config.json", + "model.safetensors.index.json", + "tokenizer.json", + "tokenizer_config.json", + "speaker_embeddings/english_sample.pt", + "codec_native/config.json", + ): + p = self.root / name + small[name] = {"path": str(p), "bytes": p.stat().st_size, "sha256": sha256(p)} + p = self.root / "codec_native/model.safetensors" + (h, offset, digest) = read_header(p) + codec = { + "path": str(p), + "bytes": p.stat().st_size, + "tensor_count": len(h) - int("__metadata__" in h), + "safetensors_header_bytes": offset, + "safetensors_header_sha256": digest, + } + if hash_weights: + codec["sha256"] = sha256(p) + return { + "schema": "native-three-mk.assets.v2", + "status": "PASS_CPU_ASSET_INVENTORY", + "root": str(self.root), + "tensor_count": len(self.headers), + "checkpoint_dtype": "BF16", + "effective_input_dtype": "BF16", + "backbone_abi_dtype": "BF16 with FP32 A/D/dt_bias", + "constructor_lineage": "BF16Max64V1 -> BF16V1 explicitly bypasses older V3 FP16 " + "constructor; " + "inherited forward methods run on BF16 tensors", + "fp16_weight_roundtrip": False, + "payload_hashes_measured": bool(hash_weights), + "assets": small, + "model_shards": shards, + "codec_weights": codec, + "weight_mapping": [asdict(s) for s in self.specs], + "scope": "Local headers/configuration verified; no GPU execution or model-output " + "qualification.", + } + + def read_tensor(self, key, *, device="cpu"): + from safetensors import safe_open + + with safe_open(str(self.root / self.weight_map[key]), framework="pt", device="cpu") as f: + return f.get_tensor(key).to(device=device) + + def inputs(self, *, device="cpu"): + import torch + + keys = ["text_embedding.weight", "phoneme_embeddings.0.weight"] + keys += [f"audio_embeddings.{i}.weight" for i in range(16)] + values = { + k: self.read_tensor(k).to(device=device, dtype=torch.bfloat16).contiguous() + for k in keys + } + payload = torch.load( + self.root / "speaker_embeddings/english_sample.pt", + map_location="cpu", + weights_only=True, + ) + speaker = payload.get("speaker_encoding") if isinstance(payload, dict) else payload + if ( + not isinstance(speaker, torch.Tensor) + or tuple(speaker.shape) != (64, 1536) + or speaker.dtype != torch.float32 + ): + raise ValueError("english_sample must be FP32[64,1536]") + values["speaker"] = speaker.to(device=device, dtype=torch.bfloat16).contiguous() + return values + + def materialize_weights(self, *, device="cpu"): + import torch + + result = {} + for spec in self.specs: + if spec.family in {"input", "unused"}: + continue + values = [self.read_tensor(k) for k in spec.sources] + if spec.transform.startswith("stack_"): + value = torch.stack(values) + elif spec.transform.startswith("concat_"): + value = torch.cat(values, dim=0) + else: + value = values[0] + if spec.transform == "negative_exp_fp32": + value = value.float().exp().neg() + value = ( + value.reshape(spec.shape) + .to(dtype=getattr(torch, spec.dtype), device=device) + .contiguous() + ) + result[spec.family, spec.field, spec.index] = value + return result + + +class PredictorWeights(C.Structure): + _fields_ = ( + [(x, C.c_void_p * 3) for x in ("qkv", "o", "ff_in", "ff_out", "norm_self", "norm_ff")] + + [("pos_emb", C.c_void_p)] + + [(x, C.c_void_p * 16) for x in ("audio_emb", "head_w", "head_b")] + ) + + +class BackboneWeights(C.Structure): + _fields_ = ( + [("norm", C.c_void_p * 31), ("norm_f", C.c_void_p)] + + [ + (x, C.c_void_p * 14) + for x in ( + "mb_in_proj", + "mb_out_proj", + "mb_conv_w", + "mb_conv_b", + "mb_A", + "mb_D", + "mb_dt_bias", + "mb_gnorm", + ) + ] + + [(x, C.c_void_p * 3) for x in ("at_qkv", "at_o")] + + [ + (x, C.c_void_p * 14) + for x in ( + "moe_exp_up", + "moe_exp_down", + "moe_sh_up", + "moe_sh_down", + "moe_gate", + "moe_gate_bias", + ) + ] + + [("phoneme_w", C.c_void_p), ("phoneme_b", C.c_void_p), ("pattern", C.c_char_p)] + ) + + +class Stage0Bindings(C.Structure): + _fields_ = [ + ("conv_state", C.c_void_p * 14), + ("conv_slot_stride", C.c_int * 14), + ("ssm_state", C.c_void_p * 14), + ("ssm_slot_stride", C.c_int * 14), + ("slot_map", C.c_void_p), + ("kv_pool_base", C.c_void_p * 3), + ("kv_block_table", C.c_void_p * 3), + ("kv_max_blocks_per_seq", C.c_int), + ("kv_tokens_per_block", C.c_int), + ("seq_len", C.c_void_p), + ("embedding", C.c_void_p), + ("gumbel", C.c_void_p), + ("temperature", C.c_float), + ("moe_w_inject_debug", C.c_void_p), + ("codes_out", C.c_void_p), + ("phoneme_out", C.c_void_p), + ] + + +def weight_bindings(owned_tensors): + objects = {"predictor": PredictorWeights(), "backbone": BackboneWeights()} + objects["backbone"].pattern = PATTERN.encode("ascii") + for spec in weight_specs(): + if spec.family in {"input", "unused"}: + continue + tensor = owned_tensors[spec.family, spec.field, spec.index] + if ( + str(tensor.device) != "cuda:0" + or tuple(tensor.shape) != spec.shape + or (not tensor.is_contiguous()) + ): + raise ValueError(f"invalid device weight binding: {spec.field}") + if str(tensor.dtype) != "torch." + spec.dtype: + raise ValueError(f"invalid weight dtype: {spec.field}") + obj = objects[spec.family] + if spec.index is None: + setattr(obj, spec.field, tensor.data_ptr()) + else: + getattr(obj, spec.field)[spec.index] = tensor.data_ptr() + return (objects["predictor"], objects["backbone"]) diff --git a/python/nemotron_native_backend/execution/native_inputs.py b/python/nemotron_native_backend/execution/native_inputs.py new file mode 100644 index 0000000..84d1589 --- /dev/null +++ b/python/nemotron_native_backend/execution/native_inputs.py @@ -0,0 +1,190 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from .native_assets import HIDDEN, WIDTHS + +MAX_GENERATIONS = 511 +RNG_CAPACITY = 64 + + +@dataclass(frozen=True) +class Request: + text: str + target_ids: tuple[int, ...] + context_ids: tuple[int, ...] + + +def requests_from_texts(assets, texts): + from tokenizers import Tokenizer + + tokenizer = Tokenizer.from_file(str(assets.root / "tokenizer.json")) + context = tuple(tokenizer.encode("[EN]", add_special_tokens=True).ids) + if len(context) != 3: + raise ValueError(f"reviewed context has3tokens; tokenizer returned{context}") + result = [] + for text in texts: + if not isinstance(text, str) or not text.strip(): + raise ValueError("text must be a nonempty already-normalized string") + ids = tuple(tokenizer.encode(text, add_special_tokens=False).ids) + (131073,) + if len(ids) > 512: + raise ValueError("target text exceeds512token state capacity") + result.append(Request(text, ids, context)) + return result + + +def validate_requests(requests): + if len(requests) not in WIDTHS: + raise ValueError("native exact cohort requires B1/B16/B32/B64") + if len({len(r.context_ids) for r in requests}) != 1: + raise ValueError("ragged context lengths are unsupported") + for r in requests: + if ( + len(r.context_ids) != 3 + or not 1 <= len(r.target_ids) <= 512 + or r.target_ids[-1] != 131073 + ): + raise ValueError("invalid native request context/target/EOS contract") + if any(type(x) is not int or not 0 <= x < 131075 for x in (*r.context_ids, *r.target_ids)): + raise ValueError("text token outside the checkpoint vocabulary") + + +class NativeInputs: + def __init__(self, assets, requests, *, device="cpu"): + import torch + + validate_requests(requests) + self.torch = torch + self.assets = assets + self.requests = tuple(requests) + self.batch = len(requests) + self.device = torch.device(device) + self.weights = assets.inputs(device=self.device) + self.offsets = torch.zeros(self.batch, dtype=torch.int64, device=self.device) + self.text_lengths = torch.tensor( + [len(r.target_ids) for r in requests], dtype=torch.int64, device=self.device + ) + self.text_tokens = torch.zeros((self.batch, 512), dtype=torch.int64, device=self.device) + for i, r in enumerate(requests): + self.text_tokens[i, : len(r.target_ids)] = torch.tensor( + r.target_ids, dtype=torch.int64, device=self.device + ) + self.slots = torch.arange(self.batch, dtype=torch.int64, device=self.device) + self.last_codes = torch.zeros((self.batch, 16), dtype=torch.int64, device=self.device) + self.last_phoneme = torch.zeros(self.batch, dtype=torch.int64, device=self.device) + self.phoneme_ended = torch.zeros(self.batch, dtype=torch.bool, device=self.device) + self.audio_valid = torch.zeros(self.batch, dtype=torch.bool, device=self.device) + self.stop = torch.zeros(self.batch, dtype=torch.bool, device=self.device) + self.control_logits = torch.zeros((self.batch, 2), dtype=torch.bfloat16, device=self.device) + self.embedding = torch.empty((self.batch, HIDDEN), dtype=torch.bfloat16, device=self.device) + self.combined_model = torch.empty( + (self.batch, HIDDEN), dtype=torch.bfloat16, device=self.device + ) + prompts = [] + for r in requests: + ids = torch.tensor(r.context_ids, dtype=torch.int64, device=self.device) + prompts.append( + torch.cat( + (self.weights["speaker"], self.weights["text_embedding.weight"][ids]), dim=0 + ) + ) + self.prompt = torch.stack(prompts, dim=1).to(torch.bfloat16).contiguous() + self.prompt_positions = self.prompt.shape[0] + self.prepared = False + self.generation = 0 + + def prepare_decode(self): + torch = self.torch + if self.prepared: + raise RuntimeError("previous decode input has not been committed") + if self.generation >= MAX_GENERATIONS: + raise RuntimeError("configured generation limit reached") + (offsets, lengths) = (self.offsets, self.text_lengths) + safe_offsets = torch.minimum(offsets, torch.clamp_min(lengths - 1, 0)) + text_ids = self.text_tokens[self.slots, safe_offsets] + combined = self.weights["text_embedding.weight"][text_ids] * (offsets < lengths).unsqueeze( + -1 + ).to(torch.bfloat16) + self.audio_valid.copy_(offsets >= 5) + audio_codes = self.last_codes.clone() + audio_codes[offsets == 5] = 1024 + audio = self.weights["audio_embeddings.0.weight"][audio_codes[:, 0]] + for i in range(1, 16): + audio = audio + self.weights[f"audio_embeddings.{i}.weight"][audio_codes[:, i]] + audio = audio / 16 + combined = combined + audio * self.audio_valid.unsqueeze(-1).to(audio.dtype) + phoneme_valid = (offsets >= 3) & ~self.phoneme_ended + phoneme = self.last_phoneme.clone() + phoneme[offsets == 3] = 2048 + value = self.weights["phoneme_embeddings.0.weight"][phoneme] + combined = combined + value * phoneme_valid.unsqueeze(-1).to(value.dtype) + self.phoneme_ended |= (offsets > 3) & phoneme_valid & (phoneme == 2049) + self.combined_model.copy_(combined) + self.embedding.copy_(combined) + self.prepared = True + return self.embedding + + def commit_outputs(self, codes, phonemes): + torch = self.torch + if not self.prepared: + raise RuntimeError("commit requires a prepared generation") + if ( + tuple(codes.shape) != (self.batch, 16) + or codes.dtype != torch.int64 + or codes.device != self.device + ): + raise ValueError("codes must be INT64[B,16] on the owned device") + if ( + tuple(phonemes.shape) != (self.batch,) + or phonemes.dtype != torch.int32 + or phonemes.device != self.device + ): + raise ValueError("phoneme output must be INT32[B] on the owned device") + self.stop.copy_((codes == 1025).any(dim=1) & self.audio_valid) + self.last_codes.copy_(codes) + self.last_phoneme.copy_(phonemes) + self.offsets.add_(1) + self.control_logits.zero_() + self.control_logits[:, 1] = torch.where(self.stop, 30.0, -30.0).to(torch.bfloat16) + self.prepared = False + self.generation += 1 + return self.stop + + +def cache_plan(batch, *, prompt_positions=67, max_generations=MAX_GENERATIONS): + if ( + batch not in WIDTHS + or not 1 <= prompt_positions + or (not 1 <= max_generations <= MAX_GENERATIONS) + ): + raise ValueError("unsupported native cohort/cache capacity") + max_tokens = prompt_positions + max_generations + blocks = (max_tokens + 31) // 32 + conv = batch * 3584 * 3 * 2 + ssm = batch * 64 * 24 * 128 * 4 + kv = batch * 2 * blocks * 4 * 32 * 128 * 2 + return { + "batch": batch, + "prompt_positions": prompt_positions, + "max_generations": max_generations, + "max_sequence_tokens": max_tokens, + "tokens_per_block": 32, + "blocks_per_sequence": blocks, + "conv_shape_per_layer": [batch, 3584, 3], + "conv_dtype": "BF16", + "conv_bytes_per_layer": conv, + "ssm_shape_per_layer": [batch, 64, 24, 128], + "ssm_dtype": "FP32", + "ssm_bytes_per_layer": ssm, + "kv_shape_per_layer": [batch * 2 * blocks, 4, 32, 128], + "kv_dtype": "BF16", + "kv_bytes_per_layer": kv, + "block_table_shape_per_layer": [batch, 2, blocks], + "block_table_dtype": "INT32", + "total_recurrent_payload_bytes": 14 * (conv + ssm) + 3 * kv, + "slot_map": "identity INT32[batch]; no reassignment while cohort is active", + "initial_past_len": prompt_positions, + "past_len_update": "mk2_run_native_stage0 increments device lengths; Python must not " + "increment", + "ownership": "native prefill handle; separate Stage0 handle borrows stable cache views", + } diff --git a/python/nemotron_native_backend/execution/native_stage1.py b/python/nemotron_native_backend/execution/native_stage1.py new file mode 100644 index 0000000..adf285e --- /dev/null +++ b/python/nemotron_native_backend/execution/native_stage1.py @@ -0,0 +1,534 @@ +from __future__ import annotations + +import ctypes as C +import hashlib +import json +import threading +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from .source_build import source_library_receipt + +HERE = Path(__file__).resolve().parent + + +class NativeStage1Error(RuntimeError): + pass + + +@dataclass(frozen=True) +class OpSpec: + prefix: str + state: str + cin: int + cout: int + kernel: int + stride: int = 1 + activation: str | None = None + + +def operation_specs(): + ops = [OpSpec("pre_conv", "pre_conv", 40, 768, 3)] + + def residual(prefix, state, channels, k): + ops.extend( + ( + OpSpec( + prefix + ".input_conv", + state + "_input", + channels, + 2 * channels, + k, + activation=prefix + ".input_conv.activation.alpha", + ), + OpSpec( + prefix + ".skip_conv", + state + "_skip", + 2 * channels, + channels, + k, + activation=prefix + ".output_activation.alpha", + ), + ) + ) + + residual("pre_resblocks.0", "pre_resblock_0", 768, 3) + ops.append( + OpSpec( + "pre_up_sample_layers.0", + "pre_upsample_0", + 768, + 768, + 4, + 2, + "pre_up_sample_layers.0.activation.alpha", + ) + ) + for i in range(6): + residual(f"conv_layers.{i}", f"hidden_resblock_{i}", 768, 3) + cin = 768 + for i, (cout, stride) in enumerate(((384, 9), (128, 7), (32, 7))): + prefix = f"resblock_up_sample_layers.{i}" + ops.append( + OpSpec( + prefix, f"upsample_{i}", cin, cout, 2 * stride, stride, prefix + ".activation.alpha" + ) + ) + residual(f"resblocks.{i}", f"upsample_resblock_{i}", cout, 7) + cin = cout + ops.append(OpSpec("post_conv", "post_conv", 32, 1, 7)) + return tuple(ops) + + +OPS = operation_specs() +STATE_SPECS = tuple((o.state, o.cin, 1 if o.stride > 1 else o.kernel - 1) for o in OPS) + + +def file_sha(path): + h = hashlib.sha256() + with Path(path).open("rb") as f: + for data in iter(lambda: f.read(1024 * 1024), b""): + h.update(data) + return h.hexdigest() + + +class WeightPointers(C.Structure): + _fields_ = [("weight", C.c_void_p), ("bias", C.c_void_p), ("alpha", C.c_void_p)] + + +class NativeCodecWeights: + def __init__(self, checkpoint, device="cuda:0", expected_sha256=None): + import torch + from safetensors.torch import load_file + + self.path = Path(checkpoint).resolve() + self.device = torch.device(device) + self.sha256 = file_sha(self.path) + if expected_sha256 and self.sha256 != expected_sha256: + raise NativeStage1Error("checkpoint hash mismatch") + weights = load_file(str(self.path), device="cpu") + codec = { + k[len("audio_decoder.") :]: v + for (k, v) in weights.items() + if k.startswith("audio_decoder.") + } + if len(codec) != 76: + raise NativeStage1Error(f"expected76 codec tensors, got{len(codec)}") + self.tensors = [] + self.pointers = (WeightPointers * 26)() + for p, o in enumerate(OPS): + w = codec[o.prefix + ".conv.weight"] + b = codec[o.prefix + ".conv.bias"] + shape = (o.cin, 1, o.kernel) if o.stride > 1 else (o.cout, o.cin, o.kernel) + if tuple(w.shape) != shape or tuple(b.shape) != (o.cout,): + raise NativeStage1Error(f"weight shape mismatch at{o.prefix}") + w = w.to(device=self.device, dtype=torch.float16).contiguous() + b = b.to(device=self.device, dtype=torch.float32).contiguous() + alpha = None + if o.activation: + alpha = codec[o.activation].reshape(-1) + if alpha.numel() != o.cout // 2: + raise NativeStage1Error("HalfSnake alpha shape mismatch") + alpha = alpha.to(device=self.device, dtype=torch.float32).contiguous() + self.tensors.extend(x for x in (w, b, alpha) if x is not None) + self.pointers[p] = WeightPointers( + w.data_ptr(), b.data_ptr(), alpha.data_ptr() if alpha is not None else None + ) + torch.cuda.current_stream(self.device).synchronize() + + def receipt(self): + return { + "checkpoint": str(self.path), + "checkpoint_sha256": self.sha256, + "tensor_count": len(self.tensors), + "bytes": sum(x.numel() * x.element_size() for x in self.tensors), + "weight_dtype": "float16", + "bias_alpha_dtype": "float32", + "immutable": True, + } + + +def load_library(path): + lib = C.CDLL(str(Path(path).resolve())) + lib.native_stage1_create.argtypes = [C.c_int] * 4 + [C.POINTER(WeightPointers)] + lib.native_stage1_create.restype = C.c_void_p + lib.native_stage1_run.argtypes = ( + [C.c_void_p] * 4 + [C.POINTER(C.c_void_p)] * 2 + [C.c_void_p] * 3 + [C.c_int, C.c_void_p] + ) + lib.native_stage1_run.restype = C.c_int + lib.native_stage1_destroy.argtypes = [C.c_void_p] + lib.native_stage1_destroy.restype = C.c_int + lib.native_stage1_last_error.restype = C.c_char_p + lib.native_stage1_scratch_bytes.argtypes = [C.c_void_p] + lib.native_stage1_scratch_bytes.restype = C.c_size_t + lib.native_stage1_resident_blocks.argtypes = [C.c_void_p] + lib.native_stage1_resident_blocks.restype = C.c_int + return lib + + +@dataclass +class Ticket: + owner: Any + event: Any + codes: Any + lengths: Any + slots: Any + states: tuple + next_states: tuple + pcm: Any + samples: Any + status: Any + finite: Any = None + + +class NativeStage1: + def __init__( + self, + weights: NativeCodecWeights, + batch=1, + max_frames=16, + state_slots=None, + resident_blocks=0, + library=None, + ): + import torch + + self.weights = weights + self.device = weights.device + self.batch = int(batch) + self.max_frames = int(max_frames) + self.state_slots = int(batch if state_slots is None else state_slots) + if library is None: + raise ValueError("An explicit qualified codec library is required") + self.path = Path(library).resolve() + self.lib = load_library(self.path) + self._lock = threading.RLock() + self._ticket = None + self._poisoned = False + self.device_finite = False + with torch.cuda.device(self.device): + self.handle = self.lib.native_stage1_create( + self.batch, + self.max_frames, + self.state_slots, + int(resident_blocks), + weights.pointers, + ) + if not self.handle: + raise NativeStage1Error(self._error()) + + def _error(self): + return (self.lib.native_stage1_last_error() or b"unknown native Stage1 error").decode() + + def zero_states(self): + import torch + + return tuple( + torch.zeros((self.state_slots, c, h), dtype=torch.float16, device=self.device) + for (_, c, h) in STATE_SPECS + ) + + def bind_buffers( + self, + codes, + lengths, + slots, + states, + next_states, + pcm, + samples, + status, + stream=None, + completion_event=None, + ): + return BoundStage1( + self, + codes, + lengths, + slots, + states, + next_states, + pcm, + samples, + status, + stream, + completion_event, + ) + + def run_async(self, codes, lengths, slots, states, stream=None): + import torch + + with self._lock: + if not self.handle or self._poisoned: + raise NativeStage1Error("closed/poisoned handle") + if self._ticket is not None: + raise NativeStage1Error("finish outstanding ticket before reuse") + frames = int(codes.shape[1]) if codes.ndim == 3 else 0 + if not 1 <= frames <= self.max_frames: + raise NativeStage1Error("invalid chunk frames") + tensors = [ + (codes, torch.int64, (self.batch, frames, 16)), + (lengths, torch.int32, (self.batch,)), + (slots, torch.int32, (self.batch,)), + ] + if len(states) != 26: + raise NativeStage1Error("expected26 input states") + tensors.extend( + (x, torch.float16, (self.state_slots, c, h)) + for (x, (_, c, h)) in zip(states, STATE_SPECS, strict=False) + ) + for x, dtype, shape in tensors: + if ( + x.device != self.device + or x.dtype != dtype + or tuple(x.shape) != shape + or (not x.is_contiguous()) + ): + raise NativeStage1Error( + f"input ABI mismatch: expected{dtype}/{shape}/{self.device}" + ) + stream = stream or torch.cuda.current_stream(self.device) + with torch.cuda.device(self.device), torch.cuda.stream(stream): + ns = tuple(torch.empty_like(x) for x in states) + pcm = torch.empty( + (self.batch, frames * 1764), dtype=torch.float32, device=self.device + ) + samples = torch.empty((self.batch,), dtype=torch.int32, device=self.device) + status = torch.empty((4,), dtype=torch.int32, device=self.device) + ptrs = (C.c_void_p * 26)(*(x.data_ptr() for x in states)) + outs = (C.c_void_p * 26)(*(x.data_ptr() for x in ns)) + code = self.lib.native_stage1_run( + self.handle, + codes.data_ptr(), + lengths.data_ptr(), + slots.data_ptr(), + ptrs, + outs, + pcm.data_ptr(), + samples.data_ptr(), + status.data_ptr(), + frames, + int(stream.cuda_stream), + ) + if code: + self._poisoned = True + raise NativeStage1Error(self._error()) + event = torch.cuda.Event(blocking=False) + event.record(stream) + self._ticket = Ticket( + self, event, codes, lengths, slots, tuple(states), ns, pcm, samples, status + ) + return self._ticket + + def finish(self, ticket, copy_pcm=False): + with self._lock: + if ticket is not self._ticket or ticket.owner is not self: + raise NativeStage1Error("ticket ownership/reuse mismatch") + try: + ticket.event.synchronize() + status = ticket.status.cpu().tolist() + except Exception: + self._poisoned = True + self._ticket = None + raise + self._ticket = None + if status[0]: + raise NativeStage1Error( + f"device input validation failed: code,row,frame,detail={status}" + ) + return { + "pcm": ticket.pcm.cpu() if copy_pcm else ticket.pcm, + "next_states": ticket.next_states, + "samples": ticket.samples, + "status": status, + "sample_rate_hz": 22050, + } + + def receipt(self): + source_build = source_library_receipt(self.path, self.batch) + if source_build is not None: + return { + "kind": "native_stage1_single_cooperative_kernel_prototype", + "qualified": False, + "library": str(self.path), + "library_sha256": file_sha(self.path), + "batch": self.batch, + **source_build, + "state_slots": self.state_slots, + "max_frames": self.max_frames, + "scratch_bytes": self.lib.native_stage1_scratch_bytes(self.handle), + "resident_blocks": self.lib.native_stage1_resident_blocks(self.handle), + "weights": self.weights.receipt(), + "tensorrt_dependency": False, + } + built_source = self.path.parent / "native_stage1.cu" + if not built_source.is_file(): + raise NativeStage1Error("library has no adjacent frozen source receipt") + build_path = self.path.parent / "build-receipt.json" + build = json.loads(build_path.read_text()) + library_hash = file_sha(self.path) + if build.get("library_sha256") != library_hash: + raise NativeStage1Error("library does not match its build receipt") + closure = [] + for entry in build["sources"]: + frozen = self.path.parent / Path(entry["path"]).name + actual = file_sha(frozen) + if actual != entry["sha256"]: + raise NativeStage1Error(f"frozen build source hash mismatch: {frozen}") + closure.append({"path": str(frozen), "sha256": actual}) + return { + "kind": "native_stage1_single_cooperative_kernel_prototype", + "qualified": False, + "library": str(self.path), + "library_sha256": library_hash, + "source": str(built_source), + "source_sha256": file_sha(built_source), + "batch": self.batch, + "build_receipt": str(build_path), + "build_receipt_sha256": file_sha(build_path), + "source_closure": closure, + "state_slots": self.state_slots, + "max_frames": self.max_frames, + "scratch_bytes": self.lib.native_stage1_scratch_bytes(self.handle), + "resident_blocks": self.lib.native_stage1_resident_blocks(self.handle), + "weights": self.weights.receipt(), + "tensorrt_dependency": False, + } + + def close(self): + with self._lock: + if self._ticket is not None: + self.finish(self._ticket) + if self.handle: + result = self.lib.native_stage1_destroy(self.handle) + self.handle = None + if result: + raise NativeStage1Error(self._error()) + + +class BoundStage1: + def __init__( + self, + engine, + codes, + lengths, + slots, + states, + next_states, + pcm, + samples, + status, + stream=None, + completion_event=None, + ): + import torch + + self.engine = engine + self.stream = stream or torch.cuda.current_stream(engine.device) + self.event = completion_event or torch.cuda.Event(blocking=False) + self.codes = codes + self.lengths = lengths + self.slots = slots + self.states = tuple(states) + self.next_states = tuple(next_states) + self.pcm = pcm + self.samples = samples + self.status = status + f = int(codes.shape[1]) if codes.ndim == 3 else 0 + if not 1 <= f <= engine.max_frames or len(states) != 26 or len(next_states) != 26: + raise NativeStage1Error("invalid fixed-buffer frame/state count") + specs = [ + (codes, torch.int64, (engine.batch, f, 16)), + (lengths, torch.int32, (engine.batch,)), + (slots, torch.int32, (engine.batch,)), + (pcm, torch.float32, (engine.batch, 1764 * f)), + (samples, torch.int32, (engine.batch,)), + (status, torch.int32, (4,)), + ] + for tensors in (states, next_states): + specs.extend( + (x, torch.float16, (engine.state_slots, c, h)) + for (x, (_, c, h)) in zip(tensors, STATE_SPECS, strict=False) + ) + for x, dtype, shape in specs: + if ( + x.dtype != dtype + or x.device != engine.device + or tuple(x.shape) != shape + or (not x.is_contiguous()) + ): + raise NativeStage1Error(f"fixed-buffer ABI mismatch: expected{dtype}/{shape}") + + def interval(x): + return (x.data_ptr(), x.data_ptr() + x.numel() * x.element_size()) + + read = [interval(x) for x in (codes, lengths, slots, *states)] + write = [interval(x) for x in (pcm, samples, status, *next_states)] + + def overlaps(a, b): + return a[0] < b[1] and b[0] < a[1] + + if any(overlaps(a, b) for a in read for b in write) or any( + overlaps(a, b) for (i, a) in enumerate(write) for b in write[:i] + ): + raise NativeStage1Error("fixed-buffer inputs/outputs must not overlap") + self.inputs = (C.c_void_p * 26)(*(x.data_ptr() for x in states)) + self.outputs = (C.c_void_p * 26)(*(x.data_ptr() for x in next_states)) + self.args = ( + engine.handle, + codes.data_ptr(), + lengths.data_ptr(), + slots.data_ptr(), + self.inputs, + self.outputs, + pcm.data_ptr(), + samples.data_ptr(), + status.data_ptr(), + f, + int(self.stream.cuda_stream), + ) + + def enqueue(self): + import torch + + e = self.engine + with e._lock: + if not e.handle or e._poisoned: + raise NativeStage1Error("closed/poisoned handle") + if e._ticket is not None: + raise NativeStage1Error("finish outstanding ticket before reuse") + with torch.cuda.device(e.device), torch.cuda.stream(self.stream): + if e.lib.native_stage1_run(*self.args): + e._poisoned = True + raise NativeStage1Error(e._error()) + finite = torch.isfinite(self.pcm).all().view(1) if e.device_finite else None + self.event.record(self.stream) + e._ticket = Ticket( + e, + self.event, + self.codes, + self.lengths, + self.slots, + self.states, + self.next_states, + self.pcm, + self.samples, + self.status, + finite, + ) + return e._ticket + + +if __name__ == "__main__": + print( + json.dumps( + { + "operations": len(OPS), + "states": STATE_SPECS, + "elements_per_request": sum(c * h for (_, c, h) in STATE_SPECS), + "status": "UNQUALIFIED_BUILD_READY_PROTOTYPE", + }, + indent=2, + ) + ) diff --git a/python/nemotron_native_backend/execution/pcm_transfer.py b/python/nemotron_native_backend/execution/pcm_transfer.py new file mode 100644 index 0000000..b4e178f --- /dev/null +++ b/python/nemotron_native_backend/execution/pcm_transfer.py @@ -0,0 +1,112 @@ +def install(engine): + import torch + + from .. import r31 + + codec = engine.core.codec + stream = torch.cuda.Stream(device=codec.device) + device_finite = r31.enabled("fast_first_chunk") + codec.device_finite = device_finite + slots = [] + for _ in range(2): + slots.append( + { + "pcm": torch.empty(engine.batch * 8 * 1764, dtype=torch.float32, pin_memory=True), + "metadata": torch.empty(engine.batch + 4, dtype=torch.int32, pin_memory=True), + "finite": torch.empty(1, dtype=torch.bool, pin_memory=True), + "done": torch.cuda.Event(), + "ticket": None, + } + ) + submit = engine._submit_chunk + finish = codec.finish + close = engine.close + cursor = 0 + pending = {} + + def submit_ticket(ticket, *, row_count=None): + nonlocal cursor + with codec._lock: + if codec._ticket is not ticket or ticket.owner is not codec: + raise RuntimeError("pinned ticket ownership mismatch") + if id(ticket) in pending: + raise RuntimeError("pinned ticket already submitted") + if row_count is not None and ( + type(row_count) is not int or not 0 <= row_count <= engine.batch + ): + raise ValueError("PCM transfer row count must fit the physical batch") + slot = slots[cursor] + if slot["ticket"] is not None: + raise RuntimeError("pinned PCM lease reused before release") + slot["ticket"] = ticket + try: + source = ( + ticket.pcm[:row_count] + if row_count is not None and device_finite + else ticket.pcm + ) + pcm = slot["pcm"][: source.numel()].view_as(source) + stream.wait_event(ticket.event) + with torch.cuda.stream(stream): + if source.numel(): + pcm.copy_(source, non_blocking=True) + slot["metadata"][:4].copy_(ticket.status, non_blocking=True) + slot["metadata"][4:].copy_(ticket.samples, non_blocking=True) + if device_finite: + if ticket.finite is None: + raise RuntimeError( + "device finiteness flag missing from the codec ticket" + ) + slot["finite"].copy_(ticket.finite, non_blocking=True) + slot["done"].record(stream) + pending[id(ticket)] = (slot, pcm) + cursor = 1 - cursor + except BaseException: + codec._poisoned = True + raise + + def submit_chunk(*args, **kwargs): + result = submit(*args, **kwargs) + submit_ticket(result["ticket"]) + return result + + def finish_ticket(ticket, copy_pcm=False): + item = pending.get(id(ticket)) + if item is None: + return finish(ticket, copy_pcm=copy_pcm) + (slot, pcm) = item + with codec._lock: + if codec._ticket is not ticket or ticket.owner is not codec: + raise RuntimeError("pinned ticket ownership mismatch") + try: + slot["done"].synchronize() + status = slot["metadata"][:4].tolist() + if status[0]: + raise RuntimeError("codec validation failed: " + str(status)) + result = { + "pcm": pcm if copy_pcm else ticket.pcm, + "next_states": ticket.next_states, + "samples": slot["metadata"][4:], + "status": status, + "sample_rate_hz": 22050, + } + if device_finite and copy_pcm: + result["finite"] = bool(slot["finite"][0]) + return result + except BaseException: + codec._poisoned = True + raise + finally: + codec._ticket = None + slot["ticket"] = None + pending.pop(id(ticket), None) + + def close_engine(): + stream.synchronize() + return close() + + engine._submit_chunk = submit_chunk + engine.submit_pcm_transfer = submit_ticket + codec.finish = finish_ticket + engine.close = close_engine + engine._r8_pcm_owners = (stream, slots, pending) diff --git a/python/nemotron_native_backend/execution/prefix_cache.py b/python/nemotron_native_backend/execution/prefix_cache.py new file mode 100644 index 0000000..1820085 --- /dev/null +++ b/python/nemotron_native_backend/execution/prefix_cache.py @@ -0,0 +1,84 @@ +import hashlib +import time + + +def install(e): + n = e.core.neural + t = e.t + canonical = getattr(n, "_r8_initial_state", None) + if not canonical: + raise RuntimeError("prefix cache requires canonical source") + + class Alias: + def __init__(self, ptr, size): + self.__cuda_array_interface__ = { + "shape": (size,), + "typestr": "|u1", + "data": (int(ptr), False), + "version": 3, + "strides": None, + } + + def view(ptr, size): + return t.as_tensor(Alias(ptr, size), device="cuda:0") + + original = n.prefill_lib.mk2_native_prefill_run + p = n.inputs.prompt_positions + b = n.batch + blocks = n.info.blocks_per_sequence + key = dict( + stage0=n.library_sha256, + prefill=n.prefill_library_sha256, + context=[list(r.context_ids) for r in n.inputs.requests], + prompt_sha256=hashlib.sha256( + n.inputs.prompt.cpu().view(t.uint8).numpy().tobytes() + ).hexdigest(), + precision="BF16", + state="FP32-canonical-private", + batch=b, + ) + version = n.inputs.prompt._version + views = [n.final_hidden, view(n.bindings.slot_map, b * 4), view(n.bindings.seq_len, b * 4)] + for i in range(14): + views.append(view(n.bindings.conv_state[i], b * 3584 * 3 * 2)) + for i in range(3): + views.append(view(n.bindings.kv_block_table[i], b * 2 * blocks * 4)) + pool = view(n.bindings.kv_pool_base[i], b * 2 * blocks * 4 * 32 * 128 * 2).reshape( + b, 2, blocks, 4, 32, 256 + ) + for block in range((p + 31) // 32): + views.append(pool[:, :, block, :, : min(32, p - block * 32), :]) + views.append(view(canonical["ready"].value, 4)) + start = time.monotonic() + with t.cuda.stream(n.stream): + n._check( + original(n.prefill_handle, n.stream.cuda_stream), + "prefix cache preparation", + n.prefill_lib, + ) + snapshots = [v.clone() for v in views] + n.stream.synchronize() + receipt = dict( + key=key, + preparation_s=time.monotonic() - start, + bytes=sum(x.numel() * x.element_size() for x in snapshots), + hits=0, + misses=0, + ) + + def run(handle, stream): + compatible = ( + n.inputs.prompt._version == version + and [list(r.context_ids) for r in n.inputs.requests] == key["context"] + ) + if not compatible: + receipt["misses"] += 1 + return original(handle, stream) + with t.cuda.stream(n.stream): + for dst, src in zip(views, snapshots, strict=False): + dst.copy_(src) + receipt["hits"] += 1 + return 0 + + n.prefill_lib.mk2_native_prefill_run = run + e._r9_prefix_cache = (views, snapshots, original, receipt) diff --git a/python/nemotron_native_backend/execution/prefix_snapshot.py b/python/nemotron_native_backend/execution/prefix_snapshot.py new file mode 100644 index 0000000..ddb5487 --- /dev/null +++ b/python/nemotron_native_backend/execution/prefix_snapshot.py @@ -0,0 +1,246 @@ +import hashlib +import time + +WIDTHS = (1, 2, 4, 8, 16, 64) + + +def require(condition, message): + if not condition: + raise RuntimeError(message) + + +class PrefixSnapshot: + def __init__(self, engine): + self.engine = engine + self.n = n = engine.core.neural + self.t = engine.t + require(engine.batch in WIDTHS and n.batch == engine.batch, "Unsupported FP32 prefix width") + require(not n.started and not n.closed, "Prefix preparation requires a fresh idle owner") + require( + not getattr(n, "_r8_initial_state", None), + "Canonical-state entries need their existing cache", + ) + require(not hasattr(engine, "_fp32_prefix_snapshot"), "Prefix owner is already installed") + require( + n.cache["ssm_dtype"] == "FP32" and n.cache["conv_dtype"] == "BF16", + "Unrecognized state contract", + ) + require( + n.info.ssm_slot_elements == 196608 and n.info.conv_slot_elements == 10752, + "Unexpected recurrent allocation geometry", + ) + require( + all( + n.bindings.ssm_slot_stride[i] == 196608 and n.bindings.conv_slot_stride[i] == 10752 + for i in range(14) + ), + "Recurrent state strides differ from the source contract", + ) + self.original = n.prefill_lib.mk2_native_prefill_run + self.original_close = n.close + self.identity = self.owner_identity() + require( + n.inputs.prompt.is_contiguous() + and n.inputs.prompt.dtype == self.t.bfloat16 + and str(n.inputs.prompt.device) == "cuda:0" + and tuple(n.inputs.prompt.shape) == (n.config.prompt_positions, n.batch, 1536), + "Native prefill requires fixed contiguous BF16 position-major CUDA0 prompt", + ) + self.prompt_allocation = self.prompt_layout() + self.prompt_identity = self.prompt_key() + self.closed = False + self.views = self.state_views() + self.snapshots = [] + started = time.monotonic() + with self.t.cuda.stream(n.stream): + n._check( + self.original(n.prefill_handle, n.stream.cuda_stream), + "FP32 prefix snapshot preparation", + n.prefill_lib, + ) + self.snapshots = [value.clone() for value in self.views] + n.stream.synchronize() + self.receipt = dict( + schema="nemotron.r22.fp32-prefix-snapshot.v1", + status="PREPARED_UNQUALIFIED", + width=engine.batch, + state="FP32-private", + conv="BF16", + kv="BF16-valid-prefix-only", + neural_implementation_changed=False, + request_text_cached=False, + output_cached=False, + context_noise_draw_removed=False, + bytes=sum(value.numel() * value.element_size() for value in self.snapshots), + preparation_seconds=time.monotonic() - started, + prefill_sha256=n.prefill_library_sha256, + stage0_sha256=n.library_sha256, + prompt_sha256=hashlib.sha256( + n.inputs.prompt.cpu().view(self.t.uint8).numpy().tobytes() + ).hexdigest(), + hits=0, + misses=0, + copied_bytes=0, + ) + self.run_hook = self.run + self.close_hook = self.close_neural + n.prefill_lib.mk2_native_prefill_run = self.run_hook + n.close = self.close_hook + engine._fp32_prefix_snapshot = self + + def owner_identity(self): + n = self.n + return ( + n.prefill_handle, + n.stage0_handle, + n.stream.cuda_stream, + n.library_sha256, + n.prefill_library_sha256, + n.final_hidden.data_ptr(), + int(n.bindings.slot_map), + int(n.bindings.seq_len), + tuple( + ( + int(n.bindings.conv_state[i]), + int(n.bindings.ssm_state[i]), + n.bindings.conv_slot_stride[i], + n.bindings.ssm_slot_stride[i], + ) + for i in range(14) + ), + tuple( + (int(n.bindings.kv_block_table[i]), int(n.bindings.kv_pool_base[i])) + for i in range(3) + ), + ) + + def prompt_key(self): + n = self.n + return ( + n.inputs.prompt.data_ptr(), + n.inputs.prompt._version, + tuple(n.inputs.prompt.shape), + str(n.inputs.prompt.dtype), + n.inputs.prompt_positions, + tuple(tuple(request.context_ids) for request in n.inputs.requests), + ) + + def prompt_layout(self): + n = self.n + return ( + n.inputs.prompt.data_ptr(), + tuple(n.inputs.prompt.shape), + tuple(n.inputs.prompt.stride()), + n.inputs.prompt.storage_offset(), + str(n.inputs.prompt.dtype), + str(n.inputs.prompt.device), + n.inputs.prompt_positions, + n.config.batch, + n.config.prompt_positions, + n.config.max_sequence_tokens, + n.config.tokens_per_block, + ) + + def byte_view(self, pointer, size): + require(pointer and size > 0, "Invalid prefix allocation") + + class Alias: + def __init__(self): + self.__cuda_array_interface__ = dict( + shape=(size,), + typestr="|u1", + data=(int(pointer), False), + version=3, + strides=None, + ) + + return self.t.as_tensor(Alias(), device="cuda:0") + + def state_views(self): + n = self.n + b, p, blocks = n.batch, n.inputs.prompt_positions, n.info.blocks_per_sequence + require(1 <= p <= blocks * 32, "Prefix exceeds allocated KV capacity") + expected_kv = b * 2 * blocks * 4 * 32 * 128 * 2 + require(n.info.kv_pool_bytes_per_layer == expected_kv, "Unexpected KV geometry") + views = [ + n.final_hidden, + self.byte_view(n.bindings.slot_map, b * 4), + self.byte_view(n.bindings.seq_len, b * 4), + ] + for i in range(14): + views.extend( + ( + self.byte_view(n.bindings.conv_state[i], b * 10752 * 2), + self.byte_view(n.bindings.ssm_state[i], b * 196608 * 4), + ) + ) + for i in range(3): + views.append(self.byte_view(n.bindings.kv_block_table[i], b * 2 * blocks * 4)) + pool = self.byte_view(n.bindings.kv_pool_base[i], expected_kv).reshape( + b, 2, blocks, 4, 32, 256 + ) + for block in range((p + 31) // 32): + views.append(pool[:, :, block, :, : min(32, p - block * 32), :]) + return views + + def run(self, handle, stream): + n = self.n + require(not self.closed and not n.closed, "Closed prefix snapshot") + require( + handle == n.prefill_handle and stream == n.stream.cuda_stream, + "Prefix call targets a different owner or stream", + ) + require(self.owner_identity() == self.identity, "Prefix allocation ownership changed") + require( + self.prompt_layout() == self.prompt_allocation, + "Native prefill prompt allocation changed; construct a new owner", + ) + if self.prompt_key() != self.prompt_identity: + self.receipt["misses"] += 1 + return self.original(handle, stream) + with self.t.cuda.stream(n.stream): + for destination, snapshot in zip(self.views, self.snapshots, strict=True): + destination.copy_(snapshot) + self.receipt["hits"] += 1 + self.receipt["copied_bytes"] += self.receipt["bytes"] + return 0 + + def detach(self): + if self.closed: + return + n, engine = self.n, self.engine + require( + getattr(engine, "_fp32_prefix_snapshot", None) is self, + "Prefix engine owner was replaced", + ) + require( + n.prefill_lib.mk2_native_prefill_run is self.run_hook, + "Prefix hook was replaced by another owner", + ) + require(n.close is self.close_hook, "Neural close hook was replaced") + n.stream.synchronize() + n.prefill_lib.mk2_native_prefill_run = self.original + n.close = self.original_close + self.snapshots.clear() + self.views.clear() + engine._fp32_prefix_receipt = dict(self.receipt) + del engine._fp32_prefix_snapshot + self.engine = None + self.n = None + self.original = None + self.original_close = None + self.run_hook = None + self.close_hook = None + self.t = None + self.closed = True + + def close_neural(self): + if self.closed: + return + original_close = self.original_close + self.detach() + original_close() + + +def install(engine): + return PrefixSnapshot(engine) diff --git a/python/nemotron_native_backend/execution/program.py b/python/nemotron_native_backend/execution/program.py new file mode 100644 index 0000000..e36144d --- /dev/null +++ b/python/nemotron_native_backend/execution/program.py @@ -0,0 +1,286 @@ +from __future__ import annotations + +import hashlib +import json +import os +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path + +from ..bundle import BundleError, NativeBundlePlan, _freeze +from ..bundle import load_plan as load_wide_plan +from .composition import PackageSupport +from .composition import create_for_qualification as create_wide +from .engine import _CONSTRUCTION_LOCK +from .readiness import load_readiness_plan +from .retained_narrow import create_for_qualification as create_narrow +from .retained_narrow import load_plan as load_narrow_plan +from .small_width import create_small_width_engine, load_small_width_plan + +FAMILIES = { + "readiness": (1,), + "small": (2, 4, 8), + "narrow": (16, 32, 64), + "wide": (128, 256, 512, 1024, 2048, 4096), +} + + +def load_program(path, family): + loaders = { + "readiness": load_readiness_plan, + "small": load_small_width_plan, + "narrow": load_narrow_plan, + "wide": load_wide_plan, + } + if family not in loaders: + raise BundleError("Unknown native program family") + plan = loaders[family](path) + if plan.physical_width not in FAMILIES[family]: + raise BundleError("Neural width does not belong to its declared family") + return plan + + +def load_serving_plan(path, neural, family): + payload = Path(path).read_bytes() + value = json.loads(payload) + if ( + value.get("schema") != "nemotron.unified-serving.r22.v1" + or value.get("enabled") is not False + or value.get("family") != family + or value.get("physical_width") != neural.physical_width + or value.get("neural_manifest_sha256") != neural.manifest_sha256 + or value["model"]["checkpoint_sha256"] != neural.contract["model"]["checkpoint_sha256"] + ): + raise BundleError("Serving and neural programs disagree") + policy = value["host"].get("fp32_prefix_snapshot", "disabled") + batching = value["host"].get("batching_mode", "fixed_cohort") + request_seed = value["host"].get("request_seed", 0) + if type(request_seed) is not int or not 0 <= request_seed <= 0xFFFFFFFFFFFFFFFF: + raise BundleError("Request seed must fit uint64") + if batching not in ("fixed_cohort", "fixed_request", "inflight"): + raise BundleError("Unsupported request batching mode or physical width") + if batching != "fixed_cohort" and "stage0_inflight" not in {a.name for a in neural.artifacts}: + raise BundleError("Request batching requires a source-bound masked producer") + if "inflight_policy" in value["host"]: + from .inflight_policy import InflightPolicy + + try: + InflightPolicy.parse(value["host"]["inflight_policy"]) + except ValueError as error: + raise BundleError(str(error)) from error + if batching == "fixed_cohort": + raise BundleError("In-flight policy cannot alter the fixed-cohort path") + if policy not in ("disabled", "immutable-fp32-v3") or ( + policy != "disabled" and neural.physical_width not in (1, 2, 4, 8, 16, 64) + ): + raise BundleError("Unsupported FP32 prefix policy or physical width") + research = value.get("research_chunk_schedule") + if research is not None: + from .research_program import load_research_program + + research_path = Path(path).with_name("research.json") + program = load_research_program(research_path, neural) + if ( + research + != { + "program_sha256": program.manifest_sha256, + "startup": list(program.schedule), + "steady": 8, + "horizon": 511, + "qualification_only": True, + } + or program.horizon != 511 + ): + raise BundleError("Changed research serving schedule or horizon") + for section in ("generation", "execution", "state"): + expected = neural.contract[section] + if section == "generation" and research is not None: + expected = dict(expected, chunk_frames=tuple(program.schedule) + (8,)) + if section in value and _freeze(value[section]) != _freeze(expected): + raise BundleError("Serving wrapper changes the neural contract: " + section) + return NativeBundlePlan( + value["id"], + neural.physical_width, + neural.artifacts, + _freeze(value), + hashlib.sha256(payload).hexdigest(), + ) + + +@dataclass(frozen=True) +class ProgramFactory: + family: str + neural: NativeBundlePlan + neural_path: Path + artifact_root: Path + model_root: Path + source_pins: Mapping[str, str] + model_identity: str + readiness: object + serving_sha256: str + token_policy: str + prefix_snapshot_policy: str = "disabled" + research_path: Path | None = None + research_sha256: str | None = None + + def verify(self, serving, readiness): + if ( + self.family not in FAMILIES + or self.neural.physical_width not in FAMILIES[self.family] + or serving.physical_width != self.neural.physical_width + or serving.manifest_sha256 != self.serving_sha256 + or serving.contract["neural_manifest_sha256"] != self.neural.manifest_sha256 + or serving.contract["host"]["tokens"] != self.token_policy + or self.token_policy not in ("pinned", "retained-engine-token-packing") + or serving.contract["host"].get("fp32_prefix_snapshot", "disabled") + != self.prefix_snapshot_policy + or self.prefix_snapshot_policy not in ("disabled", "immutable-fp32-v3") + or ( + self.prefix_snapshot_policy != "disabled" + and self.neural.physical_width not in (1, 2, 4, 8, 16, 64) + ) + or readiness is not self.readiness + or self.model_identity != self.neural.contract["model"]["checkpoint_sha256"] + ): + raise BundleError("Invalid serving program binding") + research = serving.contract.get("research_chunk_schedule") + if research is None: + if self.research_path is not None or self.research_sha256 is not None: + raise BundleError("Unexpected research factory") + elif ( + self.research_path is None + or research["program_sha256"] != self.research_sha256 + or research["horizon"] != 511 + or research["qualification_only"] is not True + ): + raise BundleError("Research factory does not match serving manifest") + self._verify_sources() + + def _verify_sources(self): + if hashlib.sha256(self.neural_path.read_bytes()).hexdigest() != self.neural.manifest_sha256: + raise BundleError("Neural program changed after binding") + PackageSupport(self.neural, {}, self.source_pins).verify_sources() + self.readiness.factory.verify_identity(self.readiness.source_sha256) + if self.research_path is not None: + from .research_program import load_research_program + + program = load_research_program(self.research_path, self.neural) + if program.manifest_sha256 != self.research_sha256 or program.horizon != 511: + raise BundleError("Research program changed after binding") + + def __call__(self, width): + if self.prefix_snapshot_policy not in ("disabled", "immutable-fp32-v3"): + raise BundleError("Unsupported FP32 prefix policy") + if width not in (1, self.neural.physical_width): + raise BundleError("Unregistered physical width") + with _CONSTRUCTION_LOCK: + self._verify_sources() + previous = os.environ.get("NEMOTRON_STAGE0_MEGAKERNEL_GRID") + try: + if width == 1 and (self.neural.physical_width != 1 or self.research_path is None): + engine = self.readiness.factory() + elif self.research_path is not None: + from .research_program import create_research, load_research_program + + program = load_research_program(self.research_path, self.neural) + engine = create_research( + program, + self.neural, + self.artifact_root, + self.model_root, + source_pins=self.source_pins, + model_identity=self.model_identity, + ) + else: + constructors = { + "small": create_small_width_engine, + "narrow": create_narrow, + "wide": create_wide, + } + engine = constructors[self.family]( + self.neural, + self.artifact_root, + self.model_root, + source_pins=self.source_pins, + model_identity=self.model_identity, + qualification_only=True, + ) + if engine.batch != width or engine.steps != 511: + engine.close() + raise BundleError("Constructed native engine differs from its program") + if width == 1 and engine.active_lanes_config: + engine.close() + raise BundleError("Readiness cannot inherit contraction") + if width == self.neural.physical_width and self.token_policy == "pinned": + from .token_staging import install + + install(engine, mode="pinned") + if ( + width == self.neural.physical_width + and self.prefix_snapshot_policy == "immutable-fp32-v3" + ): + from .prefix_snapshot import install as install_prefix + + try: + install_prefix(engine) + except BaseException: + engine.close() + raise + engine.retain_codes = False + return engine + finally: + if previous is None: + os.environ.pop("NEMOTRON_STAGE0_MEGAKERNEL_GRID", None) + else: + os.environ["NEMOTRON_STAGE0_MEGAKERNEL_GRID"] = previous + + +async def prepare_for_qualification( + plan, + neural, + neural_path, + artifact_root, + model_root, + *, + source_pins, + model_identity, + readiness, + socket_path, + receipts, + qualification_only, + slab_audit_enabled=False, + trace=None, +): + from ..worker import prepare_for_qualification as prepare_worker + + factory = ProgramFactory( + plan.contract["family"], + neural, + Path(neural_path).resolve(strict=True), + Path(artifact_root).resolve(strict=True), + Path(model_root).resolve(strict=True), + _freeze(dict(source_pins)), + model_identity, + readiness, + plan.manifest_sha256, + plan.contract["host"]["tokens"], + plan.contract["host"].get("fp32_prefix_snapshot", "disabled"), + Path(neural_path).with_name("research.json") + if "research_chunk_schedule" in plan.contract + else None, + plan.contract.get("research_chunk_schedule", {}).get("program_sha256"), + ) + return await prepare_worker( + plan, + artifact_root, + model_root, + source_pins=source_pins, + model_identity=model_identity, + readiness=readiness, + socket_path=socket_path, + receipts=receipts, + qualification_only=qualification_only, + slab_audit_enabled=slab_audit_enabled, + trace=trace, + engine_factory_override=factory, + ) diff --git a/python/nemotron_native_backend/execution/readiness.py b/python/nemotron_native_backend/execution/readiness.py new file mode 100644 index 0000000..4b18eee --- /dev/null +++ b/python/nemotron_native_backend/execution/readiness.py @@ -0,0 +1,265 @@ +from __future__ import annotations + +import ctypes as C +import hashlib +import json +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType + +from ..bundle import ( + Artifact, + BundleError, + NativeBundlePlan, + _freeze, + _relative_path, + _sha256, + resolve_artifacts, +) +from . import helper_graphs +from .composition import PackageSupport +from .engine import _CONSTRUCTION_LOCK, NativeCohortEngine +from .source_build import is_source_build, verify_source_build + +READINESS_POLICY = MappingProxyType( + dict( + mode="serial", + grid_policy="overlap-only", + stage0_grid=96, + full_stage0_grid=96, + codec_grid=48, + ) +) + + +def load_readiness_plan(path: Path) -> NativeBundlePlan: + payload = Path(path).read_bytes() + value = json.loads(payload) + if ( + value.get("schema") != "nemotron.native-readiness-plan.r22.v1" + or value.get("enabled") is not False + or value.get("status") != "PROPOSED_UNQUALIFIED" + or value.get("physical_width") != 1 + or value.get("neural_cohort_partitioning") is not False + ): + raise BundleError("Invalid isolated physical-B1 readiness plan") + if value["execution"] != dict(READINESS_POLICY): + raise BundleError("Readiness must retain its serial96/48 resource contract") + if value["generation"] != dict( + cap=511, startup_skip=True, chunk_frames=[3, 5, 8], contraction=False + ): + raise BundleError("Readiness generation contract changed") + if value["state"] != dict( + format="FP32-private", + recurrence_arithmetic="FP32", + canonical_binding=False, + prefix_cache=False, + initialization="prefill-per-cohort", + ): + raise BundleError("Readiness must not inherit wide private-state or prefix-cache hooks") + if value["hooks"] != [ + "predictor_active_mask", + "helper_graphs", + "controller_queues", + "admission_receipt", + ]: + raise BundleError("Readiness hook order changed") + if value["kernels"] != dict(prefill="prefill", stage0="stage0", codec="codec"): + raise BundleError("Readiness requires three separately bound native libraries") + _sha256(value["model"]["checkpoint_sha256"]) + artifacts = tuple( + Artifact(name, _relative_path(row["path"]), _sha256(row["sha256"])) + for name, row in sorted(value["artifacts"].items()) + ) + if len({item.relative_path for item in artifacts}) != len(artifacts): + raise BundleError("Readiness artifact destinations collide") + required = { + "prefill", + "stage0", + "codec", + "selected_candidate", + "prefill_build_receipt", + "stage0_build_receipt", + "codec_build_receipt", + } + if is_source_build(value): + required = {"prefill", "stage0", "codec", "source_build_receipt"} + if not required <= {item.name for item in artifacts}: + raise BundleError("Readiness is missing an immutable binding or build receipt") + return NativeBundlePlan( + value["id"], 1, artifacts, _freeze(value), hashlib.sha256(payload).hexdigest() + ) + + +def verify_readiness_artifacts(plan, artifacts): + if is_source_build(plan.contract): + verify_source_build(plan, artifacts) + return + selection = json.loads(artifacts["selected_candidate"].read_text()) + if selection["schema"] != "native-serving-candidate.v1" or set(selection["widths"]) != {"1"}: + raise BundleError("Wrong historical readiness selection") + row = selection["widths"]["1"] + if row["startup_predictor_skip"] is not True: + raise BundleError("Readiness startup mask qualification is missing") + for region, legacy in (("prefill", "prefill"), ("stage0", "stage0"), ("codec", "stage1")): + actual = hashlib.sha256(artifacts[region].read_bytes()).hexdigest() + receipt = json.loads(artifacts[region + "_build_receipt"].read_text()) + if actual != row[legacy]["library"]["sha256"] or actual != receipt["library_sha256"]: + raise BundleError( + "Readiness library differs from its pinned selection/build: " + region + ) + qualification = artifacts[region + "_qualification"] + if ( + hashlib.sha256(qualification.read_bytes()).hexdigest() + != row[legacy]["qualification"]["sha256"] + ): + raise BundleError("Readiness qualification source changed: " + region) + + +def bind_readiness_mask(neural): + function = neural.lib.mk2_set_native_predictor_active_mask + function.argtypes = [C.c_void_p, C.c_void_p] + function.restype = C.c_int + neural._check( + function(neural.stage0_handle, neural.inputs.audio_valid.data_ptr()), "startup mask" + ) + neural._startup_skip_mask_owner = neural.inputs.audio_valid + + +def bind_admission_receipt(engine, admission): + original = engine.stream + frozen = dict(admission) + + def stream(*values, **keywords): + iterator = original(*values, **keywords) + try: + for event in iterator: + if event["type"] == "receipt": + event["receipt"]["r21_admission"] = dict(frozen) + yield event + finally: + iterator.close() + + engine.stream = stream + + +class ConfiguredReadinessEngine(NativeCohortEngine): + def __init__( + self, + plan, + artifacts, + artifact_root, + model_root, + source_pins, + admission, + *, + research_schedule=(3, 5), + research_steps=None, + ): + self.native_plan = plan + self.native_artifacts = artifacts + self.readiness_plan = plan + support = PackageSupport(plan, artifacts, source_pins) + super().__init__( + 1, + model_root=model_root, + steps=511 if research_steps is None else research_steps, + retain_codes=False, + support=support, + artifact_root=artifact_root, + qualification_mode=True, + execution=dict(plan.contract["execution"]), + active_lanes={"mode": "disabled"}, + research_schedule=research_schedule, + ) + bind_admission_receipt(self, admission) + + def _allocate_queues(self): + bind_readiness_mask(self.core.neural) + self.startup_predictor_skip = True + self.helper_graphs = helper_graphs.install(self.core.neural) + super()._allocate_queues() + + +@dataclass(frozen=True) +class ReadinessFactory: + plan: NativeBundlePlan + plan_path: Path + artifact_root: Path + model_root: Path + source_pins: Mapping[str, str] + model_identity: str + admission: Mapping[str, object] + source_sha256: str + + def verify_identity(self, expected_source_sha256): + source = Path(__file__).resolve() + actual = hashlib.sha256(source.read_bytes()).hexdigest() + if ( + expected_source_sha256 != self.source_sha256 + or actual != self.source_sha256 + or self.source_pins.get("execution/readiness.py") != actual + ): + raise BundleError("Readiness factory source is not the frozen package source") + if hashlib.sha256(self.plan_path.read_bytes()).hexdigest() != self.plan.manifest_sha256: + raise BundleError("Readiness manifest changed after binding") + if self.model_identity != self.plan.contract["model"]["checkpoint_sha256"]: + raise BundleError("Readiness model identity differs from the wide model") + artifacts = resolve_artifacts(self.plan, self.artifact_root) + verify_readiness_artifacts(self.plan, artifacts) + PackageSupport(self.plan, artifacts, self.source_pins).verify_sources() + return artifacts + + def __call__(self): + with _CONSTRUCTION_LOCK: + artifacts = self.verify_identity(self.source_sha256) + return ConfiguredReadinessEngine( + self.plan, + artifacts, + self.artifact_root, + self.model_root, + self.source_pins, + self.admission, + ) + + +def make_readiness_binding( + path: Path, + artifact_root: Path, + model_root: Path, + *, + source_pins: Mapping[str, str], + wide_plan: NativeBundlePlan, + model_identity: str, + qualification_only: bool, +): + if qualification_only is not True or wide_plan.contract["enabled"] is not False: + raise BundleError("The isolated readiness port is qualification-only") + plan_path = Path(path).resolve(strict=True) + plan = load_readiness_plan(plan_path) + if ( + plan.contract["model"]["checkpoint_sha256"] + != wide_plan.contract["model"]["checkpoint_sha256"] + ): + raise BundleError("Readiness and wide engine use different checkpoints") + source_sha = hashlib.sha256(Path(__file__).read_bytes()).hexdigest() + admission = dict( + initial_singleton_readiness=True, + post_readiness_width=wide_plan.physical_width, + incomplete_deadline_ms=wide_plan.contract["publication"]["incomplete_deadline_ms"], + ) + factory = ReadinessFactory( + plan, + plan_path, + Path(artifact_root).resolve(strict=True), + Path(model_root).resolve(strict=True), + MappingProxyType(dict(source_pins)), + model_identity, + MappingProxyType(admission), + source_sha, + ) + factory.verify_identity(source_sha) + from ..worker import ReadinessBinding + + return ReadinessBinding(factory=factory, source_sha256=source_sha, physical_width=1) diff --git a/python/nemotron_native_backend/execution/request_noise.py b/python/nemotron_native_backend/execution/request_noise.py new file mode 100644 index 0000000..d532d2b --- /dev/null +++ b/python/nemotron_native_backend/execution/request_noise.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +import ctypes as C + + +class RequestNoise: + def __init__(self, neural, progress): + t = neural.torch + self.batch = neural.batch + self.neural = neural + self.progress = progress + self.seeds = t.zeros(self.batch, dtype=t.int64, device=neural.inputs.device) + self.buffer = t.zeros((self.batch, 16, 80), dtype=t.float32, device=neural.inputs.device) + self.generations = t.zeros(self.batch, dtype=t.int32, device=neural.inputs.device) + self.launch = neural.lib.mk2_inflight_noise + self.launch.argtypes = [C.c_void_p] * 5 + self.launch.restype = C.c_int + + def bind(self, slot, seed): + if type(slot) is not int or not 0 <= slot < self.batch: + raise ValueError("invalid request noise slot") + if type(seed) is not int or not 0 <= seed < 1 << 64: + raise ValueError("seed must be an unsigned 64-bit integer") + self.seeds[slot] = seed if seed < 1 << 63 else seed - (1 << 64) + + def generation(self): + n = self.neural + self.generations.copy_(n.inputs.offsets) + n._check( + self.launch( + self.seeds.data_ptr(), + self.generations.data_ptr(), + self.progress.data_ptr(), + self.buffer.data_ptr(), + n.stream.cuda_stream, + ), + "request-local Philox noise", + ) + + @staticmethod + def provenance(): + return { + "schema": "nemotron.request-local-philox.v1", + "key": "request seed; logical generation; codebook and sample index", + "distribution": "float32 Gumbel -log(-log(U))", + "neighbour_independent": True, + "legacy_batch_rng_equivalent": False, + } diff --git a/python/nemotron_native_backend/execution/research_program.py b/python/nemotron_native_backend/execution/research_program.py new file mode 100644 index 0000000..f65fd0f --- /dev/null +++ b/python/nemotron_native_backend/execution/research_program.py @@ -0,0 +1,131 @@ +import hashlib +import json +from dataclasses import dataclass +from pathlib import Path + +from ..bundle import BundleError, _freeze, resolve_artifacts + +SUPPORTED = { + "readiness": (1,), + "small": (2, 4, 8), + "narrow": (16, 32, 64), + "wide": (128, 256, 512, 1024, 2048, 4096), +} +ALLOWED = ((3, 5), (2, 3, 4, 5)) + + +@dataclass(frozen=True) +class ResearchProgram: + identity: str + schedule: tuple[int, ...] + horizon: int + parent_sha256: str + manifest_sha256: str + physical_width: int + family: str + + +def load_research_program(path, base_plan): + payload = Path(path).read_bytes() + v = json.loads(payload) + family = v.get("family") + schedule = tuple(v.get("startup", ())) + if ( + v.get("schema") != "nemotron.chunk-schedule.program.r30.v1" + or family not in SUPPORTED + or v.get("physical_width") not in SUPPORTED[family] + or v.get("physical_width") != base_plan.physical_width + or schedule not in ALLOWED + or v.get("steady") != 8 + or v.get("speech_delay") != 5 + or v.get("horizon") not in (28, 511) + or v.get("parent_generation_cap") != 511 + or v.get("parent_manifest_sha256") != base_plan.manifest_sha256 + or v.get("qualification_only") is not True + or v.get("serving_enabled") is not False + or v.get("continuation_proof") is not True + or v.get("neural_dso_changes") is not False + or _freeze(v.get("parent_execution")) != base_plan.contract["execution"] + or _freeze(v.get("parent_state")) != base_plan.contract["state"] + ): + raise BundleError("Research width, family, schedule or selected parent contract mismatch") + return ResearchProgram( + v["id"], + schedule, + v["horizon"], + base_plan.manifest_sha256, + hashlib.sha256(payload).hexdigest(), + base_plan.physical_width, + family, + ) + + +def create_research(program, base_plan, artifact_root, model_root, *, source_pins, model_identity): + from .composition import ConfiguredEngine, PackageSupport, verify_compiled_contract + from .engine import _CONSTRUCTION_LOCK + + if ( + program.parent_sha256 != base_plan.manifest_sha256 + or program.physical_width != base_plan.physical_width + or model_identity != base_plan.contract["model"]["checkpoint_sha256"] + ): + raise BundleError("Research/base/model identity mismatch") + artifacts = resolve_artifacts(base_plan, artifact_root) + PackageSupport(base_plan, artifacts, source_pins).verify_sources() + paths = (base_plan, artifacts, Path(artifact_root), Path(model_root), source_pins) + options = dict(research_schedule=program.schedule, research_steps=program.horizon) + with _CONSTRUCTION_LOCK: + if program.family == "readiness": + from .readiness import ConfiguredReadinessEngine, verify_readiness_artifacts + + verify_readiness_artifacts(base_plan, artifacts) + engine = ConfiguredReadinessEngine( + *paths, + dict( + initial_singleton_readiness=True, + post_readiness_width=1, + incomplete_deadline_ms=10000, + ), + **options, + ) + elif program.family == "small": + from .small_width import SmallWidthEngine, verify_builds + + verify_builds(base_plan, artifacts) + engine = SmallWidthEngine(*paths, False, **options) + elif program.family == "narrow": + from .retained_narrow import ConfiguredNarrowEngine, verify + + engine = ConfiguredNarrowEngine(*paths, verify(base_plan, artifacts), **options) + elif program.family == "wide": + engine = ConfiguredEngine( + *paths, verify_compiled_contract(base_plan, artifacts), **options + ) + else: + raise BundleError("Unregistered research family") + original = engine.stream + + def stream(*args, **kwargs): + iterator = original(*args, **kwargs) + try: + for event in iterator: + if event["type"] == "receipt": + event["receipt"]["research_program"] = dict( + id=program.identity, + manifest_sha256=program.manifest_sha256, + startup=list(program.schedule), + steady=8, + speech_delay=5, + horizon=program.horizon, + parent_cap=511, + physical_width=program.physical_width, + family=program.family, + serving_qualified=False, + authoritative_HTTP_measurement=False, + ) + yield event + finally: + iterator.close() + + engine.stream = stream + return engine diff --git a/python/nemotron_native_backend/execution/retained_narrow.py b/python/nemotron_native_backend/execution/retained_narrow.py new file mode 100644 index 0000000..0134115 --- /dev/null +++ b/python/nemotron_native_backend/execution/retained_narrow.py @@ -0,0 +1,198 @@ +from __future__ import annotations + +import hashlib +import json +from pathlib import Path + +from ..bundle import ( + Artifact, + BundleError, + NativeBundlePlan, + _freeze, + _relative_path, + _sha256, + resolve_artifacts, +) +from . import helper_graphs, narrow_prefix_cache, narrow_state_binding +from .composition import PackageSupport +from .engine import _CONSTRUCTION_LOCK, NativeCohortEngine +from .readiness import bind_readiness_mask +from .source_build import is_source_build, verify_source_build + + +def load_plan(path): + payload = Path(path).read_bytes() + v = json.loads(payload) + width = v.get("physical_width") + if ( + v.get("schema") != "nemotron.retained-narrow.r22.v1" + or width not in (16, 32, 64) + or v.get("enabled") is not False + or v.get("status") != "PROPOSED_UNQUALIFIED" + or v.get("neural_cohort_partitioning") is not False + ): + raise BundleError("Unsupported retained narrow contract") + execution = dict( + mode="overlap", + grid_policy="overlap-only", + stage0_grid=80, + full_stage0_grid=132, + codec_grid=48 if width == 64 else 52, + ) + if v["execution"] != execution: + raise BundleError("Retained resource policy changed") + if v["generation"] != dict(cap=511, startup_skip=True, chunk_frames=[3, 5, 8]): + raise BundleError("Retained generation contract changed") + if v["state"] != dict( + format="BF16-private" if width == 32 else "FP32-private", + recurrence="FP32", + prefix_cache=width == 32, + ): + raise BundleError("Retained state/cache contract changed") + if v["predictor"] != dict(full_width="FP8-FFN" if width == 32 else "BF16", contracted="BF16"): + raise BundleError("Retained predictor precision changed") + if width == 16 and v["active_lanes"] != {"mode": "disabled"}: + raise BundleError("B16 must not contract") + if width != 16: + active = v["active_lanes"] + if any( + active.get(k) != x + for k, x in dict( + mode="fixed-subset", + target_batch=16, + after_chunk0=True, + narrow_grid=80, + grid_policy="follow-stage0", + ).items() + ): + raise BundleError("Retained contraction policy changed") + artifacts = tuple( + Artifact(name, _relative_path(row["path"]), _sha256(row["sha256"])) + for name, row in sorted(v["artifacts"].items()) + ) + if len({a.relative_path for a in artifacts}) != len(artifacts): + raise BundleError("Artifact destinations collide") + required = { + "prefill", + "stage0", + "codec", + "selected_candidate", + "compiled_binding", + "retained_runtime", + } + if is_source_build(v): + required = {"prefill", "stage0", "codec", "source_build_receipt"} + if width == 32: + required.add("contraction_bf16") + if not required <= {a.name for a in artifacts}: + raise BundleError("Missing retained binding evidence") + return NativeBundlePlan( + v["id"], width, artifacts, _freeze(v), hashlib.sha256(payload).hexdigest() + ) + + +def verify(plan, artifacts): + if is_source_build(plan.contract): + receipt = verify_source_build(plan, artifacts) + if not receipt.get("expected_residency"): + raise BundleError("Fresh narrow build requires an explicit resource contract") + return receipt["expected_residency"] + binding = json.loads(artifacts["compiled_binding"].read_text()) + if ( + binding["physical_batch"] != plan.physical_width + or binding["neural_cohort_partitioning"] is not False + ): + raise BundleError("Actual compiled width differs from plan") + selection = json.loads(artifacts["selected_candidate"].read_text())["widths"][ + str(plan.physical_width) + ] + config = ( + json.loads(artifacts["numerical_config"].read_text()) if plan.physical_width == 32 else None + ) + for region in ("prefill", "stage0", "codec"): + legacy = "stage1" if region == "codec" else region + actual = hashlib.sha256(artifacts[region].read_bytes()).hexdigest() + expected = ( + config["components"][legacy]["sha256"] + if config + else selection[legacy]["library"]["sha256"] + ) + receipt = json.loads(artifacts[region + "_build_receipt"].read_text()) + if actual != expected or receipt["library_sha256"] != actual: + raise BundleError("Retained library identity mismatch: " + region) + if config: + expected = config["hooks"][0]["kwargs"]["narrow_sha256"] + if hashlib.sha256(artifacts["contraction_bf16"].read_bytes()).hexdigest() != expected: + raise BundleError("B32 contraction precision artifact mismatch") + return json.loads(artifacts["retained_runtime"].read_text())["native"]["residency"] + + +class ConfiguredNarrowEngine(NativeCohortEngine): + def __init__( + self, + plan, + artifacts, + artifact_root, + model_root, + source_pins, + expected, + *, + research_schedule=(3, 5), + research_steps=None, + ): + self.native_plan = plan + self.native_artifacts = artifacts + active = { + k: dict(v) if hasattr(v, "items") else v + for k, v in plan.contract["active_lanes"].items() + } + super().__init__( + plan.physical_width, + model_root=model_root, + steps=511 if research_steps is None else research_steps, + retain_codes=False, + support=PackageSupport(plan, artifacts, source_pins), + artifact_root=artifact_root, + qualification_mode=True, + execution=dict(plan.contract["execution"]), + active_lanes=active, + research_schedule=research_schedule, + ) + try: + actual = self.residency + for name, value in expected["stage0"].items(): + if name != "device" and actual["stage0"][name] != value: + raise BundleError("Actual producer ABI differs from retained contract: " + name) + if actual["codec_blocks"] != expected["codec_blocks"]: + raise BundleError("Actual codec resource binding changed") + if self.batch == 32: + narrow = artifacts["contraction_bf16"] + narrow_state_binding.install( + self, narrow, hashlib.sha256(narrow.read_bytes()).hexdigest(), artifact_root + ) + narrow_prefix_cache.install(self) + self._narrow_retained_recipe = dict(plan.contract["state"]) + except BaseException: + self.close() + raise + + def _allocate_queues(self): + bind_readiness_mask(self.core.neural) + self.startup_predictor_skip = True + self.helper_graphs = helper_graphs.install(self.core.neural) + super()._allocate_queues() + + +def create_for_qualification( + plan, artifact_root, model_root, *, source_pins, model_identity, qualification_only +): + if qualification_only is not True or plan.contract["enabled"] is not False: + raise BundleError("Narrow port is qualification-only") + if model_identity != plan.contract["model"]["checkpoint_sha256"]: + raise BundleError("Model identity mismatch") + artifacts = resolve_artifacts(plan, artifact_root) + expected = verify(plan, artifacts) + with _CONSTRUCTION_LOCK: + return ConfiguredNarrowEngine( + plan, artifacts, Path(artifact_root), Path(model_root), source_pins, expected + ) diff --git a/python/nemotron_native_backend/execution/reusable_inputs.py b/python/nemotron_native_backend/execution/reusable_inputs.py new file mode 100644 index 0000000..bea1331 --- /dev/null +++ b/python/nemotron_native_backend/execution/reusable_inputs.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +from .native_inputs import MAX_GENERATIONS, NativeInputs + + +class ReusableInputs(NativeInputs): + @classmethod + def adopt(cls, existing): + obj = object.__new__(cls) + obj.__dict__ = existing.__dict__.copy() + obj._allocate_scratch() + return obj + + def _allocate_scratch(self): + torch = self.torch + B = self.batch + + def empty(shape, dtype): + return torch.empty(shape, dtype=dtype, device=self.device) + + self.max_text_offsets = (self.text_lengths - 1).clamp_min_(0) + self.safe_offsets = empty((B,), torch.int64) + self.text_ids = empty((B, 1), torch.int64) + self.text_mask = empty((B,), torch.bool) + self.bos_mask = empty((B,), torch.bool) + self.audio_codes = empty((B, 16), torch.int64) + self.audio_indices = empty((B, 16), torch.int64) + self.audio_offsets = ( + torch.arange(16, dtype=torch.int64, device=self.device).reshape(1, 16) * 1032 + ) + self.audio_bos = torch.tensor(1024, dtype=torch.int64, device=self.device) + self.audio_table = torch.cat( + [self.weights[f"audio_embeddings.{i}.weight"] for i in range(16)], dim=0 + ).contiguous() + self.audio_values = empty((B * 16, 1536), torch.bfloat16) + self.audio_frame = empty((B, 1536), torch.bfloat16) + self.phoneme_valid = empty((B,), torch.bool) + self.not_ended = empty((B,), torch.bool) + self.phoneme_ids = empty((B,), torch.int64) + self.phoneme_bos = torch.tensor(2048, dtype=torch.int64, device=self.device) + self.phoneme_value = empty((B, 1536), torch.bfloat16) + self.feed_eos = empty((B,), torch.bool) + self.after_phoneme_delay = empty((B,), torch.bool) + self.code_eos = empty((B, 16), torch.bool) + self.any_eos = empty((B,), torch.bool) + self.positive_logit = torch.tensor(30.0, dtype=torch.bfloat16, device=self.device) + self.negative_logit = torch.tensor(-30.0, dtype=torch.bfloat16, device=self.device) + self._audio_values3 = self.audio_values.view(B, 16, 1536) + self._audio_values_rows = tuple(self._audio_values3[:, i, :] for i in range(16)) + self._audio_index_flat = self.audio_indices.view(-1) + self._text_ids_flat = self.text_ids.view(-1) + self._safe_offsets_column = self.safe_offsets.view(B, 1) + self._text_mask_column = self.text_mask.view(B, 1) + self._bos_mask_column = self.bos_mask.view(B, 1) + self._audio_valid_column = self.audio_valid.view(B, 1) + self._phoneme_valid_column = self.phoneme_valid.view(B, 1) + self._stop_logit = self.control_logits[:, 1] + self._continue_logit = self.control_logits[:, 0] + + def prepare_decode(self): + torch = self.torch + if self.prepared or self.generation >= MAX_GENERATIONS: + raise RuntimeError( + "prepare requires uncommitted-free state and remaining generation capacity" + ) + torch.minimum(self.offsets, self.max_text_offsets, out=self.safe_offsets) + torch.gather(self.text_tokens, 1, self._safe_offsets_column, out=self.text_ids) + torch.index_select( + self.weights["text_embedding.weight"], 0, self._text_ids_flat, out=self.combined_model + ) + torch.lt(self.offsets, self.text_lengths, out=self.text_mask) + self.combined_model.mul_(self._text_mask_column) + torch.ge(self.offsets, 5, out=self.audio_valid) + torch.eq(self.offsets, 5, out=self.bos_mask) + torch.where(self._bos_mask_column, self.audio_bos, self.last_codes, out=self.audio_codes) + torch.add(self.audio_codes, self.audio_offsets, out=self.audio_indices) + torch.index_select(self.audio_table, 0, self._audio_index_flat, out=self.audio_values) + self.audio_frame.copy_(self._audio_values_rows[0]) + for value in self._audio_values_rows[1:]: + torch.add(self.audio_frame, value, out=self.audio_frame) + self.audio_frame.div_(16) + self.audio_frame.mul_(self._audio_valid_column) + torch.add(self.combined_model, self.audio_frame, out=self.combined_model) + torch.ge(self.offsets, 3, out=self.phoneme_valid) + torch.logical_not(self.phoneme_ended, out=self.not_ended) + torch.logical_and(self.phoneme_valid, self.not_ended, out=self.phoneme_valid) + torch.eq(self.offsets, 3, out=self.bos_mask) + torch.where(self.bos_mask, self.phoneme_bos, self.last_phoneme, out=self.phoneme_ids) + torch.index_select( + self.weights["phoneme_embeddings.0.weight"], 0, self.phoneme_ids, out=self.phoneme_value + ) + self.phoneme_value.mul_(self._phoneme_valid_column) + torch.add(self.combined_model, self.phoneme_value, out=self.combined_model) + torch.eq(self.phoneme_ids, 2049, out=self.feed_eos) + torch.gt(self.offsets, 3, out=self.after_phoneme_delay) + torch.logical_and(self.feed_eos, self.after_phoneme_delay, out=self.feed_eos) + torch.logical_and(self.feed_eos, self.phoneme_valid, out=self.feed_eos) + torch.logical_or(self.phoneme_ended, self.feed_eos, out=self.phoneme_ended) + self.embedding.copy_(self.combined_model) + self.prepared = True + return self.embedding + + def commit_outputs(self, codes, phonemes): + torch = self.torch + if not self.prepared: + raise RuntimeError("commit requires prepared inputs") + if ( + codes.dtype != torch.int64 + or tuple(codes.shape) != (self.batch, 16) + or codes.device != self.device + ): + raise ValueError("codes ABI mismatch") + if ( + phonemes.dtype != torch.int32 + or tuple(phonemes.shape) != (self.batch,) + or phonemes.device != self.device + ): + raise ValueError("phoneme ABI mismatch") + torch.eq(codes, 1025, out=self.code_eos) + torch.any(self.code_eos, dim=1, out=self.any_eos) + torch.logical_and(self.any_eos, self.audio_valid, out=self.stop) + self.last_codes.copy_(codes) + self.last_phoneme.copy_(phonemes) + self.offsets.add_(1) + self._continue_logit.zero_() + torch.where(self.stop, self.positive_logit, self.negative_logit, out=self._stop_logit) + self.prepared = False + self.generation += 1 + return self.stop diff --git a/python/nemotron_native_backend/execution/reusable_runner.py b/python/nemotron_native_backend/execution/reusable_runner.py new file mode 100644 index 0000000..fa59b68 --- /dev/null +++ b/python/nemotron_native_backend/execution/reusable_runner.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import ctypes as C + +from .reusable_inputs import ReusableInputs +from .split_native_runtime import NativeNeuralRunner + + +class ReusableNeuralRunner(NativeNeuralRunner): + def __init__(self, *, diagnostic_steps=30, **kwargs): + super().__init__(**kwargs) + try: + self.inputs = ReusableInputs.adopt(self.inputs) + self.step_events = [ + ( + self.torch.cuda.Event(enable_timing=True), + self.torch.cuda.Event(enable_timing=True), + ) + for _ in range(diagnostic_steps) + ] + for a, b in self.step_events: + a.record() + b.record() + self.stream.wait_stream(self.torch.cuda.current_stream(0)) + except BaseException: + self.close() + raise + + def step(self): + if ( + self.closed + or not self.started + or self.inputs.generation >= min(self.cache["max_generations"], len(self.step_events)) + ): + raise RuntimeError("native step exceeds prepared diagnostic capacity") + with self.torch.cuda.stream(self.stream): + self.inputs.prepare_decode() + self.noise.generation(self.batch) + (a, b) = self.step_events[self.inputs.generation] + a.record(self.stream) + self._check( + self.lib.mk2_run_native_stage0( + self.stage0_handle, C.byref(self.bindings), self.stream.cuda_stream + ), + "native Stage0", + ) + b.record(self.stream) + self.inputs.commit_outputs(self.codes, self.phonemes) + self.history.append((a, b)) + return (self.codes, self.phonemes, self.inputs.stop) diff --git a/python/nemotron_native_backend/execution/small_width.py b/python/nemotron_native_backend/execution/small_width.py new file mode 100644 index 0000000..fecc3ce --- /dev/null +++ b/python/nemotron_native_backend/execution/small_width.py @@ -0,0 +1,353 @@ +from __future__ import annotations + +import ctypes as C +import hashlib +import json +from collections.abc import Mapping +from pathlib import Path + +from ..bundle import ( + Artifact, + BundleError, + NativeBundlePlan, + _freeze, + _relative_path, + _sha256, + resolve_artifacts, +) +from . import core, helper_graphs +from .engine import _CONSTRUCTION_LOCK, NativeCohortEngine +from .lane_control import update_first_eos +from .native_assets import ModelAssets +from .native_inputs import requests_from_texts +from .source_build import is_source_build, verify_source_build + +WIDTHS = (2, 4, 8) +HOOKS = ("predictor_active_mask", "helper_graphs", "controller_queues", "publication_receipt") +GENERATION = dict(cap=511, startup_skip=True, chunk_frames=[3, 5, 8], contraction=False) +STATE = dict( + format="FP32-private", + recurrence_arithmetic="FP32", + canonical_binding=False, + prefix_cache=False, + initialization="prefill-per-cohort", +) +PUBLICATION = dict( + continuation_proof_required=True, + sample_rate=22050, + neural_output="pcm_f32_le", + http_output="pcm_s16_le", +) + + +def _digest(path): + value = hashlib.sha256() + with Path(path).open("rb") as stream: + for block in iter(lambda: stream.read(1 << 20), b""): + value.update(block) + return value.hexdigest() + + +def load_small_width_plan(path: Path) -> NativeBundlePlan: + payload = Path(path).read_bytes() + value = json.loads(payload) + width = value.get("physical_width") + if ( + value.get("schema") != "nemotron.small-width-plan.r22.v1" + or type(width) is not int + or width not in WIDTHS + or value.get("status") != "BUILT_UNQUALIFIED" + or value.get("enabled") is not False + or value.get("neural_cohort_partitioning") is not False + ): + raise BundleError("Expected a disabled built physical-B2/B4/B8 qualification plan") + if value.get("generation") != GENERATION or value.get("state") != STATE: + raise BundleError("Small-width generation or FP32 state contract changed") + if value.get("publication") != PUBLICATION or tuple(value.get("hooks", ())) != HOOKS: + raise BundleError("Small-width publication dependencies or hook order changed") + execution = value["execution"] + if execution != dict(mode="serial", grid_policy="static", stage0_grid=132, codec_grid=132): + raise BundleError("This first small-width port supports only proposed serial132/132") + if value["resources"] != dict( + occupancy_qualified=False, + admission_checked_at_construction=True, + compiled_arch="sm_90a", + producer_barrier="actual-gridDim.x", + ): + raise BundleError("Resource qualification status or barrier contract changed") + if value["geometry"] != dict( + physical_requests=width, + producer_tile_rows=16, + prefill_tile_rows=64, + sparse_expert_tile_rows=16, + ): + raise BundleError("Physical width and padded arithmetic geometry disagree") + if value["rng"] != dict( + groups=1, + generated_rows_per_draw=64, + consumed_rows=width, + unused_rows=64 - width, + context_draws=1, + all_physical_generations_advance=True, + regrouping_equivalence=False, + ): + raise BundleError("Small-width independent64 RNG contract changed") + if value["kernels"] != dict(prefill="prefill", stage0="stage0", codec="codec"): + raise BundleError("Three explicit neural library bindings are required") + _sha256(value["model"]["checkpoint_sha256"]) + artifacts = tuple( + Artifact(name, _relative_path(row["path"]), _sha256(row["sha256"])) + for name, row in sorted(value["artifacts"].items()) + ) + required = { + "prefill", + "stage0", + "codec", + "source_manifest", + "external_source_manifest", + "prefill_build_receipt", + "stage0_build_receipt", + "codec_build_receipt", + } + if is_source_build(value): + required = {"prefill", "stage0", "codec", "source_build_receipt"} + if not required <= {item.name for item in artifacts}: + raise BundleError("Missing source, build or library provenance") + if len({item.relative_path for item in artifacts}) != len(artifacts): + raise BundleError("Small-width artifact destinations collide") + return NativeBundlePlan( + value["id"], width, artifacts, _freeze(value), hashlib.sha256(payload).hexdigest() + ) + + +def verify_builds(plan, artifacts): + if is_source_build(plan.contract): + verify_source_build(plan, artifacts) + return + manifest = json.loads(artifacts["source_manifest"].read_text()) + if manifest["batch"] != plan.physical_width: + raise BundleError("Source manifest physical width differs from the program") + expected_manifest = _digest(artifacts["source_manifest"]) + expected_external = _digest(artifacts["external_source_manifest"]) + for region in ("prefill", "stage0", "codec"): + build = json.loads(artifacts[region + "_build_receipt"].read_text()) + if ( + build.get("schema") != "nemotron.small-width-build.r22.v1" + or build.get("status") != "COMPILED" + or build.get("compile_returncode") != 0 + or build.get("physical_batch") != plan.physical_width + or build.get("region") != region + or build.get("source_manifest_sha256") != expected_manifest + or build.get("external_source_manifest_sha256") != expected_external + or build.get("library_sha256") != _digest(artifacts[region]) + ): + raise BundleError("Unverified or mismatched small-width build: " + region) + if not build.get("command") or not build.get("image") or not build.get("compiler"): + raise BundleError("Missing compiler invocation provenance: " + region) + expected = { + Path(row["path"]).name: row["sha256"] + for row in manifest["files"] + if row["region"] == region + } + observed = {Path(row["path"]).name: row["sha256"] for row in build["sources"]} + if expected != observed or len(observed) != len(build["sources"]): + raise BundleError("Build source closure differs from prepared sources: " + region) + for name, digest in expected.items(): + source = artifacts.get(region + "_source_" + name) + if source is None or _digest(source) != digest: + raise BundleError( + "Packaged build source missing or changed: " + region + "/" + name + ) + if region == "codec": + if ( + artifacts["codec_build_receipt"].parent != artifacts["codec"].parent + or artifacts["codec_build_receipt"].name != "build-receipt.json" + ): + raise BundleError("Codec requires its adjacent immutable build receipt") + + +class SmallWidthSupport: + ModelAssets = ModelAssets + Cohort = core.Cohort + requests_from_texts = staticmethod(requests_from_texts) + update_first_eos = staticmethod(update_first_eos) + + def __init__(self, plan, artifacts, source_pins): + self.plan = plan + self.artifacts = artifacts + self.source_pins = dict(source_pins) + self.source_sha256 = _digest(core.__file__) + + def verify_sources(self): + package = Path(__file__).resolve().parent.parent + actual = {str(path.relative_to(package)) for path in package.rglob("*.py")} + if not self.source_pins or actual != set(self.source_pins): + raise BundleError("Small-width package source inventory differs from its freeze") + for name, expected in self.source_pins.items(): + path = (package / name).resolve(strict=True) + if not path.is_relative_to(package) or _digest(path) != expected: + raise BundleError("Small-width package source changed: " + name) + + def select(self, batch): + if batch != self.plan.physical_width or batch not in WIDTHS: + raise BundleError("No small-width fallback or physical cohort partitioning is allowed") + paths = { + name: self.artifacts[key] + for name, key in (("prefill", "prefill"), ("stage0", "stage0"), ("stage1", "codec")) + } + return paths, dict( + physical_batch=batch, + variant=self.plan.identity, + neural_cohort_partitioning=False, + bundle_manifest_sha256=self.plan.manifest_sha256, + ) + + +def bind_startup_mask(neural): + function = neural.lib.mk2_set_native_predictor_active_mask + function.argtypes = [C.c_void_p, C.c_void_p] + function.restype = C.c_int + neural._check( + function(neural.stage0_handle, neural.inputs.audio_valid.data_ptr()), "startup mask" + ) + neural._startup_skip_mask_owner = neural.inputs.audio_valid + + +class SmallWidthEngine(NativeCohortEngine): + def __init__( + self, + plan, + artifacts, + artifact_root, + model_root, + source_pins, + retain_codes, + *, + research_schedule=(3, 5), + research_steps=None, + ): + self.native_plan = plan + self.native_artifacts = artifacts + self.small_width_plan = plan + self.small_width_resources = None + support = SmallWidthSupport(plan, artifacts, source_pins) + super().__init__( + plan.physical_width, + model_root=model_root, + steps=511 if research_steps is None else research_steps, + retain_codes=retain_codes, + support=support, + artifact_root=artifact_root, + qualification_mode=True, + execution=dict(plan.contract["execution"]), + active_lanes={"mode": "disabled"}, + research_schedule=research_schedule, + ) + + def _allocate_queues(self): + n = self.core.neural + info = self.residency["stage0"] + expected = dict( + batch=self.batch, + batch_pad=16, + stages=3, + kchunk=256, + direct_input=1, + skip_hidden_snapshot=1, + threads=256, + grid=132, + ) + if any(info[key] != value for key, value in expected.items()): + raise BundleError( + "Compiled producer geometry differs from the small-width source contract" + ) + if info["blocks_per_sm"] < 1 or info["grid"] > info["sms"] * info["blocks_per_sm"]: + raise BundleError("Producer cooperative launch is not resident") + if tuple(n.torch.cuda.get_device_capability(0)) != (9, 0): + raise BundleError("This source variant is compiled for Hopper sm_90a") + if self.core.codec.batch != self.batch or self.core.codec.state_slots != self.batch: + raise BundleError("Codec width or state-slot allocation differs from physical requests") + if self.residency["codec_blocks"] != 132: + raise BundleError("Codec resource request was clamped") + if any( + n.bindings.ssm_slot_stride[i] != 196608 or n.bindings.conv_slot_stride[i] != 10752 + for i in range(14) + ): + raise BundleError("FP32 recurrence or convolution state stride changed") + if ( + n.noise.batch != self.batch + or len(n.noise.generators) != 1 + or tuple(n.noise.buffer.shape) != (64, 16, 80) + ): + raise BundleError( + "Small physical cohorts must retain one complete independent64 RNG group" + ) + bind_startup_mask(n) + self.startup_predictor_skip = True + self.helper_graphs = helper_graphs.install(n) + super()._allocate_queues() + self.small_width_resources = dict( + admission_checked=True, + execution_qualified=False, + stage0=dict(info), + codec_blocks=self.residency["codec_blocks"], + prefill_batch=n.info.batch, + prefill_padded_rows=n.info.padded_prompt_rows, + prefill_scratch_bytes=n.info.prefill_scratch_bytes, + state_format="FP32-private", + producer_barrier_participants=info["grid"], + precision_verified_by="pinned-source-and-runtime-strides", + ) + + def stream(self, *values, **keywords): + iterator = super().stream(*values, **keywords) + try: + for event in iterator: + if event["type"] == "receipt": + event["receipt"]["small_width"] = dict( + plan_sha256=self.small_width_plan.manifest_sha256, + physical_batch=self.batch, + neural_cohort_partitioning=False, + resources=self.small_width_resources, + hooks=list(HOOKS), + generated_rng_rows=64, + consumed_rng_rows=self.batch, + unused_rng_rows=64 - self.batch, + context_draws=1, + regrouping_equivalence=False, + publication_boundary="engine float PCM; not HTTP TTFA", + qualification_claim=False, + ) + yield event + finally: + iterator.close() + + +def create_small_width_engine( + plan: NativeBundlePlan, + artifact_root: Path, + model_root: Path, + *, + source_pins: Mapping[str, str], + model_identity: str, + qualification_only: bool, + retain_codes: bool = False, +): + if ( + qualification_only is not True + or plan.contract.get("enabled") is not False + or plan.contract.get("schema") != "nemotron.small-width-plan.r22.v1" + or plan.physical_width not in WIDTHS + or type(retain_codes) is not bool + ): + raise BundleError("Only explicitly bound small-width qualification is enabled") + if model_identity != plan.contract["model"]["checkpoint_sha256"]: + raise BundleError("Caller model identity differs from the frozen plan") + artifacts = resolve_artifacts(plan, artifact_root) + verify_builds(plan, artifacts) + support = SmallWidthSupport(plan, artifacts, source_pins) + support.verify_sources() + with _CONSTRUCTION_LOCK: + return SmallWidthEngine( + plan, artifacts, artifact_root, model_root, source_pins, retain_codes + ) diff --git a/python/nemotron_native_backend/execution/source_build.py b/python/nemotron_native_backend/execution/source_build.py new file mode 100644 index 0000000..2d3b994 --- /dev/null +++ b/python/nemotron_native_backend/execution/source_build.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +import hashlib +import json +from pathlib import Path + +from ..bundle import BundleError, _relative_path, _sha256 + +SCHEMA = "nemotron.native-source-build.v1" +LIBRARY_SCHEMA = "nemotron.native-source-library.v1" + + +def is_source_build(contract): + return type(contract.get("source_build")) is int and contract["source_build"] == 1 + + +def _digest(path): + value = hashlib.sha256() + with Path(path).open("rb") as stream: + for data in iter(lambda: stream.read(1 << 20), b""): + value.update(data) + return value.hexdigest() + + +def _resolve(root, value): + path = (root / _relative_path(value)).resolve(strict=True) + if not path.is_relative_to(root) or not path.is_file(): + raise BundleError("Source-build artifact escapes its package") + return path + + +def _verify_receipt(path, width): + path = Path(path).resolve(strict=True) + root = path.parent + receipt = json.loads(path.read_text()) + if ( + receipt.get("schema") != SCHEMA + or receipt.get("status") != "COMPILED" + or type(receipt.get("physical_width")) is not int + or receipt["physical_width"] != width + ): + raise BundleError("Invalid unqualified source-build receipt") + sources = receipt.get("sources") + if not isinstance(sources, list) or not sources: + raise BundleError("Source build requires its complete source closure") + seen = set() + for row in sources: + relative = _relative_path(row["path"]) + if not relative.startswith("sources/") or relative in seen: + raise BundleError("Source closure contains duplicate or non-source entries") + seen.add(relative) + source = _resolve(root, relative) + if _digest(source) != _sha256(row["sha256"]): + raise BundleError("Source-build source changed: " + relative) + libraries = receipt.get("libraries") + if not isinstance(libraries, dict) or not {"prefill", "stage0", "codec"} <= set(libraries): + raise BundleError("Source build requires three neural library bindings") + destinations = set() + for role, row in libraries.items(): + library = _resolve(root, row["path"]) + if library in destinations: + raise BundleError("Source-build library destinations collide") + destinations.add(library) + digest = _digest(library) + if digest != _sha256(row["sha256"]): + raise BundleError("Source-build library changed: " + role) + build_path = _resolve(root, row["receipt"]) + if _digest(build_path) != _sha256(row["receipt_sha256"]): + raise BundleError("Source-build compiler receipt changed: " + role) + build = json.loads(build_path.read_text()) + if ( + build.get("schema") != LIBRARY_SCHEMA + or build.get("status") != "COMPILED" + or build.get("compile_returncode") != 0 + or build.get("physical_width") != width + or build.get("region") != role + or build.get("library_sha256") != digest + or not build.get("command") + or not build.get("compiler") + ): + raise BundleError("Invalid source-build compiler receipt: " + role) + pointer = (build_path.parent / build["source_build_receipt"]).resolve(strict=True) + if pointer != path: + raise BundleError("Compiler receipt refers to another source package") + return receipt + + +def verify_source_build(plan, artifacts): + if ( + not is_source_build(plan.contract) + or plan.contract.get("enabled") is not False + or plan.contract.get("status") not in {"PROPOSED_UNQUALIFIED", "BUILT_UNQUALIFIED"} + ): + raise BundleError("Fresh source builds are qualification-only") + source_receipt = artifacts.get("source_build_receipt") + if source_receipt is None: + raise BundleError("Missing source-build receipt") + source_receipt = Path(source_receipt).resolve(strict=True) + receipt = _verify_receipt(source_receipt, plan.physical_width) + root = source_receipt.parent + for role, row in receipt["libraries"].items(): + if role not in artifacts or Path(artifacts[role]).resolve() != _resolve(root, row["path"]): + raise BundleError("Plan and source-build library binding disagree: " + role) + shared_objects = {name for name, path in artifacts.items() if Path(path).suffix == ".so"} + if shared_objects != set(receipt["libraries"]): + raise BundleError("Source build omits a plan library") + active = plan.contract.get("active_lanes", {}) + if active.get("mode") == "fixed-subset": + known = { + (_resolve(root, row["path"]), row["sha256"]) for row in receipt["libraries"].values() + } + for role in ("stage0", "transfer"): + row = active[role] + if (_resolve(root, row["path"]), row["sha256"]) not in known: + raise BundleError("Contraction library is outside the source build") + expected = receipt.get("expected_residency") + if expected is not None: + producer = expected.get("stage0", {}) + execution = plan.contract["execution"] + if ( + producer.get("batch") != plan.physical_width + or producer.get("grid") != execution.get("full_stage0_grid", execution["stage0_grid"]) + or expected.get("codec_blocks") != execution["codec_grid"] + ): + raise BundleError("Source-build resources differ from the plan") + return receipt + + +def source_library_receipt(library, width): + library = Path(library).resolve(strict=True) + build_path = library.parent / "build-receipt.json" + build = json.loads(build_path.read_text()) + if build.get("schema") != LIBRARY_SCHEMA: + return None + pointer = build.get("source_build_receipt") + if not isinstance(pointer, str) or Path(pointer).is_absolute(): + raise BundleError("Invalid source-build receipt pointer") + source_path = (build_path.parent / pointer).resolve(strict=True) + if not library.is_relative_to(source_path.parent): + raise BundleError("Source-build receipt is outside the library package") + receipt = _verify_receipt(source_path, width) + codec = receipt["libraries"]["codec"] + if _resolve(source_path.parent, codec["path"]) != library: + raise BundleError("Codec does not match its source-build package") + return dict( + build_receipt=str(build_path), + build_receipt_sha256=_digest(build_path), + source_build_receipt=str(source_path), + source_build_receipt_sha256=_digest(source_path), + source_closure=[ + {"path": str(_resolve(source_path.parent, row["path"])), "sha256": row["sha256"]} + for row in receipt["sources"] + ], + ) diff --git a/python/nemotron_native_backend/execution/split_native_runtime.py b/python/nemotron_native_backend/execution/split_native_runtime.py new file mode 100644 index 0000000..4cba7bb --- /dev/null +++ b/python/nemotron_native_backend/execution/split_native_runtime.py @@ -0,0 +1,316 @@ +from __future__ import annotations + +import ctypes as C +from pathlib import Path + +from .native_assets import ( + BackboneWeights, + PredictorWeights, + Stage0Bindings, + sha256, + weight_bindings, +) +from .native_inputs import NativeInputs, cache_plan +from .wide_noise import CohortNoise + + +class PrefillConfig(C.Structure): + _fields_ = [ + (x, C.c_int) + for x in ("batch", "prompt_positions", "max_sequence_tokens", "tokens_per_block") + ] + + +class PrefillInfo(C.Structure): + _fields_ = [ + (x, C.c_int) + for x in ( + "version", + "batch", + "prompt_positions", + "total_prompt_rows", + "padded_prompt_rows", + "max_sequence_tokens", + "tokens_per_block", + "blocks_per_sequence", + "conv_slot_elements", + "ssm_slot_elements", + ) + ] + [("kv_pool_bytes_per_layer", C.c_uint64), ("prefill_scratch_bytes", C.c_uint64)] + + +def bind_prefill_library(path): + lib = C.CDLL(str(Path(path).resolve())) + lib.mk2_last_error.argtypes = [] + lib.mk2_last_error.restype = C.c_char_p + lib.mk2_native_prefill_create.argtypes = [ + C.POINTER(PredictorWeights), + C.POINTER(BackboneWeights), + C.POINTER(PrefillConfig), + C.c_void_p, + C.c_void_p, + ] + lib.mk2_native_prefill_create.restype = C.c_void_p + lib.mk2_native_prefill_run.argtypes = [C.c_void_p, C.c_void_p] + lib.mk2_native_prefill_run.restype = C.c_int + lib.mk2_native_prefill_get_state.argtypes = [ + C.c_void_p, + C.POINTER(Stage0Bindings), + C.POINTER(PrefillInfo), + ] + lib.mk2_native_prefill_get_state.restype = C.c_int + lib.mk2_native_prefill_destroy.argtypes = [C.c_void_p] + lib.mk2_native_prefill_destroy.restype = None + return lib + + +def bind_stage0_library(path): + lib = C.CDLL(str(Path(path).resolve())) + lib.mk2_last_error.argtypes = [] + lib.mk2_last_error.restype = C.c_char_p + lib.mk2_create_stage0.argtypes = [ + C.POINTER(PredictorWeights), + C.POINTER(BackboneWeights), + C.c_int, + ] + lib.mk2_create_stage0.restype = C.c_void_p + lib.mk2_run_native_stage0.argtypes = [C.c_void_p, C.POINTER(Stage0Bindings), C.c_void_p] + lib.mk2_run_native_stage0.restype = C.c_int + lib.mk2_destroy.argtypes = [C.c_void_p] + lib.mk2_destroy.restype = None + lib.mk2_batch.argtypes = [C.c_void_p] + lib.mk2_batch.restype = C.c_int + return lib + + +class NativeNeuralRunner: + def __init__( + self, + *, + assets, + requests, + library, + prefill_library=None, + base_seed=0, + epoch=1, + max_generations=511, + ): + import torch + + self.torch = torch + torch.cuda.set_device(0) + self.library = Path(library).resolve() + self.library_sha256 = sha256(self.library) + self.lib = bind_stage0_library(self.library) + self.prefill_library = Path(prefill_library or library).resolve() + self.prefill_library_sha256 = sha256(self.prefill_library) + self.prefill_lib = bind_prefill_library(self.prefill_library) + self.prefill_handle = None + self.stage0_handle = None + self.stream = None + self.started = False + self.closed = False + self.history = [] + try: + self.inputs = NativeInputs(assets, requests, device="cuda:0") + self.batch = self.inputs.batch + self.cache = cache_plan( + self.batch, + prompt_positions=self.inputs.prompt_positions, + max_generations=max_generations, + ) + self.owned_weights = assets.materialize_weights(device="cuda:0") + (self.pw, self.bw) = weight_bindings(self.owned_weights) + self.final_hidden = torch.empty_like(self.inputs.prompt) + self.noise = CohortNoise( + device="cuda:0", base_seed=base_seed, epoch=epoch, batch=self.batch + ) + self.codes = torch.empty((self.batch, 16), dtype=torch.int64, device="cuda:0") + self.phonemes = torch.empty(self.batch, dtype=torch.int32, device="cuda:0") + self.config = PrefillConfig( + self.batch, self.inputs.prompt_positions, self.cache["max_sequence_tokens"], 32 + ) + self.prefill_handle = self.prefill_lib.mk2_native_prefill_create( + C.byref(self.pw), + C.byref(self.bw), + C.byref(self.config), + self.inputs.prompt.data_ptr(), + self.final_hidden.data_ptr(), + ) + self._require_handle(self.prefill_handle, "prefill create", self.prefill_lib) + self.bindings = Stage0Bindings() + self.info = PrefillInfo() + self._check( + self.prefill_lib.mk2_native_prefill_get_state( + self.prefill_handle, C.byref(self.bindings), C.byref(self.info) + ), + "prefill state", + self.prefill_lib, + ) + if ( + self.info.version, + self.info.batch, + self.info.prompt_positions, + self.info.conv_slot_elements, + self.info.ssm_slot_elements, + self.info.blocks_per_sequence, + self.info.kv_pool_bytes_per_layer, + ) != ( + 1, + self.batch, + self.inputs.prompt_positions, + 10752, + 196608, + self.cache["blocks_per_sequence"], + self.cache["kv_bytes_per_layer"], + ): + raise RuntimeError( + "native prefill allocation geometry disagrees with source contract" + ) + self.stage0_handle = self.lib.mk2_create_stage0( + C.byref(self.pw), C.byref(self.bw), self.batch + ) + self._require_handle(self.stage0_handle, "Stage0 create") + if self.lib.mk2_batch(self.stage0_handle) != self.batch: + raise RuntimeError("native Stage0 library batch mismatch") + self.bindings.embedding = self.inputs.embedding.data_ptr() + self.bindings.gumbel = self.noise.buffer.data_ptr() + self.bindings.temperature = 0.7 + self.bindings.moe_w_inject_debug = None + self.bindings.codes_out = self.codes.data_ptr() + self.bindings.phoneme_out = self.phonemes.data_ptr() + self.stream = torch.cuda.Stream(device=0) + self.stream.wait_stream(torch.cuda.current_stream(0)) + except BaseException: + self.close() + raise + + def _require_handle(self, handle, where, library=None): + if not handle: + raise RuntimeError( + where + + ": " + + ((library or self.lib).mk2_last_error() or b"unknown native error").decode() + ) + + def _check(self, rc, where, library=None): + if rc: + raise RuntimeError( + where + + ": " + + ((library or self.lib).mk2_last_error() or b"unknown native error").decode() + ) + + def start(self): + if self.closed or self.started: + raise RuntimeError("native runner start requires a fresh open instance") + with self.torch.cuda.stream(self.stream): + self.noise.context() + (a, b) = ( + self.torch.cuda.Event(enable_timing=True), + self.torch.cuda.Event(enable_timing=True), + ) + a.record(self.stream) + self._check( + self.prefill_lib.mk2_native_prefill_run( + self.prefill_handle, self.stream.cuda_stream + ), + "native prefill", + self.prefill_lib, + ) + b.record(self.stream) + self.prefill_events = (a, b) + self.started = True + + def step(self): + if ( + self.closed + or not self.started + or self.inputs.generation >= self.cache["max_generations"] + ): + raise RuntimeError("native step requires a started instance with remaining capacity") + with self.torch.cuda.stream(self.stream): + self.inputs.prepare_decode() + self.noise.generation(self.batch) + (a, b) = ( + self.torch.cuda.Event(enable_timing=True), + self.torch.cuda.Event(enable_timing=True), + ) + a.record(self.stream) + self._check( + self.lib.mk2_run_native_stage0( + self.stage0_handle, C.byref(self.bindings), self.stream.cuda_stream + ), + "native Stage0", + ) + b.record(self.stream) + self.inputs.commit_outputs(self.codes, self.phonemes) + self.history.append((a, b)) + return (self.codes, self.phonemes, self.inputs.stop) + + def receipt(self): + self.stream.synchronize() + return { + "schema": "native-three-mk.neural-diagnostic.v4", + "status": "COMPLETED_NATIVE_PREFILL_STAGE0_DIAGNOSTIC", + "batch": self.batch, + "library": str(self.library), + "library_sha256": self.library_sha256, + "components": { + "prefill": { + "library": str(self.prefill_library), + "library_sha256": self.prefill_library_sha256, + "entrypoint": "mk2_native_prefill_run", + "qualification_claim": False, + }, + "stage0": { + "library": str(self.library), + "library_sha256": self.library_sha256, + "entrypoint": "mk2_run_native_stage0", + "qualification_claim": False, + }, + }, + "cache_owner": "prefill handle; Stage0 borrows stable views and advances cache lengths", + "cache": self.cache, + "rng": self.noise.provenance(), + "generation_steps": self.inputs.generation, + "prefill_cuda_event_ms": self.prefill_events[0].elapsed_time(self.prefill_events[1]), + "stage0_cuda_event_ms": [a.elapsed_time(b) for (a, b) in self.history], + "scope": "Neural integration only; no codec/PCM/HTTP or independent arithmetic parity " + "claim. CUDA-event intervals include launch scheduling and first descriptor " + "setup.", + "extra_gpu_work": [ + "BF16 text/audio/phoneme embedding composition", + "full64-row exponential/log/neg noise generation", + "output stop/phoneme/code/offset control updates", + ], + "termination_policy": "fixed slots continue until all lanes EOS or cap; post-EOS " + "per-lane output " + "ignored, computation continues; slot reuse unsupported", + "trt_imported_modules": [ + k + for k in __import__("sys").modules + if k == "tensorrt" or k.startswith(("tensorrt.", "tensorrt_llm")) + ], + } + + def close(self): + if getattr(self, "closed", False): + return + self.closed = True + if self.stream is not None: + self.stream.synchronize() + elif self.prefill_handle or self.stage0_handle: + self.torch.cuda.synchronize(0) + if self.stage0_handle: + self.lib.mk2_destroy(self.stage0_handle) + self.stage0_handle = None + if self.prefill_handle: + self.prefill_lib.mk2_native_prefill_destroy(self.prefill_handle) + self.prefill_handle = None + + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() diff --git a/python/nemotron_native_backend/execution/startup_binding.py b/python/nemotron_native_backend/execution/startup_binding.py new file mode 100644 index 0000000..6983c35 --- /dev/null +++ b/python/nemotron_native_backend/execution/startup_binding.py @@ -0,0 +1,24 @@ +def install(n, *, allow_legacy_full_entry=False): + full = n.lib.mk2_run_native_stage0 + try: + startup = n.lib.mk2_run_native_startup_stage0 + except AttributeError: + if allow_legacy_full_entry: + return + raise + startup.argtypes = full.argtypes + startup.restype = full.restype + n._r8_full_stage0 = full + n._r8_startup_stage0 = startup + n._r8_startup_launches = 0 + n._r8_full_launches = 0 + + def enqueue(*args): + use_startup = getattr(n, "_r8_force_startup", n.inputs.generation < 5) + if use_startup: + n._r8_startup_launches += 1 + return startup(*args) + n._r8_full_launches += 1 + return full(*args) + + n.lib.mk2_run_native_stage0 = enqueue diff --git a/python/nemotron_native_backend/execution/state_binding.py b/python/nemotron_native_backend/execution/state_binding.py new file mode 100644 index 0000000..686bc04 --- /dev/null +++ b/python/nemotron_native_backend/execution/state_binding.py @@ -0,0 +1,64 @@ +import ctypes as C +from pathlib import Path + +import numpy as np + +H = Path(__file__).resolve().parents[1] + + +def install(n, *, state_mode, scales_path, map_path, recipe_identity): + if getattr(n, "_r9_state", None) is not None: + return n._r9_state + import torch + + from .startup_binding import install as startup_install + + mode = state_mode + if n.started or n.batch not in (128, 256, 512, 1024, 2048, 4096): + raise ValueError("fresh supported wide state owner required") + initial = (C.c_void_p * 14)() + ready = C.c_void_p() + get = n.prefill_lib.mk2_r8_initial_state + get.argtypes = [C.c_void_p, C.POINTER(C.c_void_p), C.POINTER(C.c_void_p)] + get.restype = C.c_int + n._check(get(n.prefill_handle, initial, C.byref(ready)), "canonical state source") + bind = n.lib.mk2_r8_bind_initial_state + bind.argtypes = [C.c_void_p, C.POINTER(C.c_void_p), C.c_void_p, C.c_void_p, C.c_int] + bind.restype = C.c_int + n._check( + bind(n.stage0_handle, initial, ready, n.bindings.seq_len, n.inputs.prompt_positions), + "canonical state bind", + ) + recipe = recipe_identity + scale_np = np.fromfile(scales_path, np.float32).reshape(14, 16384) + map_np = np.fromfile(map_path, np.int32).reshape(14, 1536) + scales = torch.tensor(scale_np, device="cuda") + maps = torch.tensor(map_np, device="cuda") + a = (C.c_void_p * 14)(*[scales[i].data_ptr() for i in range(14)]) + b = (C.c_void_p * 14)(*[maps[i].data_ptr() for i in range(14)]) + for lib, fn, handle in [ + (n.lib, "mk2_r9_state_bind", n.stage0_handle), + (n.prefill_lib, "mk2_r9_prefill_bind", n.prefill_handle), + ]: + f = getattr(lib, fn) + f.argtypes = [C.c_void_p, C.POINTER(C.c_void_p), C.POINTER(C.c_void_p)] + f.restype = C.c_int + n._check(f(handle, a, b), fn) + n._r8_initial_state = { + "initial": initial, + "ready": ready, + "bytes": 14 * 196608 * 4, + "prefill_owner": n.prefill_handle, + "initial_positions": n.inputs.prompt_positions, + } + n._r9_state = { + "mode": mode, + "scales": scales, + "maps": maps, + "recipe": str(recipe), + "state_bytes": 2 if mode == 1 else 1, + "canonical": n._r8_initial_state, + "private_bytes": 14 * n.batch * 196608 * (2 if mode == 1 else 1), + } + startup_install(n) + return n._r9_state diff --git a/python/nemotron_native_backend/execution/token_staging.py b/python/nemotron_native_backend/execution/token_staging.py new file mode 100644 index 0000000..40e31ce --- /dev/null +++ b/python/nemotron_native_backend/execution/token_staging.py @@ -0,0 +1,102 @@ +import types + +import numpy as np + + +def validate_lengths(requests, batch, capacity): + if len(requests) != batch: + raise ValueError("token staging requires the complete physical cohort") + lengths = [len(r.target_ids) for r in requests] + if any(length < 1 or length > capacity for length in lengths): + raise ValueError("token length outside prepared capacity") + return lengths + + +def pack_dense(requests, tokens, lengths, max_offsets): + sizes = validate_lengths(requests, tokens.shape[0], tokens.shape[1]) + tokens.fill(0) + lengths[:] = sizes + np.subtract(lengths, 1, out=max_offsets) + for row, (request, size) in enumerate(zip(requests, sizes, strict=False)): + tokens[row, :size] = request.target_ids + return int(sum(sizes)) + + +class TokenStaging: + def __init__(self, engine, mode): + if mode != "pinned": + raise ValueError("unsupported token staging mode") + t = engine.t + n = engine.core.neural + self.mode = mode + (self.batch, self.capacity) = n.inputs.text_tokens.shape + self.lease = t.cuda.Event() + self.lease.record(n.stream) + self.calls = 0 + self.bytes_transferred = 0 + self.last_tokens = 0 + self.scatter_binding = None + self.tokens = t.empty((self.batch, self.capacity), dtype=t.int64, pin_memory=True) + self.lengths = t.empty(self.batch, dtype=t.int64, pin_memory=True) + self.max_offsets = t.empty(self.batch, dtype=t.int64, pin_memory=True) + self.arrays = (self.tokens.numpy(), self.lengths.numpy(), self.max_offsets.numpy()) + self.lease.synchronize() + + def pack(self, requests): + self.lease.synchronize() + self.last_tokens = pack_dense(requests, *self.arrays) + + def enqueue(self, engine): + n = engine.core.neural + n.inputs.text_tokens.copy_(self.tokens, non_blocking=True) + n.inputs.text_lengths.copy_(self.lengths, non_blocking=True) + n.inputs.max_text_offsets.copy_(self.max_offsets, non_blocking=True) + transferred = (self.batch * self.capacity + 2 * self.batch) * 8 + self.lease.record(n.stream) + self.calls += 1 + self.bytes_transferred += transferred + + def receipt(self): + return dict( + mode=self.mode, + batch=self.batch, + capacity=self.capacity, + calls=self.calls, + bytes_transferred=self.bytes_transferred, + last_tokens=self.last_tokens, + scatter_binding=self.scatter_binding, + buffer_ownership="exclusive engine; lease event before CPU buffer reuse", + ) + + +def install(engine, mode="pinned"): + existing = getattr(engine, "_r11_token_staging", None) + if existing is not None: + if existing.mode != mode: + raise RuntimeError("token staging already installed with a different mode") + return existing + staging = TokenStaging(engine, mode) + original = engine._reset + + def reset(self, requests, seed): + (t, c, n) = (self.t, self.core, self.core.neural) + staging.pack(requests) + c.requests = requests + n.inputs.requests = tuple(requests) + if self.contraction is not None: + self.contraction.reset() + c.reset(seed) + with t.cuda.stream(n.stream): + staging.enqueue(self) + self.counts.zero_() + self.consumed.zero_() + self.consumed_host = [0] * self.batch + self._chunk_events = [] + self._overlap_records = [] + self._deferred_serial_grid = False + self._origin.record(n.stream) + + engine._r11_token_staging = staging + engine._r11_token_staging_original_reset = original + engine._reset = types.MethodType(reset, engine) + return staging diff --git a/python/nemotron_native_backend/execution/wide_noise.py b/python/nemotron_native_backend/execution/wide_noise.py new file mode 100644 index 0000000..26b04cf --- /dev/null +++ b/python/nemotron_native_backend/execution/wide_noise.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +MAX_SEED = 9223372036854775807 +GROUP_STRIDE = 521288629 + + +def group_seed(effective_seed, group): + if ( + type(effective_seed) is not int + or not 0 <= effective_seed <= MAX_SEED + or type(group) is not int + or (group < 0) + ): + raise ValueError("invalid effective seed or RNG group") + return effective_seed + group * GROUP_STRIDE & MAX_SEED + + +class CohortNoise: + def __init__(self, *, device="cpu", base_seed=0, epoch=1, batch=64): + import torch + + if ( + batch not in (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096) + or not 0 <= base_seed <= MAX_SEED + or epoch < 1 + ): + raise ValueError("unsupported batch/seed/epoch") + self.torch = torch + self.device = torch.device(device) + self.batch = batch + self.buffer = torch.empty((max(64, batch), 16, 80), dtype=torch.float32, device=self.device) + self.groups = tuple(self.buffer[i : i + 64] for i in range(0, max(64, batch), 64)) + self.generators = tuple(torch.Generator(device=self.device) for _ in self.groups) + self.generator = self.generators[0] + self._ring = None + self._predrawn = None + self._replay = 0 + self.predraw_hits = 0 + self.reset(base_seed + epoch - 1 & MAX_SEED) + + def reset(self, seed): + self.effective_seed = seed + self.draws = 0 + self.context_drawn = False + if self._predrawn is not None and self._predrawn[0] == seed: + self._replay = self._predrawn[1] + self._predrawn = None + self.predraw_hits += 1 + return + self._predrawn = None + self._replay = 0 + for group, generator in enumerate(self.generators): + generator.manual_seed(group_seed(seed, group)) + + def predraw(self, seed, count): + torch = self.torch + if self._ring is None or self._ring.shape[0] < count: + self._ring = torch.empty( + (count, *self.buffer.shape), dtype=self.buffer.dtype, device=self.device + ) + self._predrawn = None + self._replay = 0 + for group, generator in enumerate(self.generators): + generator.manual_seed(group_seed(seed, group)) + self._draw_into_buffer() + for index in range(count): + self._draw_into_buffer() + self._ring[index].copy_(self.buffer) + self._predrawn = (seed, count) + + def _draw_into_buffer(self): + for value, generator in zip(self.groups, self.generators, strict=False): + value.exponential_(generator=generator) + value.log_().neg_() + + def _draw(self): + self._draw_into_buffer() + self.draws += 1 + + def context(self): + if self.context_drawn: + raise RuntimeError("context RNG draw occurs exactly once") + if self._replay: + self.draws += 1 + else: + self._draw() + self.context_drawn = True + + def generation(self, batch): + if not self.context_drawn or batch != self.batch: + raise RuntimeError("RNG cohort width/context mismatch") + index = self.draws - 1 + if index < self._replay: + self.buffer.copy_(self._ring[index]) + self.draws += 1 + else: + self._replay = 0 + self._draw() + return self.buffer[:batch] + + def provenance(self): + result = { + "algorithm": "torch.exponential_(generator).log_().neg_()", + "capacity_shape": list(self.buffer.shape), + "effective_seed": self.effective_seed, + "device": str(self.device), + "draws": self.draws, + "context_discarded_draws": int(self.context_drawn), + "scope": "CPU and CUDA generator streams are not bit-identical. Exact CUDA stream " + "equivalence requires GPU qualification.", + } + if self.batch > 64: + result.update( + schema="native.rng-independent64.v1", + group_rows=64, + group_seeds=[ + group_seed(self.effective_seed, g) for g in range(len(self.generators)) + ], + group_stride=GROUP_STRIDE, + dead_lanes="Every group advances through every physical generation, including " + "finished lanes.", + ) + return result diff --git a/python/nemotron_native_backend/hardware.py b/python/nemotron_native_backend/hardware.py new file mode 100644 index 0000000..9b67375 --- /dev/null +++ b/python/nemotron_native_backend/hardware.py @@ -0,0 +1,61 @@ +import json +import re +from pathlib import Path + +from .bundle import BundleError + + +def validate_native_device(cuda, receipt_dir=None): + report = { + "schema": "nemotron-native-device.v1", + "status": "FAIL", + "device_index": 0, + "required": { + "family": "H100", + "compute_capability": [9, 0], + "sm_count": 132, + "compiled_architecture": "sm_90a", + "mig": False, + }, + "scope": "Qualified device profile; kernel occupancy and allocation checks still apply", + } + try: + if not cuda.is_available(): + raise BundleError("CUDA device unavailable; check the driver and GPU container access") + properties = cuda.get_device_properties(0) + name = str(properties.name) + capability = [int(properties.major), int(properties.minor)] + sms = int(properties.multi_processor_count) + report["detected"] = { + "name": name, + "compute_capability": capability, + "sm_count": sms, + "total_memory_bytes": int(properties.total_memory), + } + if capability != [9, 0]: + raise BundleError( + "The native kernels require Hopper sm_90a; other architectures need a port" + ) + if not re.search(r"\bH100\b", name, re.IGNORECASE) or re.search( + r"\bMIG\b", name, re.IGNORECASE + ): + raise BundleError("This native release is qualified only for a full H100 device") + if sms != 132: + raise BundleError( + "The native launch and synchronization profiles require 132 SMs; " + "a smaller device or MIG partition needs separately qualified profiles" + ) + report["status"] = "PASS" + except Exception as error: + detected = report.get("detected", {}) + report["error"] = ( + f"Unsupported native GPU: {detected or 'not detected'}. {error}. " + "No model weights were loaded and no neural kernel was launched." + ) + if receipt_dir is not None: + path = Path(receipt_dir) / "hardware.json" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(report, indent=2) + "\n") + if report["status"] != "PASS": + raise BundleError(report["error"]) + return report diff --git a/python/nemotron_native_backend/host/__init__.py b/python/nemotron_native_backend/host/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/nemotron_native_backend/host/cpu_pools.py b/python/nemotron_native_backend/host/cpu_pools.py new file mode 100644 index 0000000..78c82e7 --- /dev/null +++ b/python/nemotron_native_backend/host/cpu_pools.py @@ -0,0 +1,49 @@ +import os +import threading +from concurrent.futures import ThreadPoolExecutor + + +def snapshot(backend, loop=None): + import torch + + executors = { + key: value._max_workers + for key, value in vars(backend).items() + if isinstance(value, ThreadPoolExecutor) + } + default = getattr(loop, "_default_executor", None) + return dict( + scope="preparation_only", + affinity=sorted(os.sched_getaffinity(0)), + os_cpu_count=os.cpu_count(), + affinity_cpu_count=len(os.sched_getaffinity(0)), + torch_intraop=torch.get_num_threads(), + torch_interop=torch.get_num_interop_threads(), + backend_executor_max_workers=executors, + asyncio_default_executor_max_workers=getattr(default, "_max_workers", None), + environment={ + key: os.environ.get(key) + for key in ( + "OMP_NUM_THREADS", + "MKL_NUM_THREADS", + "TOKIO_WORKER_THREADS", + "NEMOTRON_R21_TORCH_THREADS", + "NEMOTRON_R21_TORCH_INTEROP_THREADS", + ) + }, + python_threads=[ + dict(name=thread.name, native_id=thread.native_id) for thread in threading.enumerate() + ], + ) + + +def configure(*, intraop_threads, interop_threads): + import torch + + for name, count, setter in [ + ("intraop", intraop_threads, torch.set_num_threads), + ("interop", interop_threads, torch.set_num_interop_threads), + ]: + if type(count) is not int or count < 1: + raise ValueError("Invalid thread-pool size: " + name) + setter(count) diff --git a/python/nemotron_native_backend/host/diagnostic.py b/python/nemotron_native_backend/host/diagnostic.py new file mode 100644 index 0000000..e96584b --- /dev/null +++ b/python/nemotron_native_backend/host/diagnostic.py @@ -0,0 +1,33 @@ +import json +import threading +import time +from pathlib import Path + + +class Trace: + def __init__(self, path=None): + self.stream = None + if path: + target = Path(path) + target.parent.mkdir(parents=True, exist_ok=True) + self.stream = target.open("a", buffering=1) + self.lock = threading.Lock() + + def event(self, name, **fields): + if self.stream is not None: + with self.lock: + self.stream.write( + json.dumps( + dict( + event=name, + monotonic_ns=time.monotonic_ns(), + thread_cpu_ns=time.thread_time_ns(), + thread=threading.get_native_id(), + **fields, + ) + ) + + "\n" + ) + + +TRACE = Trace() diff --git a/python/nemotron_native_backend/host/exact_admission.py b/python/nemotron_native_backend/host/exact_admission.py new file mode 100644 index 0000000..881b32c --- /dev/null +++ b/python/nemotron_native_backend/host/exact_admission.py @@ -0,0 +1,169 @@ +import asyncio +import json +import math +import time +from collections import deque +from pathlib import Path + +from .. import r31 + + +def install(base, *, incomplete_timeout_seconds=10.0): + timeout = float(incomplete_timeout_seconds) + if not math.isfinite(timeout) or not 0.0 < timeout <= 60.0: + raise ValueError("Invalid full-physical-cohort deadline") + installed = getattr(base, "_r21_exact_admission_timeout", None) + if installed is not None: + if installed != timeout: + raise ValueError("Exact admission is already installed with another deadline") + return + parent = base.TimedBatcher + error_type = base.BackendError + + class PhysicalWidthBatcher(parent): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if not self._holding or self.expected_cohort_size != self.backend.batch: + raise ValueError("Exact admission requires the physical-width holding path") + if self.expected_cohort_size <= 1 or not isinstance(self.backend.epoch, int): + raise ValueError("Exact admission requires a wide backend with an epoch") + self._r21_initial_dispatch = False + self._r21_incomplete_rejections = 0 + self._r21_incomplete_timeout_seconds = timeout + self._r31_incremental = r31.enabled("staged_admission") + self._r31_groups = {} + self._r31_abort_seen = False + self._r31_dispatch = {} if r31.enabled("accept_at_dispatch") else None + + def submit(self, request): + handle = super().submit(request) + if self._r31_incremental: + item = self._queue[-1] + if item.request is not request: + raise RuntimeError("submit did not append the admitted request") + self._r31_groups.setdefault(request.compatibility_key, []).append(item) + if self._r31_dispatch is not None: + self._r31_dispatch[request.request_id] = asyncio.Event() + return handle + + def r31_dispatch_event(self, request_id): + return None if self._r31_dispatch is None else self._r31_dispatch.get(request_id) + + def abort(self, request_id): + self._r31_abort_seen = True + if self._r31_dispatch is not None: + self._r31_dispatch.pop(request_id, None) + return super().abort(request_id) + + def _r31_regroup(self): + groups = {} + for item in self._queue: + groups.setdefault(item.request.compatibility_key, []).append(item) + self._r31_groups = groups + + def _purge_aborted(self): + if not self._r31_incremental: + return super()._purge_aborted() + if self._r31_abort_seen: + self._r31_abort_seen = False + super()._purge_aborted() + self._r31_regroup() + + def _grouped_queue(self): + if not self._r31_incremental: + return super()._grouped_queue() + return {key: items for (key, items) in self._r31_groups.items() if items} + + def _reject_unrepresentable(self, items): + super()._reject_unrepresentable(items) + if self._r31_incremental: + self._r31_regroup() + + def _finish_queued_error(self, item, error): + self.backend.arrivals.pop(item.request.request_id, None) + if self._r31_dispatch is not None: + self._r31_dispatch.pop(item.request.request_id, None) + return super()._finish_queued_error(item, error) + + def _reject_incomplete_physical_cohort(self, items): + items = list(items) + error = error_type( + "incomplete_physical_cohort", + f"{len(items)} compatible request(s) did not form physical " + f"B{self.expected_cohort_size} within {timeout:g} seconds; " + "partial neural execution is disabled", + retryable=True, + ) + chosen = {id(item) for item in items} + self._queue = deque(item for item in self._queue if id(item) not in chosen) + for item in items: + self._finish_queued_error(item, error) + self._r21_incomplete_rejections += len(items) + if self._r31_incremental: + self._r31_regroup() + receipt = dict( + code="incomplete_physical_cohort", + request_ids=[item.request.request_id for item in items], + epoch=self.backend.epoch, + pending_after=len(self._pending), + arrivals_after=len(self.backend.arrivals), + monotonic_ns=time.monotonic_ns(), + ) + with (Path(self.backend.receipt_dir) / "r21-incomplete-cohorts.jsonl").open( + "a" + ) as stream: + stream.write(json.dumps(receipt, separators=(",", ":")) + "\n") + + def _plan_cohorts(self, groups, now): + groups = list(groups.values()) + if ( + self.backend.epoch == 0 + and not self._r21_initial_dispatch + and self._committed_batches == 0 + and len(groups) == 1 + and len(groups[0]) == 1 + ): + return super()._plan_cohorts( + {groups[0][0].request.compatibility_key: groups[0]}, now + ) + plan = [] + waits = [] + committed = self._committed_batches + expected = self.expected_cohort_size + for items in groups: + if not items: + continue + if len(items) >= expected: + if committed < self.max_inflight_batches: + plan.append(items[:expected]) + committed += 1 + continue + remaining = timeout - (now - items[0].enqueued_at) + if remaining <= 0.0: + self._reject_incomplete_physical_cohort(items) + else: + waits.append(remaining) + return plan, None if plan else min(waits) if waits else None + + def _launch_cohort(self, cohort): + readiness = ( + len(cohort) == 1 + and self.backend.epoch == 0 + and not self._r21_initial_dispatch + and self._committed_batches == 0 + ) + if len(cohort) != self.expected_cohort_size and not readiness: + raise RuntimeError("Exact admission refused a partial neural cohort") + self._r21_initial_dispatch = True + result = super()._launch_cohort(cohort) + if self._r31_incremental: + self._r31_regroup() + if self._r31_dispatch is not None: + for item in cohort: + event = self._r31_dispatch.pop(item.request.request_id, None) + if event is not None: + event.set() + return result + + base.TimedBatcher = PhysicalWidthBatcher + base._r21_exact_admission_timeout = timeout diff --git a/python/nemotron_native_backend/host/first_fence.py b/python/nemotron_native_backend/host/first_fence.py new file mode 100644 index 0000000..0e553ce --- /dev/null +++ b/python/nemotron_native_backend/host/first_fence.py @@ -0,0 +1,130 @@ +import asyncio +import hashlib +import json +import threading +import time + +REGISTRY = {} + + +class FirstFence: + def __init__(self, request_ids, timeout=0.050): + if len(set(request_ids)) != len(request_ids) or not 0 < timeout <= 0.050: + raise ValueError("Invalid first-audio fence contract") + self.ids = tuple(request_ids) + self.timeout = timeout + self.condition = threading.Condition() + self.owned = set() + self.retired = set() + self.used = False + self.closed = False + self.waits = [] + + def acknowledge(self, request_id): + with self.condition: + if request_id not in self.ids: + raise ValueError("ACK belongs to another cohort") + self.owned.add(request_id) + self.condition.notify_all() + + def retire(self, request_id): + with self.condition: + if request_id in self.ids: + self.retired.add(request_id) + self.condition.notify_all() + + def close(self): + with self.condition: + self.closed = True + self.condition.notify_all() + + def wait(self, lanes): + if self.used: + return + required = {self.ids[lane] for lane in lanes} + if not required: + return + self.used = True + begin = time.monotonic_ns() + deadline = time.monotonic() + self.timeout + with self.condition: + while not self.closed and not required <= self.owned | self.retired: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + self.condition.wait(remaining) + self.waits.append( + dict( + requests=len(required), + start_ns=begin, + end_ns=time.monotonic_ns(), + missing=len(required - self.owned - self.retired), + closed=self.closed, + ) + ) + + def receipt(self): + return dict( + policy="once-after-original-continuation-before-G9-frontend-owned-ACK", + maximum_wait_ms=self.timeout * 1000, + waits=list(self.waits), + acknowledged=len(self.owned), + retired=len(self.retired), + public_http_ack=False, + requests_sha256=hashlib.sha256(json.dumps(self.ids).encode()).hexdigest(), + ) + + +def acknowledge(request_id): + gate = REGISTRY.get(request_id) + if gate is not None: + gate.acknowledge(request_id) + + +def retire(request_id): + gate = REGISTRY.get(request_id) + if gate is not None: + gate.retire(request_id) + + +def install_backend(base): + previous = base.NativeBackend.stream_batch + + async def stream_batch(self, requests, emit): + if not hasattr(self, "_r21_fence_owner"): + self._r21_fence_owner = asyncio.Lock() + async with self._r21_fence_owner: + target = self.engines[len(requests)] + gate = FirstFence([request.request_id for request in requests]) + if any(request_id in REGISTRY for request_id in gate.ids): + raise RuntimeError("Duplicate active fence request") + target._r21_first_fence = gate + REGISTRY.update({request_id: gate for request_id in gate.ids}) + try: + return await previous(self, requests, emit) + finally: + gate.close() + target._r21_first_fence = None + for request_id in gate.ids: + REGISTRY.pop(request_id, None) + + base.NativeBackend.stream_batch = stream_batch + + +def install_engine(engine): + previous = engine.stream + + def stream(*args, **kwargs): + iterator = previous(*args, **kwargs) + try: + for event in iterator: + gate = getattr(engine, "_r21_first_fence", None) + if gate is not None and event["type"] == "receipt": + event["receipt"]["r21_first_fence"] = gate.receipt() + yield event + if gate is not None and event["type"] == "continuation": + gate.wait(event["lanes"]) + finally: + iterator.close() + + engine.stream = stream diff --git a/python/nemotron_native_backend/host/ready_writer.py b/python/nemotron_native_backend/host/ready_writer.py new file mode 100644 index 0000000..e6d7fdd --- /dev/null +++ b/python/nemotron_native_backend/host/ready_writer.py @@ -0,0 +1,118 @@ +import asyncio + + +class ReadyWriter: + def __init__(self, write, capacity=4096, group=64): + if not 1 <= group <= 64 or capacity < group: + raise ValueError("Invalid ready-writer bounds") + self.write = write + self.queue = asyncio.Queue(maxsize=capacity) + self.group = group + self.closed = False + self.failure = None + self.peak = 0 + self.frames = 0 + self.records = 0 + self.maximum_group = 0 + self.task = asyncio.create_task(self.run()) + + async def send(self, value): + if self.closed: + raise ConnectionError("Ready writer closed") from self.failure + future = asyncio.get_running_loop().create_future() + put = None + try: + try: + self.queue.put_nowait((value, future)) + except asyncio.QueueFull: + put = asyncio.create_task(self.queue.put((value, future))) + done, _ = await asyncio.wait((put, self.task), return_when=asyncio.FIRST_COMPLETED) + if self.task in done: + raise ConnectionError("Ready writer stopped") from self.failure + await put + self.peak = max(self.peak, self.queue.qsize()) + await future + finally: + if put is not None and not put.done(): + put.cancel() + await asyncio.gather(put, return_exceptions=True) + if not future.done(): + future.cancel() + elif not future.cancelled(): + future.exception() + + async def run(self): + pending = None + active = [] + try: + while True: + first = pending if pending is not None else await self.queue.get() + pending = None + active = [first] + if first[0].get("kind") == "audio_slab": + while len(active) < self.group: + try: + following = self.queue.get_nowait() + except asyncio.QueueEmpty: + break + if following[0].get("kind") != "audio_slab": + pending = following + break + active.append(following) + value = ( + active[0][0] + if len(active) == 1 + else dict( + kind="audio_slab_batch", + protocol_version=1, + records=[item[0] for item in active], + ) + ) + await self.write(value) + self.frames += 1 + self.records += len(active) + self.maximum_group = max(self.maximum_group, len(active)) + for _, future in active: + if not future.done(): + future.set_result(None) + active = [] + except BaseException as error: + self.failure = error + finally: + self.closed = True + if pending is not None: + active.append(pending) + while not self.queue.empty(): + active.append(self.queue.get_nowait()) + for _, future in active: + if not future.done(): + future.set_exception(ConnectionError("Ready writer stopped")) + + async def close(self): + self.task.cancel() + await asyncio.gather(self.task, return_exceptions=True) + + +def ack_records(incoming): + if incoming.get("kind") == "audio_ack": + records = [incoming] + if set(incoming) != {"kind", "protocol_version", "arena", "slot", "generation"}: + raise ValueError("Invalid single ACK") + elif incoming.get("kind") == "audio_ack_batch": + if set(incoming) != {"kind", "protocol_version", "records"}: + raise ValueError("Invalid ACK batch") + records = incoming["records"] + if not isinstance(records, list) or not 1 <= len(records) <= 64: + raise ValueError("Invalid ACK batch extent") + if any( + not isinstance(item, dict) or set(item) != {"arena", "slot", "generation"} + for item in records + ): + raise ValueError("Invalid ACK record") + else: + raise ValueError("Not an ACK frame") + if incoming["protocol_version"] != 1: + raise ValueError("Invalid ACK version") + if len({(item["arena"], item["slot"], item["generation"]) for item in records}) != len(records): + raise ValueError("Duplicate ACK in one frame") + return records diff --git a/python/nemotron_native_backend/host/serialization.py b/python/nemotron_native_backend/host/serialization.py new file mode 100644 index 0000000..db87dcb --- /dev/null +++ b/python/nemotron_native_backend/host/serialization.py @@ -0,0 +1,145 @@ +import hashlib +import math +from collections.abc import Mapping, Sequence +from pathlib import Path + +VERSION = "1.2.2" + + +def install(protocol): + import msgpack + import msgpack._cmsgpack + + if msgpack.__version__ != VERSION or msgpack.Packer.__module__ != "msgpack._cmsgpack": + raise RuntimeError("required native msgpack version is " + VERSION) + if getattr(protocol, "_r11_native_serializer", None): + return protocol._r11_native_serializer + original = protocol.packb + native = msgpack.packb + error_type = protocol.ProtocolError + limit = protocol.MAX_COLLECTION_ENTRIES + + def validate(value, depth=0): + if depth > 64: + raise error_type("MessagePack nesting exceeds 64") + kind = type(value) + if value is None or kind in (bool, str, bytes, bytearray, memoryview): + return + if kind is int: + if not -(1 << 63) <= value <= (1 << 64) - 1: + raise error_type("integer is outside MessagePack int64/uint64") + return + if kind is float: + if not math.isfinite(value): + raise error_type("non-finite floats are forbidden") + return + if kind is dict: + if len(value) > limit: + raise error_type("MessagePack map is unreasonably large") + for key, child in value.items(): + if type(key) is not str: + raise error_type("protocol map keys must be strings") + validate(key, depth + 1) + validate(child, depth + 1) + return + if kind in (list, tuple): + if len(value) > limit: + raise error_type("MessagePack array is unreasonably large") + for child in value: + validate(child, depth + 1) + return + if isinstance(value, (Mapping, Sequence, int, float, str)): + raise TypeError("legacy compatible value") + raise error_type("unsupported MessagePack type: " + kind.__name__) + + def packb(value): + try: + validate(value) + except TypeError: + return original(value) + try: + return native(value, use_bin_type=True, use_single_float=False) + except (ValueError, TypeError, OverflowError, BufferError): + return original(value) + + path = Path(msgpack._cmsgpack.__file__) + record = { + "version": VERSION, + "implementation": msgpack.Packer.__module__, + "extension": str(path), + "sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + "wire_contract": "legacy exact bytes, validation and frame limits preserved", + } + protocol.packb = packb + from .. import r31 + + if r31.enabled("native_unpack"): + protocol.unpackb = native_unpackb(protocol) + record["unpack"] = ( + "native msgpack.unpackb with the pure-Python validation contract; " + "float32 encodings decode instead of being rejected (never sent by the frontend)" + ) + protocol._r11_native_serializer = record + return record + + +def native_unpackb(protocol): + import msgpack + + error_type = protocol.ProtocolError + limit = protocol.MAX_COLLECTION_ENTRIES + + def pairs(items): + result = {} + for key, child in items: + if type(key) is not str: + raise error_type("protocol map keys must be strings") + if key in result: + raise error_type(f"duplicate protocol map key {key!r}") + result[key] = child + return result + + def check(value, depth): + if depth > 64: + raise error_type("MessagePack nesting exceeds 64") + kind = type(value) + if kind is float: + if not math.isfinite(value): + raise error_type("non-finite floats are forbidden") + elif kind is dict or kind is list: + if value and depth + 1 > 64: + raise error_type("MessagePack nesting exceeds 64") + for child in value.values() if kind is dict else value: + if type(child) in (dict, list, float): + check(child, depth + 1) + + def unpackb(payload): + try: + value = msgpack.unpackb( + payload, + raw=False, + use_list=True, + strict_map_key=False, + object_pairs_hook=pairs, + max_array_len=limit, + max_map_len=limit, + max_ext_len=0, + max_str_len=len(payload), + max_bin_len=len(payload), + ) + except error_type: + raise + except msgpack.ExtraData as error: + raise error_type("trailing bytes after MessagePack value") from error + except ( + ValueError, + TypeError, + UnicodeDecodeError, + msgpack.FormatError, + msgpack.StackError, + ) as error: + raise error_type("invalid MessagePack value: " + str(error)) from error + check(value, 0) + return value + + return unpackb diff --git a/python/nemotron_native_backend/host/slab_audit.py b/python/nemotron_native_backend/host/slab_audit.py new file mode 100644 index 0000000..442e78e --- /dev/null +++ b/python/nemotron_native_backend/host/slab_audit.py @@ -0,0 +1,61 @@ +import hashlib +import json + + +def pcm16(payload): + import numpy as np + + x = np.frombuffer(payload, dtype="= 1, 32767, y)) + return y.astype(" None: + async with self._client_lock: + if self._stopping or self._client_active: + writer.close() + await writer.wait_closed() + return + self._client_active = True + self._client_writer = writer + pool = None + audit = None + send_lock = asyncio.Lock() + held_requests = {} + held_tasks = set() + request_ids: set[str] = set() + completed_ids = CompletedRequestIdWindow(self.completed_request_id_capacity) + first_leases = {} + config = self.native_host_options + queued_input = deque() + input_batches = 0 + input_records = 0 + maximum_input_group = 0 + ready_writer = None + dma = self.native_dma + + async def write(value): + self.native_trace.event( + "descriptor_write_begin", request_id=value.get("request_id"), kind=value["kind"] + ) + await write_frame(writer, value, self.max_frame_bytes) + self.native_trace.event( + "descriptor_write_end", request_id=value.get("request_id"), kind=value["kind"] + ) + + if config.get("f2"): + ready_writer = ReadyWriter(write) + + async def send_held(value, held): + held.before_send(value) + original = value + lease = None + advertised = False + try: + if value.get("kind") == "audio": + payload = value["data"] + self.native_trace.event( + "slab_publish_begin", + request_id=value["request_id"], + sequence=value["sequence"], + bytes=len(payload), + ) + lease = await held.wait_owned(pool.publish(payload)) + self.native_trace.event( + "slab_publish_end", + request_id=value["request_id"], + sequence=value["sequence"], + slot=lease["slot"], + generation=lease["generation"], + ) + value = {key: item for key, item in value.items() if key != "data"} + value.update(kind="audio_slab", **lease) + self.native_trace.event( + "descriptor_lock_begin", request_id=value.get("request_id"), kind=value["kind"] + ) + async with send_lock: + held.before_send(value) + if lease is not None: + if original["sequence"] == 0: + first_leases[(lease["arena"], lease["slot"], lease["generation"])] = ( + original["request_id"] + ) + if audit is not None: + audit.audio(original) + held.begin_send(value) + advertised = True + await write(value) + held.after_send(value) + finally: + if lease is not None and not advertised and not pool.closed: + first_leases.pop((lease["arena"], lease["slot"], lease["generation"]), None) + pool.acknowledge(lease["arena"], lease["slot"], lease["generation"]) + pool.counts["locally_reclaimed_unsent"] += 1 + + def abort_held(request_id): + held = held_requests.get(request_id) if isinstance(request_id, str) else None + if held is not None and held.terminal_started and not held.delivery_task.done(): + held.late_abort_count += 1 + return "KNOWN_TERMINAL_PENDING" + + def abort_owned(ident): + owned = held_requests.get(ident) + if owned is not None and owned.terminal_sent: + return False + accepted = owned.abort() if owned is not None else False + return self.batcher.abort(ident) or accepted + + return abort_response(request_id, completed_ids, abort_owned) + + def held_done(task, handle): + held_tasks.discard(task) + if held_requests.get(handle.request_id) is handle: + held_requests.pop(handle.request_id) + + async def send(value: dict[str, Any]) -> None: + held = held_requests.get(value.get("request_id")) + if held is not None and asyncio.current_task() is held.delivery_task: + return await send_held(value, held) + if value.get("kind") == "audio": + payload = value["data"] + self.native_trace.event( + "slab_publish_begin", + request_id=value["request_id"], + sequence=value["sequence"], + bytes=len(payload), + ) + lease = dma.publish_direct(payload) if dma is not None else None + if lease is None: + lease = await pool.publish(payload) + self.native_trace.event( + "slab_publish_end", + request_id=value["request_id"], + sequence=value["sequence"], + slot=lease["slot"], + generation=lease["generation"], + ) + if value["sequence"] == 0: + first_leases[(lease["arena"], lease["slot"], lease["generation"])] = value[ + "request_id" + ] + if audit is not None: + audit.audio(value) + value = {k: v for k, v in value.items() if k != "data"} + value.update(kind="audio_slab", **lease) + if ready_writer is not None: + await ready_writer.send(value) + return + self.native_trace.event( + "descriptor_lock_begin", request_id=value.get("request_id"), kind=value["kind"] + ) + async with send_lock: + await write(value) + + try: + pool = SlabPool(slots=1024) + if dma is not None: + dma.establish(pool, self.backend) + if self.slab_audit_enabled: + audit = SlabAudit( + Path(self.backend.receipt_dir, pool.arena + "-expected-pcm16.jsonl") + ) + hello = validate_envelope( + await read_frame(reader, self.max_frame_bytes), from_client=True + ) + self._validate_hello(hello) + await send( + message( + "worker_ready", + worker_instance_id=self.worker_instance_id, + **self.manifest.worker_ready_fields(), + cuda_device_name=self.backend.identity.cuda_device_name, + cuda_compute_capability=list(self.backend.identity.cuda_compute_capability), + ) + ) + while True: + incoming = ( + queued_input.popleft() + if queued_input + else await read_frame(reader, self.max_frame_bytes) + ) + if isinstance(incoming, dict) and incoming.get("kind") == "submit_batch": + if ( + not config.get("f4") + or set(incoming) != {"kind", "protocol_version", "records"} + or incoming["protocol_version"] != 1 + ): + raise ProtocolError("invalid submit batch envelope") + records = incoming["records"] + if ( + not isinstance(records, list) + or not 1 <= len(records) <= 64 + or any( + not isinstance(item, dict) or item.get("kind") != "submit" + for item in records + ) + ): + raise ProtocolError("invalid submit batch records") + queued_input.extend(records) + input_batches += 1 + input_records += len(records) + maximum_input_group = max(maximum_input_group, len(records)) + continue + if isinstance(incoming, dict) and incoming.get("kind") in ( + "audio_ack", + "audio_ack_batch", + ): + if incoming["kind"] == "audio_ack_batch" and not config.get("f2"): + raise ProtocolError("unregistered ACK batch") + try: + records = ack_records(incoming) + except (ValueError, KeyError, TypeError) as error: + raise ProtocolError(str(error)) from error + for ack in records: + handled = ( + dma.acknowledge_direct(ack["arena"], ack["slot"], ack["generation"]) + if dma is not None + else False + ) + if not handled: + pool.acknowledge(ack["arena"], ack["slot"], ack["generation"]) + request_id = first_leases.pop( + (ack["arena"], ack["slot"], ack["generation"]), None + ) + self.native_trace.event( + "slab_ack", + request_id=request_id, + slot=ack["slot"], + generation=ack["generation"], + ) + if request_id is not None: + first_fence.acknowledge(request_id) + continue + incoming = validate_envelope(incoming, from_client=True) + kind = incoming["kind"] + if kind == "submit": + self.native_trace.event("submit_read", request_id=incoming.get("request_id")) + request = StreamingSynthesisRequest.from_message(incoming) + if request.request_id in request_ids or request.request_id in completed_ids: + await send( + self._failed( + request.request_id, + BackendError( + "duplicate_request_id", + "request ID was already used on this connection", + ), + ) + ) + continue + try: + handle = self.batcher.submit(request) + except BackendError as error: + await send(self._failed(request.request_id, error)) + continue + request_ids.add(request.request_id) + if self._initial_hold_seconds: + handle = InitialAudioHold( + handle, self._initial_hold_seconds, self._initial_hold_budget + ) + held_requests[request.request_id] = handle + self.native_trace.event("submit_enqueued", request_id=request.request_id) + lookup = getattr(self.batcher, "r31_dispatch_event", None) + dispatched = lookup(request.request_id) if lookup is not None else None + if dispatched is None: + await send(message("accepted", request_id=request.request_id)) + task = asyncio.create_task( + self._deliver_stream(handle, send, request_ids, completed_ids, dispatched), + name=f"deliver-stream-{request.request_id}", + ) + self._delivery_tasks.add(task) + task.add_done_callback(self._delivery_done) + if self._initial_hold_seconds: + handle.delivery_task = task + held_tasks.add(task) + task.add_done_callback(lambda task, handle=handle: held_done(task, handle)) + elif kind == "abort": + self._validate_exact_keys(incoming, {"request_id"}) + request_id = incoming["request_id"] + action = ( + abort_held(request_id) + if self._initial_hold_seconds + else abort_response(request_id, completed_ids, self.batcher.abort) + ) + if action != "UNKNOWN_REQUEST": + first_fence.retire(request_id) + if audit is not None: + with Path( + self.backend.receipt_dir, pool.arena + "-cancellation.jsonl" + ).open("a") as stream: + stream.write( + json.dumps(dict(request_id=request_id, action=action)) + "\n" + ) + if action == "KNOWN_TERMINAL": + await send(message("aborted", request_id=request_id)) + elif action == "UNKNOWN_REQUEST": + await send( + self._failed( + str(request_id), + BackendError("unknown_request", "request is not pending"), + ) + ) + elif kind == "ping": + self._validate_exact_keys(incoming, {"nonce"}) + nonce = self._validate_ping_nonce(incoming["nonce"]) + await send( + message( + "pong", + nonce=nonce, + queue_depth=self.batcher.queue_depth, + active_batches=self.batcher.active_batches, + ) + ) + else: + raise ProtocolError(f"unexpected post-handshake kind {kind!r}") + except (asyncio.IncompleteReadError, ConnectionError, BrokenPipeError): + pass + finally: + if self._initial_hold_seconds: + for request_id in tuple(request_ids): + if request_id in held_requests: + held_requests[request_id].abort() + first_fence.retire(request_id) + self.batcher.abort(request_id) + writer.close() + with contextlib.suppress(Exception): + await writer.wait_closed() + remaining = tuple(held_tasks) + if remaining: + done, pending = await asyncio.wait(remaining, timeout=5.0) + for task in pending: + task.cancel() + await asyncio.gather(*remaining, return_exceptions=True) + self._initial_hold_budget.snapshot(self.backend.receipt_dir) + if ready_writer is not None: + await ready_writer.close() + if dma is not None: + dma.retire_direct() + for request_id in request_ids: + first_fence.retire(request_id) + if audit is not None: + audit.close() + if pool is not None: + pool.close() + directory = getattr(self.backend, "receipt_dir", None) + if directory is not None: + receipt = pool.receipt() + receipt["r21"] = dict( + options=config, + input_batch_records=input_records, + input_batches=input_batches, + maximum_input_batch=maximum_input_group, + output_frames=ready_writer.frames if ready_writer else None, + output_records=ready_writer.records if ready_writer else None, + maximum_output_batch=ready_writer.maximum_group if ready_writer else None, + ) + Path(directory, pool.arena + ".json").write_text( + json.dumps(receipt, indent=2) + "\n" + ) + for request_id in tuple(request_ids): + self.batcher.abort(request_id) + writer.close() + with contextlib.suppress(Exception): + await writer.wait_closed() + async with self._client_lock: + if self._client_writer is writer: + self._client_writer = None + self._client_active = False + + +def install_server(base): + base.NativeServer = type("R20LeasedPCMServer", (SlabServerMixin, base.NativeServer), {}) diff --git a/python/nemotron_native_backend/inflight_batcher.py b/python/nemotron_native_backend/inflight_batcher.py new file mode 100644 index 0000000..9f0173d --- /dev/null +++ b/python/nemotron_native_backend/inflight_batcher.py @@ -0,0 +1,502 @@ +from __future__ import annotations + +import asyncio +import json +import time +from collections import deque +from dataclasses import dataclass + +from .backend import BackendError +from .batcher import QueueFull, RequestAborted +from .execution.native_assets import WIDTHS +from .streaming_batcher import StreamingHandle +from .streaming_contract import OneChunkLookaheadPublisher +from .streaming_models import StreamingAudioChunk, StreamingResult + + +class InflightHandle(StreamingHandle): + def __init__(self, request_id, result, chunks, wake): + super().__init__(request_id, result, chunks) + self._wake = wake + + async def next_chunk(self): + try: + return await super().next_chunk() + finally: + self._wake.set() + + +@dataclass +class PendingRequest: + request: object + handle: InflightHandle + submitted_ns: int + slot: int | None = None + epoch: int = 0 + publisher: object = None + admitted_ns: int | None = None + first_audio_ns: int | None = None + chunks: int = 0 + samples: int = 0 + aborted: bool = False + + +class InflightBatcher: + def __init__( + self, + backend, + *, + mode="inflight", + max_queue_depth=None, + chunk_queue_depth=4, + max_active_requests=None, + ): + if mode not in ("inflight", "fixed_request") or backend.batch not in WIDTHS: + raise ValueError("In-flight execution requires a supported physical width") + self.batch = backend.batch + if max_active_requests is not None and ( + mode != "inflight" + or type(max_active_requests) is not int + or not 1 <= max_active_requests <= self.batch + ): + raise ValueError("active admission cap requires inflight mode and 1..physical width") + self.max_active_requests = max_active_requests + max_queue_depth = max(128, 4 * self.batch) if max_queue_depth is None else max_queue_depth + if max_queue_depth < self.batch or chunk_queue_depth < 2: + raise ValueError("Insufficient bounded request or chunk capacity") + self.backend = backend + self.mode = mode + self._admission_signal = ( + getattr(backend.inflight_engine, "admission_requested", None) + if mode == "inflight" + else None + ) + self.max_queue_depth = max_queue_depth + self.chunk_queue_depth = chunk_queue_depth + self._pending = {} + self._queue = deque() + self._slots = [None] * self.batch + self._epochs = [0] * self.batch + self._retire = [] + self._defer_admission_reset = mode == "inflight" and getattr( + getattr(backend.inflight_engine, "policy", None), "defer_admission_reset", False + ) + self._deferred_admissions = {} + self._deferred_admission_count = 0 + self._wake = asyncio.Event() + self._runner = None + self._closed = False + self._failure = None + self._shutdown_task = None + self._readiness_done = False + self._records = [] + self.on_fatal = None + + @property + def queue_depth(self): + return sum(not item.aborted for item in self._queue) + + @property + def active_batches(self): + return int(any(item is not None for item in self._slots)) + + def start(self): + if self._runner is not None: + raise RuntimeError("In-flight controller is already started") + self._runner = asyncio.create_task(self._run(), name=f"native-b{self.batch}-inflight") + + def submit(self, request): + if self._closed or self._failure is not None: + raise BackendError( + "worker_stopping", "In-flight execution is unavailable", retryable=True + ) + if request.request_id in self._pending: + raise BackendError("duplicate_request_id", "Request ID is already pending") + if not 1 <= len(request.token_ids) <= 512: + raise BackendError( + "unsupported_text_length", "Native execution supports 1..512 text tokens" + ) + if len(self._pending) >= self.max_queue_depth: + raise QueueFull() + loop = asyncio.get_running_loop() + handle = InflightHandle( + request.request_id, + loop.create_future(), + asyncio.Queue(maxsize=self.chunk_queue_depth), + self._wake, + ) + item = PendingRequest(request, handle, time.monotonic_ns()) + self._pending[request.request_id] = item + self._queue.append(item) + if getattr(getattr(self.backend.inflight_engine, "policy", None), "diagnostics", False): + self._record(item, "queued") + if self.max_active_requests is not None: + self._refresh_admission_signal() + elif self._admission_signal is not None: + self._admission_signal.set() + self._wake.set() + return handle + + def abort(self, request_id): + item = self._pending.get(request_id) + if item is None: + return False + item.aborted = True + self._error(item, RequestAborted()) + if self.max_active_requests is not None: + self._refresh_admission_signal() + self._wake.set() + return True + + def _error(self, item, error): + while not item.handle._chunks.empty(): + item.handle._chunks.get_nowait() + item.handle._chunks.put_nowait(error) + if not item.handle.result.done(): + item.handle.result.set_exception(error) + if item.slot is None: + if isinstance(error, RequestAborted) and getattr( + getattr(self.backend.inflight_engine, "policy", None), "diagnostics", False + ): + self._record(item, "cancelled", error=str(error)) + self._pending.pop(item.request.request_id, None) + + def _record(self, item, status, **extra): + row = dict( + request_id=item.request.request_id, + seed=item.request.seed, + normalized_text=item.request.normalized_text, + token_count=len(item.request.token_ids), + mode=self.mode, + physical_batch=self.batch, + slot=item.slot, + epoch=item.epoch, + submitted_ns=item.submitted_ns, + admitted_ns=item.admitted_ns, + first_audio_ns=item.first_audio_ns, + observed_ns=time.monotonic_ns(), + chunks=item.chunks, + samples=item.samples, + status=status, + **extra, + ) + if self.max_active_requests is not None: + row["max_active_requests"] = self.max_active_requests + self._records.append(row) + + def _flush_records(self): + if self._records: + payload = "".join( + json.dumps(row, separators=(",", ":")) + "\n" for row in self._records + ) + with (self.backend.receipt_dir / "inflight-requests.jsonl").open("a") as stream: + stream.write(payload) + self._records.clear() + engine = self.backend.inflight_engine + if getattr(getattr(engine, "policy", None), "diagnostics", False): + rows = engine.drain_diagnostics() + if rows: + payload = "".join(json.dumps(row, separators=(",", ":")) + "\n" for row in rows) + with (self.backend.receipt_dir / "inflight-diagnostics.jsonl").open("a") as stream: + stream.write(payload) + + def _free(self, item): + self._pending.pop(item.request.request_id, None) + if item.slot is not None: + self._slots[item.slot] = None + self._retire.append((item.slot, item.epoch)) + if self.max_active_requests is not None: + self._refresh_admission_signal() + + def _publish(self, item, chunk): + audio = StreamingAudioChunk( + chunk.request_id, chunk.sequence, chunk.payload, chunk.samples, chunk.is_final + ) + audio.validate() + item.handle._chunks.put_nowait(audio) + item.chunks += 1 + item.samples += chunk.samples + if chunk.samples and item.first_audio_ns is None: + item.first_audio_ns = time.monotonic_ns() + if getattr(getattr(self.backend.inflight_engine, "policy", None), "diagnostics", False): + self._record(item, "publication_enqueued", chunk_samples=chunk.samples) + + def _consume(self, events): + for event in events: + slot = event["slot"] + if not 0 <= slot < self.batch: + raise RuntimeError("Execution returned an invalid physical slot") + item = self._slots[slot] + if item is None or item.epoch != event["epoch"]: + continue + if item.request.request_id != event["request_id"]: + raise RuntimeError("Execution returned another request's slot ownership") + if item.aborted: + continue + if event["type"] == "chunk": + chunks = item.publisher.decoded( + item.request.request_id, event["pcm_f32_le"], event["samples"] + ) + for chunk in chunks: + self._publish(item, chunk) + elif event["type"] == "finish": + if not item.samples: + self._error(item, BackendError("empty_audio", "Request ended without audio")) + self._record(item, "failed", error="empty_audio") + self._free(item) + continue + self._publish(item, item.publisher.finish(item.request.request_id)) + result = StreamingResult(item.request.request_id, item.chunks, item.samples) + result.validate() + if not item.handle.result.done(): + item.handle.result.set_result(result) + self._record( + item, + "finished", + termination=event.get("reason"), + engine_startup_priority_launches=event.get("receipt", {}).get( + "engine_startup_priority_launches" + ), + engine_admission_yields=event.get("receipt", {}).get("engine_admission_yields"), + ) + self._free(item) + self._readiness_done = True + else: + raise RuntimeError("Unsupported in-flight execution event") + + def _refresh_admission_signal(self): + if self._admission_signal is not None: + waiting = any(not item.aborted for item in self._queue) + owners = sum(item is not None for item in self._slots) + if waiting and owners < self.max_active_requests: + self._admission_signal.set() + else: + self._admission_signal.clear() + + def _select_admissions(self): + if not self._queue: + if self._admission_signal is not None: + self._admission_signal.clear() + return [] + self._queue = deque(item for item in self._queue if not item.aborted) + if not self._queue and self._admission_signal is not None: + self._admission_signal.clear() + free = [slot for slot, item in enumerate(self._slots) if item is None] + if self.mode == "fixed_request": + if len(free) != self.batch: + return [] + if self._readiness_done and len(self._queue) < self.batch: + return [] + if ( + not self._readiness_done + and len(self._queue) not in (1,) + and len(self._queue) < self.batch + ): + return [] + selected = [] + peers = self.batch - len(free) + for slot in free: + if self.max_active_requests is not None and peers >= self.max_active_requests: + break + if not self._queue: + break + item = self._queue.popleft() + self._epochs[slot] += 1 + item.slot = slot + item.epoch = self._epochs[slot] + item.admitted_ns = time.monotonic_ns() + item.publisher = OneChunkLookaheadPublisher( + [item.request.request_id], eager=True, completion_marker=True + ) + self._slots[slot] = item + selected.append(item) + self._record(item, "admitted", peers=peers) + peers += 1 + if self.max_active_requests is not None: + self._refresh_admission_signal() + elif not self._queue and self._admission_signal is not None: + self._admission_signal.clear() + return selected + + def _tick_deferred(self, retirements, cancellations, admissions, progress): + engine = self.backend.inflight_engine + dropped = set() + for slot, epoch in cancellations: + key = (slot, epoch) + if self._deferred_admissions.pop(key, None) is not None: + dropped.add(key) + else: + engine.cancel(slot, epoch) + initialized_retirements = [] + for slot, epoch in retirements: + key = (slot, epoch) + if key in dropped or self._deferred_admissions.pop(key, None) is not None: + dropped.add(key) + else: + initialized_retirements.append(key) + if getattr(getattr(engine, "policy", None), "retire_batch", False): + engine.retire_many(initialized_retirements) + else: + for slot, epoch in initialized_retirements: + engine.retire(slot, epoch) + pending_slots = {slot for slot, _ in self._deferred_admissions} + for item in admissions: + if item.slot in pending_slots: + raise RuntimeError("physical slot has conflicting deferred admission ownership") + if self._slots[item.slot] is not item or self._epochs[item.slot] != item.epoch: + raise RuntimeError("deferred admission lost slot ownership") + self._deferred_admissions[item.slot, item.epoch] = dict( + slot=item.slot, + epoch=item.epoch, + text=item.request.normalized_text, + token_ids=tuple(item.request.token_ids), + seed=item.request.seed, + request_id=item.request.request_id, + ) + pending_slots.add(item.slot) + self._deferred_admission_count += 1 + if not progress or self._closed: + return [] + if ( + admissions + and self._deferred_admission_count < self.batch + and self._admission_signal is not None + and self._admission_signal.is_set() + and any(item is None for item in self._slots) + ): + return [] + if self._deferred_admissions: + records = list(self._deferred_admissions.values()) + engine.admit_many(records) + self._deferred_admissions.clear() + before = (engine.producer_launches, engine.codec_launches) + events = engine.advance(progress, allow_admission_yield=False) + if before != (engine.producer_launches, engine.codec_launches): + self._deferred_admission_count = 0 + return events + + def _tick(self, retirements, cancellations, admissions, progress): + if self._defer_admission_reset: + return self._tick_deferred(retirements, cancellations, admissions, progress) + engine = self.backend.inflight_engine + for slot, epoch in cancellations: + engine.cancel(slot, epoch) + if getattr(getattr(engine, "policy", None), "retire_batch", False): + engine.retire_many(retirements) + else: + for slot, epoch in retirements: + engine.retire(slot, epoch) + if admissions: + records = [ + dict( + slot=item.slot, + epoch=item.epoch, + text=item.request.normalized_text, + token_ids=item.request.token_ids, + seed=item.request.seed, + request_id=item.request.request_id, + ) + for item in admissions + ] + admit_many = getattr(engine, "admit_many", None) + if admit_many is not None: + admit_many(records) + else: + for record in records: + engine.admit(**record) + return engine.advance(progress) if progress else [] + + def _progress_slots(self): + return [ + slot + for slot, item in enumerate(self._slots) + if item is not None + and not item.aborted + and item.handle._chunks.qsize() <= self.chunk_queue_depth - 2 + ] + + async def _run(self): + loop = asyncio.get_running_loop() + try: + while not self._closed: + self._wake.clear() + cancellations = [] + for item in self._slots: + if item is not None and item.aborted: + cancellations.append((item.slot, item.epoch)) + self._record(item, "cancelled") + self._free(item) + retirements, self._retire = self._retire, [] + admissions = self._select_admissions() + progress = self._progress_slots() + if not (retirements or admissions or progress or cancellations): + await self._wake.wait() + continue + while admissions and not self._defer_admission_reset: + await loop.run_in_executor( + self.backend.executor, + self._tick, + retirements, + cancellations, + admissions, + [], + ) + retirements, cancellations, admissions = [], [], [] + if not self._closed: + await asyncio.sleep(0) + admissions = self._select_admissions() + progress = [] if self._closed else self._progress_slots() + events = await loop.run_in_executor( + self.backend.executor, + self._tick, + retirements, + cancellations, + admissions, + progress, + ) + self._consume(events) + self._flush_records() + if events: + await asyncio.sleep(0) + except Exception as error: + self._failure = error + self._closed = True + for item in tuple(self._pending.values()): + self._error(item, error) + self._record(item, "failed", error=str(error)) + self._pending.clear() + if self.on_fatal is not None: + self.on_fatal() + finally: + self._flush_records() + active = [(slot, item.epoch) for slot, item in enumerate(self._slots) if item] + if self._failure is None: + await loop.run_in_executor( + self.backend.executor, self._tick, self._retire + active, active, [], [] + ) + self._flush_records() + self._slots = [None] * self.batch + self._retire = [] + self._deferred_admissions.clear() + self._deferred_admission_count = 0 + + async def shutdown(self): + if self._shutdown_task is None: + self._shutdown_task = asyncio.create_task(self._shutdown()) + await asyncio.shield(self._shutdown_task) + + async def _shutdown(self): + self._closed = True + for item in tuple(self._pending.values()): + item.aborted = True + self._error(item, RequestAborted()) + if item.slot is not None: + self._record(item, "cancelled") + self._wake.set() + try: + if self._runner is not None: + await self._runner + finally: + self._pending.clear() + self._queue.clear() + await self.backend.shutdown() diff --git a/python/nemotron_native_backend/initial_publication.py b/python/nemotron_native_backend/initial_publication.py new file mode 100644 index 0000000..e698392 --- /dev/null +++ b/python/nemotron_native_backend/initial_publication.py @@ -0,0 +1,207 @@ +from __future__ import annotations + +import asyncio +import json +import math +import os +from collections import Counter +from pathlib import Path + +from .backend import BackendError +from .batcher import RequestAborted + +MAX_REQUEST_PACKETS = 8 +MAX_HELD_BYTES = 128 * 1024 * 1024 + + +def initial_hold_seconds(): + value = float(os.environ.get("NEMOTRON_INITIAL_AUDIO_HOLD_MS", "0")) + if not math.isfinite(value) or value not in (0.0, 40.0, 300.0): + raise ValueError("initial audio hold must be 0, 40 or 300 milliseconds") + return value / 1000.0 + + +class InitialHoldBudget: + def __init__(self, limit=MAX_HELD_BYTES, binding=None, configured_hold_ms=0): + if type(configured_hold_ms) is not int or configured_hold_ms not in (0, 40, 300): + raise ValueError("invalid initial audio hold receipt duration") + self.configured_hold_ms = configured_hold_ms + if type(limit) is not int or limit < 0: + raise ValueError("invalid initial audio staging limit") + self.limit = limit + self.binding = binding + self.used = 0 + self.peak = 0 + self.reasons = Counter() + self.delay_histogram_ms = Counter() + self.count = 0 + self.delay_sum = 0.0 + self.delay_min = None + self.delay_max = 0.0 + + def reserve(self, size): + if size < 0 or self.used + size > self.limit: + return False + self.used += size + self.peak = max(self.peak, self.used) + return True + + def release(self, size): + if not 0 <= size <= self.used: + raise RuntimeError("initial audio staging accounting drift") + self.used -= size + + def record(self, delay, reason): + self.count += 1 + self.reasons[reason] += 1 + self.delay_sum += delay + self.delay_min = delay if self.delay_min is None else min(self.delay_min, delay) + self.delay_max = max(self.delay_max, delay) + self.delay_histogram_ms[min(1000, int(delay * 1000) // 5 * 5)] += 1 + + def snapshot(self, directory): + row = { + "schema": "nemotron.initial-publication-hold.v6", + "activation_binding": self.binding, + "configured_hold_ms": self.configured_hold_ms, + "max_request_packets": MAX_REQUEST_PACKETS, + "extra_staging_limit_bytes": self.limit, + "extra_staging_current_bytes": self.used, + "extra_staging_peak_bytes": self.peak, + "release_count": self.count, + "release_reasons": dict(self.reasons), + "hold_seconds_sum": self.delay_sum, + "hold_seconds_min": self.delay_min, + "hold_seconds_max": self.delay_max, + "hold_histogram_5ms_lower_edges_1000_overflow": dict(self.delay_histogram_ms), + "scope": "cumulative process hold-release delays; excludes subsequent send time", + } + with (Path(directory) / "initial-publication-hold.jsonl").open("a") as stream: + stream.write(json.dumps(row, separators=(",", ":")) + "\n") + + +class InitialAudioHold: + def __init__(self, handle, seconds, budget): + if not 0 < handle._chunks.maxsize < MAX_REQUEST_PACKETS: + raise RuntimeError("initial hold requires an existing queue below eight packets") + self.handle = handle + self.request_id = handle.request_id + self.result = handle.result + self.seconds = seconds + self.budget = budget + self.aborted = asyncio.Event() + self.terminal_sent = False + self.terminal_started = False + self.late_abort_count = 0 + self.delivery_task = None + self.reading = None + self.first = True + self.held_bytes = 0 + + def abort(self): + if self.terminal_sent: + return False + already_aborted = self.aborted.is_set() + self.aborted.set() + if not already_aborted and self.reading is not None and not self.reading.done(): + self.reading.cancel(self) + return True + + def check(self): + if self.aborted.is_set(): + raise RequestAborted() + if self.result.cancelled(): + raise RequestAborted() + if self.result.done(): + error = self.result.exception() + if error is not None: + raise error + + def before_send(self, value): + if value["kind"] in ("audio", "audio_slab", "stream_complete_v1"): + self.check() + + def begin_send(self, value): + if value["kind"] in ("stream_complete_v1", "failed", "aborted") or ( + value["kind"] in ("audio", "audio_slab") and value.get("is_final") + ): + self.terminal_started = True + + def after_send(self, value): + if value["kind"] in ("stream_complete_v1", "failed", "aborted") or ( + value["kind"] in ("audio", "audio_slab") and value.get("is_final") + ): + self.terminal_sent = True + + def release(self): + self.budget.release(self.held_bytes) + self.held_bytes = 0 + + async def wait_owned(self, awaitable): + if self.reading is not None: + raise RuntimeError("overlapping initial hold ownership wait") + self.reading = asyncio.current_task() + try: + return await awaitable + except asyncio.CancelledError as error: + if error.args and error.args[0] is self: + if self.reading.uncancel(): + raise + raise RequestAborted() from None + raise + finally: + self.reading = None + + async def next_chunk(self): + self.release() + self.check() + chunk = await self.wait_owned(self.handle.next_chunk()) + self.check() + if not self.first: + return chunk + self.first = False + chunk.validate() + if chunk.request_id != self.request_id or chunk.sequence != 0: + raise BackendError("backend_contract", "streaming audio identity/sequence drift") + if chunk.is_final or self.result.done(): + self.budget.record(0.0, "successful_terminal") + return chunk + size = len(chunk.pcm_f32_le) + if not self.budget.reserve(size): + self.budget.record(0.0, "staging_capacity") + return chunk + self.held_bytes = size + loop = asyncio.get_running_loop() + started = loop.time() + deadline = started + self.seconds + released = loop.create_future() + + def wake(_result=None): + if not released.done(): + released.set_result(None) + + def deadline_reached(): + nonlocal timer + if loop.time() < deadline: + timer = loop.call_at(deadline, deadline_reached) + else: + wake() + + timer = loop.call_at(deadline, deadline_reached) + self.result.add_done_callback(wake) + reason = "cancelled_or_failed" + try: + await self.wait_owned(released) + self.check() + reason = "successful_terminal" if self.result.done() else "timer" + return chunk + finally: + self.budget.record(loop.time() - started, reason) + timer.cancel() + self.result.remove_done_callback(wake) + released.cancel() + + async def close(self): + if self.reading is not None: + raise RuntimeError("initial audio reader remains active during delivery cleanup") + self.release() diff --git a/python/nemotron_native_backend/lifecycle.py b/python/nemotron_native_backend/lifecycle.py new file mode 100644 index 0000000..dd42d4c --- /dev/null +++ b/python/nemotron_native_backend/lifecycle.py @@ -0,0 +1,132 @@ +from __future__ import annotations + +import json +import os +import stat +import sys +import time +from pathlib import Path +from typing import Any + +SCHEMA = "nemotron-speech.worker-lifecycle.v1" +ENVIRONMENT = "NEMOTRON_TRT_LIFECYCLE_RECEIPT" + + +class LifecycleRecorder: + def __init__(self, path: Path | None) -> None: + self.path = path + self.events: list[dict[str, Any]] = [] + self._last_monotonic_ns = 0 + self.worker_exit: int | None = None + if path is None: + return + if not path.is_absolute(): + raise RuntimeError("lifecycle receipt path must be absolute") + if path.exists() or path.is_symlink(): + raise RuntimeError("lifecycle receipt must be absent at startup") + if not path.parent.is_dir(): + raise RuntimeError("lifecycle receipt parent must exist") + for parent in reversed(path.parents): + if parent.is_symlink(): + raise RuntimeError("lifecycle receipt path has a symlink component") + + @classmethod + def from_environment(cls, *, required: bool) -> LifecycleRecorder: + value = os.environ.get(ENVIRONMENT) + if value is None: + if required: + raise RuntimeError(f"{ENVIRONMENT} must be set") + return cls(None) + return cls(Path(value)) + + @staticmethod + def sentinel_stat(path: Path) -> dict[str, Any]: + observed = os.lstat(path) + if not stat.S_ISREG(observed.st_mode): + raise RuntimeError(f"stop sentinel must be a regular file: {path}") + return { + "path": str(path), + "mode": stat.S_IMODE(observed.st_mode), + "size": observed.st_size, + "uid": observed.st_uid, + "gid": observed.st_gid, + "inode": observed.st_ino, + "device": observed.st_dev, + "mtime_ns": observed.st_mtime_ns, + } + + def record(self, phase: str, **fields: Any) -> None: + monotonic_ns = max(time.monotonic_ns(), self._last_monotonic_ns + 1) + self._last_monotonic_ns = monotonic_ns + event = {"phase": phase, "monotonic_ns": monotonic_ns, **fields} + self.events.append(event) + if phase == "worker_exit": + self.worker_exit = int(fields["exit_code"]) + print( + "nemotron_trt_lifecycle " + json.dumps(event, sort_keys=True, separators=(",", ":")), + file=sys.stderr, + flush=True, + ) + if self.path is not None: + self._write() + + def _write(self) -> None: + if self.path is None: + raise RuntimeError("lifecycle recorder has no receipt path") + payload = { + "schema": SCHEMA, + "status": "PASS" if self.worker_exit == 0 else "IN_PROGRESS", + "pid": os.getpid(), + "worker_exit": self.worker_exit, + "forced": False, + "events": self.events, + } + encoded = (json.dumps(payload, sort_keys=True, separators=(",", ":")) + "\n").encode( + "utf-8" + ) + temporary = self.path.with_name(f".{self.path.name}.{os.getpid()}.{len(self.events)}.tmp") + descriptor = os.open( + temporary, + os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, + 0o644, + ) + try: + offset = 0 + while offset < len(encoded): + offset += os.write(descriptor, encoded[offset:]) + os.fchmod(descriptor, 0o644) + os.fsync(descriptor) + finally: + os.close(descriptor) + try: + os.replace(temporary, self.path) + parent = os.open(self.path.parent, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(parent) + finally: + os.close(parent) + finally: + if temporary.exists(): + temporary.unlink() + + +_RECORDER: LifecycleRecorder | None = None + + +def install(recorder: LifecycleRecorder) -> None: + global _RECORDER + _RECORDER = recorder + + +def record(phase: str, **fields: Any) -> None: + if _RECORDER is not None: + _RECORDER.record(phase, **fields) + + +__all__ = [ + "ENVIRONMENT", + "LifecycleRecorder", + "SCHEMA", + "install", + "record", +] diff --git a/python/nemotron_native_backend/manifest.py b/python/nemotron_native_backend/manifest.py new file mode 100644 index 0000000..5bd1491 --- /dev/null +++ b/python/nemotron_native_backend/manifest.py @@ -0,0 +1,15 @@ +from __future__ import annotations + +import re +from typing import Any, Protocol, runtime_checkable + +SHA256_RE = re.compile(r"^[0-9a-f]{64}$") + + +@runtime_checkable +class WorkerManifest(Protocol): + sha256: str + raw: dict[str, Any] + enabled_buckets: tuple[int, ...] + + def worker_ready_fields(self) -> dict[str, Any]: ... diff --git a/python/nemotron_native_backend/models.py b/python/nemotron_native_backend/models.py new file mode 100644 index 0000000..225d5df --- /dev/null +++ b/python/nemotron_native_backend/models.py @@ -0,0 +1,162 @@ +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any + +SAMPLE_RATE_HZ = 22_050 +SAMPLES_PER_REQUEST = 17_640 +PCM_BYTES_PER_REQUEST = SAMPLES_PER_REQUEST * 4 +TEXT_EOS_TOKEN_ID = 131_073 + + +class RequestError(ValueError): + pass + + +_SUBMIT_FIELDS = frozenset( + { + "protocol_version", + "kind", + "request_id", + "session_id", + "normalized_text", + "token_ids", + "speaker_id", + "max_tokens", + "temperature", + "top_k", + "sample_rate_hz", + "seed", + } +) +_STRING_SUBMIT_FIELDS = ("request_id", "normalized_text", "speaker_id") +_INTEGER_SUBMIT_FIELDS = ("max_tokens", "top_k", "sample_rate_hz", "seed") + + +def _strict_submit_fields(value: dict[str, Any]) -> dict[str, Any]: + if not isinstance(value, dict) or any(not isinstance(key, str) for key in value): + raise RequestError("submit payload must be an object with string keys") + observed = set(value) + if observed != _SUBMIT_FIELDS: + raise RequestError( + f"submit fields differ: missing={sorted(_SUBMIT_FIELDS - observed)}, " + f"extra={sorted(observed - _SUBMIT_FIELDS)}" + ) + if type(value["protocol_version"]) is not int or value["protocol_version"] != 1: + raise RequestError("submit protocol_version must be integer 1") + if type(value["kind"]) is not str or value["kind"] != "submit": + raise RequestError("submit kind must be the string submit") + for name in _STRING_SUBMIT_FIELDS: + if type(value[name]) is not str: + raise RequestError(f"{name} must be a string") + session_id = value["session_id"] + if session_id is not None and type(session_id) is not str: + raise RequestError("session_id must be a string or null") + if session_id is not None and (not session_id or len(session_id.encode("utf-8")) > 1024): + raise RequestError("session_id must contain 1..1024 UTF-8 bytes when present") + token_ids = value["token_ids"] + if not isinstance(token_ids, list) or not token_ids: + raise RequestError("token_ids must be a nonempty array") + if any(type(item) is not int for item in token_ids): + raise RequestError("token_ids values must be integers") + for name in _INTEGER_SUBMIT_FIELDS: + if type(value[name]) is not int: + raise RequestError(f"{name} must be an integer") + if type(value["temperature"]) is not float: + raise RequestError("temperature must be a float") + return { + "request_id": value["request_id"], + "session_id": session_id, + "normalized_text": value["normalized_text"], + "token_ids": tuple(token_ids), + "speaker_id": value["speaker_id"], + "max_tokens": value["max_tokens"], + "temperature": value["temperature"], + "top_k": value["top_k"], + "sample_rate_hz": value["sample_rate_hz"], + "seed": value["seed"], + } + + +@dataclass(frozen=True, slots=True) +class SynthesisRequest: + request_id: str + session_id: str | None + normalized_text: str + token_ids: tuple[int, ...] + speaker_id: str + max_tokens: int + temperature: float + top_k: int + sample_rate_hz: int + seed: int + + @classmethod + def from_message(cls, value: dict[str, Any]) -> SynthesisRequest: + request = cls(**_strict_submit_fields(value)) + request.validate() + return request + + def validate(self) -> None: + if not self.request_id or len(self.request_id.encode("utf-8")) > 1024: + raise RequestError("request_id must contain 1..1024 UTF-8 bytes") + if not self.normalized_text.strip(): + raise RequestError("normalized_text must not be empty") + if len(self.normalized_text.encode("utf-8")) > 1_000_000: + raise RequestError("normalized_text is too large") + if len(self.token_ids) > 4096: + raise RequestError("token_ids exceeds 4096") + if any(item < 0 or item >= 131_075 for item in self.token_ids): + raise RequestError("token_ids contains an out-of-vocabulary ID") + if self.token_ids[-1] != TEXT_EOS_TOKEN_ID: + raise RequestError("functional v1 requires the configured text EOS token") + if self.speaker_id != "english_sample": + raise RequestError("functional v1 supports only speaker_id=english_sample") + if self.max_tokens != 16: + raise RequestError("functional v1 requires max_tokens=16 (D15 plus sentinel)") + if not math.isfinite(self.temperature) or self.temperature != 0.7: + raise RequestError("functional v1 requires temperature=0.7") + if self.top_k != 80: + raise RequestError("functional v1 requires top_k=80") + if self.sample_rate_hz != SAMPLE_RATE_HZ: + raise RequestError("functional v1 requires sample_rate_hz=22050") + if self.seed < 0 or self.seed > 0xFFFFFFFFFFFFFFFF: + raise RequestError("seed must fit uint64") + + @property + def compatibility_key(self) -> tuple[Any, ...]: + return ( + self.speaker_id, + self.max_tokens, + self.temperature, + self.top_k, + self.sample_rate_hz, + ) + + +@dataclass(frozen=True, slots=True) +class AudioResult: + request_id: str + pcm_f32_le: bytes + samples: int = SAMPLES_PER_REQUEST + sample_rate_hz: int = SAMPLE_RATE_HZ + channels: int = 1 + + def validate(self) -> None: + if type(self.request_id) is not str or not self.request_id: + raise RuntimeError("backend output request_id is empty or not a string") + if type(self.samples) is not int: + raise RuntimeError("backend output sample count is not an integer") + if self.samples != SAMPLES_PER_REQUEST: + raise RuntimeError("backend output sample count is not 17640") + if type(self.sample_rate_hz) is not int or type(self.channels) is not int: + raise RuntimeError("backend output rate and channels must be integers") + if self.sample_rate_hz != SAMPLE_RATE_HZ or self.channels != 1: + raise RuntimeError("backend output is not 22050 Hz mono") + if type(self.pcm_f32_le) is not bytes: + raise RuntimeError("backend PCM-F32LE payload must be bytes") + if len(self.pcm_f32_le) != PCM_BYTES_PER_REQUEST: + raise RuntimeError( + f"backend PCM-F32LE byte count {len(self.pcm_f32_le)} != {PCM_BYTES_PER_REQUEST}" + ) diff --git a/python/nemotron_native_backend/protocol.py b/python/nemotron_native_backend/protocol.py new file mode 100644 index 0000000..80e52cb --- /dev/null +++ b/python/nemotron_native_backend/protocol.py @@ -0,0 +1,285 @@ +from __future__ import annotations + +import asyncio +import math +import struct +from collections.abc import Mapping, Sequence +from typing import Any + +PROTOCOL_VERSION = 1 +DEFAULT_MAX_FRAME_BYTES = 64 * 1024 * 1024 +MAX_COLLECTION_ENTRIES = 16_384 +OUTPUT_CONTRACT = "pcm_f32_le_22050_mono_17640" +STREAMING_OUTPUT_CONTRACT = "pcm_f32_le_22050_mono_variable_streaming_v2" +TWO_PHASE_OUTPUT_CONTRACT = "pcm_f32_le_22050_mono_variable_streaming_v3" + +CLIENT_KINDS = frozenset({"client_hello", "submit", "abort", "ping"}) +WORKER_KINDS = frozenset( + {"worker_ready", "accepted", "audio", "stream_complete_v1", "failed", "aborted", "pong"} +) + + +class ProtocolError(ValueError): + pass + + +def _pack_length(tag8: int, tag16: int, tag32: int, length: int) -> bytes: + if length <= 0xFF: + return bytes((tag8, length)) + if length <= 0xFFFF: + return bytes((tag16,)) + struct.pack(">H", length) + if length <= 0xFFFFFFFF: + return bytes((tag32,)) + struct.pack(">I", length) + raise ProtocolError("MessagePack value is too large") + + +def packb(value: Any) -> bytes: + + output = bytearray() + + def emit(item: Any, depth: int = 0) -> None: + if depth > 64: + raise ProtocolError("MessagePack nesting exceeds 64") + if item is None: + output.append(0xC0) + elif item is False: + output.append(0xC2) + elif item is True: + output.append(0xC3) + elif isinstance(item, int): + if 0 <= item <= 0x7F: + output.append(item) + elif -32 <= item < 0: + output.append(item & 0xFF) + elif 0 <= item <= 0xFF: + output.extend((0xCC, item)) + elif 0 <= item <= 0xFFFF: + output.append(0xCD) + output.extend(struct.pack(">H", item)) + elif 0 <= item <= 0xFFFFFFFF: + output.append(0xCE) + output.extend(struct.pack(">I", item)) + elif 0 <= item <= 0xFFFFFFFFFFFFFFFF: + output.append(0xCF) + output.extend(struct.pack(">Q", item)) + elif -(1 << 7) <= item < 0: + output.append(0xD0) + output.extend(struct.pack(">b", item)) + elif -(1 << 15) <= item < 0: + output.append(0xD1) + output.extend(struct.pack(">h", item)) + elif -(1 << 31) <= item < 0: + output.append(0xD2) + output.extend(struct.pack(">i", item)) + elif -(1 << 63) <= item < 0: + output.append(0xD3) + output.extend(struct.pack(">q", item)) + else: + raise ProtocolError("integer is outside MessagePack int64/uint64") + elif isinstance(item, float): + if not math.isfinite(item): + raise ProtocolError("non-finite floats are forbidden") + output.append(0xCB) + output.extend(struct.pack(">d", item)) + elif isinstance(item, str): + encoded = item.encode("utf-8") + length = len(encoded) + if length <= 31: + output.append(0xA0 | length) + else: + output.extend(_pack_length(0xD9, 0xDA, 0xDB, length)) + output.extend(encoded) + elif isinstance(item, (bytes, bytearray, memoryview)): + encoded = bytes(item) + output.extend(_pack_length(0xC4, 0xC5, 0xC6, len(encoded))) + output.extend(encoded) + elif isinstance(item, Mapping): + length = len(item) + if length > MAX_COLLECTION_ENTRIES: + raise ProtocolError("MessagePack map is unreasonably large") + if not all(isinstance(key, str) for key in item): + raise ProtocolError("protocol map keys must be strings") + if length <= 15: + output.append(0x80 | length) + elif length <= 0xFFFF: + output.append(0xDE) + output.extend(struct.pack(">H", length)) + else: + output.append(0xDF) + output.extend(struct.pack(">I", length)) + for key, child in item.items(): + emit(key, depth + 1) + emit(child, depth + 1) + elif isinstance(item, Sequence): + length = len(item) + if length > MAX_COLLECTION_ENTRIES: + raise ProtocolError("MessagePack array is unreasonably large") + if length <= 15: + output.append(0x90 | length) + elif length <= 0xFFFF: + output.append(0xDC) + output.extend(struct.pack(">H", length)) + else: + output.append(0xDD) + output.extend(struct.pack(">I", length)) + for child in item: + emit(child, depth + 1) + else: + raise ProtocolError(f"unsupported MessagePack type: {type(item).__name__}") + + emit(value) + return bytes(output) + + +class _Unpacker: + def __init__(self, payload: bytes) -> None: + self.data = memoryview(payload) + self.offset = 0 + + def take(self, size: int) -> memoryview: + end = self.offset + size + if size < 0 or end > len(self.data): + raise ProtocolError("truncated MessagePack value") + value = self.data[self.offset : end] + self.offset = end + return value + + def number(self, fmt: str) -> int | float: + return struct.unpack(fmt, self.take(struct.calcsize(fmt)))[0] + + def value(self, depth: int = 0) -> Any: + if depth > 64: + raise ProtocolError("MessagePack nesting exceeds 64") + tag = int(self.take(1)[0]) + if tag <= 0x7F: + return tag + if tag >= 0xE0: + return tag - 256 + if 0xA0 <= tag <= 0xBF: + return self.string(tag & 0x1F) + if 0x90 <= tag <= 0x9F: + return self.array(tag & 0x0F, depth) + if 0x80 <= tag <= 0x8F: + return self.mapping(tag & 0x0F, depth) + if tag == 0xC0: + return None + if tag == 0xC2: + return False + if tag == 0xC3: + return True + if tag == 0xC4: + return bytes(self.take(int(self.number(">B")))) + if tag == 0xC5: + return bytes(self.take(int(self.number(">H")))) + if tag == 0xC6: + return bytes(self.take(int(self.number(">I")))) + if tag == 0xCB: + value = float(self.number(">d")) + if not math.isfinite(value): + raise ProtocolError("non-finite floats are forbidden") + return value + numeric = { + 0xCC: ">B", + 0xCD: ">H", + 0xCE: ">I", + 0xCF: ">Q", + 0xD0: ">b", + 0xD1: ">h", + 0xD2: ">i", + 0xD3: ">q", + } + if tag in numeric: + return int(self.number(numeric[tag])) + if tag == 0xD9: + return self.string(int(self.number(">B"))) + if tag == 0xDA: + return self.string(int(self.number(">H"))) + if tag == 0xDB: + return self.string(int(self.number(">I"))) + if tag == 0xDC: + return self.array(int(self.number(">H")), depth) + if tag == 0xDD: + return self.array(int(self.number(">I")), depth) + if tag == 0xDE: + return self.mapping(int(self.number(">H")), depth) + if tag == 0xDF: + return self.mapping(int(self.number(">I")), depth) + raise ProtocolError(f"unsupported MessagePack tag 0x{tag:02x}") + + def array(self, length: int, depth: int) -> list[Any]: + if length > MAX_COLLECTION_ENTRIES: + raise ProtocolError("MessagePack array is unreasonably large") + return [self.value(depth + 1) for _ in range(length)] + + def string(self, length: int) -> str: + try: + return bytes(self.take(length)).decode("utf-8") + except UnicodeDecodeError as error: + raise ProtocolError("invalid UTF-8 string") from error + + def mapping(self, length: int, depth: int) -> dict[str, Any]: + if length > MAX_COLLECTION_ENTRIES: + raise ProtocolError("MessagePack map is unreasonably large") + result: dict[str, Any] = {} + for _ in range(length): + key = self.value(depth + 1) + if not isinstance(key, str): + raise ProtocolError("protocol map keys must be strings") + if key in result: + raise ProtocolError(f"duplicate protocol map key {key!r}") + result[key] = self.value(depth + 1) + return result + + +def unpackb(payload: bytes) -> Any: + unpacker = _Unpacker(payload) + value = unpacker.value() + if unpacker.offset != len(unpacker.data): + raise ProtocolError("trailing bytes after MessagePack value") + return value + + +def validate_envelope(message: Any, *, from_client: bool) -> dict[str, Any]: + if not isinstance(message, dict): + raise ProtocolError("protocol frame root must be a map") + protocol_version = message.get("protocol_version") + if type(protocol_version) is not int or protocol_version != PROTOCOL_VERSION: + raise ProtocolError("unsupported protocol_version") + kind = message.get("kind") + allowed = CLIENT_KINDS if from_client else WORKER_KINDS + if kind not in allowed: + raise ProtocolError(f"unsupported {'client' if from_client else 'worker'} kind {kind!r}") + return message + + +def frame(message: Mapping[str, Any], max_frame_bytes: int = DEFAULT_MAX_FRAME_BYTES) -> bytes: + payload = packb(message) + if not payload or len(payload) > max_frame_bytes: + raise ProtocolError(f"frame payload size {len(payload)} is outside allowed range") + return struct.pack(">I", len(payload)) + payload + + +async def read_frame( + reader: asyncio.StreamReader, max_frame_bytes: int = DEFAULT_MAX_FRAME_BYTES +) -> dict[str, Any]: + header = await reader.readexactly(4) + length = struct.unpack(">I", header)[0] + if length == 0 or length > max_frame_bytes: + raise ProtocolError(f"frame length {length} is outside allowed range") + value = unpackb(await reader.readexactly(length)) + if not isinstance(value, dict): + raise ProtocolError("frame root must be a map") + return value + + +async def write_frame( + writer: asyncio.StreamWriter, + message: Mapping[str, Any], + max_frame_bytes: int = DEFAULT_MAX_FRAME_BYTES, +) -> None: + writer.write(frame(message, max_frame_bytes)) + await writer.drain() + + +def message(kind: str, **fields: Any) -> dict[str, Any]: + return {"protocol_version": PROTOCOL_VERSION, "kind": kind, **fields} diff --git a/python/nemotron_native_backend/r31.py b/python/nemotron_native_backend/r31.py new file mode 100644 index 0000000..701b991 --- /dev/null +++ b/python/nemotron_native_backend/r31.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +import gc +import json +import os +import time + +KNOWN = ( + "gc_trace", + "gc_quiet", + "native_unpack", + "staged_admission", + "eager_publication", + "fast_first_chunk", + "rng_predraw", + "accept_at_dispatch", +) + + +def _parse(): + raw = os.environ.get("NEMOTRON_R31", "{}") + value = json.loads(raw) + if ( + not isinstance(value, dict) + or set(value) - set(KNOWN) + or not all(type(v) is bool for v in value.values()) + ): + raise RuntimeError("NEMOTRON_R31 must be a JSON object of known boolean options: " + raw) + return {key: bool(value.get(key, key != "gc_trace")) for key in KNOWN} + + +OPTIONS = _parse() + + +def enabled(name): + return OPTIONS[name] + + +class GcTrace: + def __init__(self): + self.events = [] + self._open = {} + + def __call__(self, phase, info): + now = time.monotonic_ns() + generation = info.get("generation") + if phase == "start": + self._open[generation] = now + else: + self.events.append( + ( + self._open.pop(generation, now), + now, + generation, + info.get("collected", 0), + info.get("uncollectable", 0), + ) + ) + + def drain(self): + events, self.events = self.events, [] + return [ + dict(start_ns=a, end_ns=b, generation=g, collected=c, uncollectable=u, ms=(b - a) / 1e6) + for (a, b, g, c, u) in events + ] + + +TRACE = GcTrace() +_installed = False + + +def after_prepare(): + global _installed + if _installed: + return + _installed = True + if OPTIONS["gc_trace"] or OPTIONS["gc_quiet"]: + gc.callbacks.append(TRACE) + if OPTIONS["gc_quiet"]: + gc.collect() + gc.freeze() + gc.disable() + + +def cohort_boundary(): + if OPTIONS["gc_quiet"]: + gc.collect() + + +def receipt(): + return dict( + options=dict(OPTIONS), + gc_enabled=gc.isenabled(), + gc_frozen=gc.get_freeze_count(), + gc_thresholds=list(gc.get_threshold()), + gc_events=TRACE.drain(), + ) diff --git a/python/nemotron_native_backend/receipts.py b/python/nemotron_native_backend/receipts.py new file mode 100644 index 0000000..0d0b365 --- /dev/null +++ b/python/nemotron_native_backend/receipts.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +import asyncio +import json +import os +import stat +import time +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any + +from .models import SynthesisRequest + + +def cohort_id(request_ids: Sequence[str], batch: int) -> str: + if not request_ids: + raise ValueError("cohort must contain requests") + prefix = os.path.commonprefix(list(request_ids)).rstrip("-_/.") + if prefix: + return prefix + return f"b{batch}-{time.time_ns()}" + + +class ReceiptWriter: + def __init__(self, path: Path, *, manifest_sha256: str) -> None: + if not path.is_absolute() or not path.parent.is_dir(): + raise ValueError("receipt path must be absolute with an existing parent") + if path.is_symlink() or (path.exists() and not path.is_file()): + raise ValueError("receipt path must be a regular file or absent, not a symlink") + if any(component.is_symlink() for component in path.parents): + raise ValueError("receipt path must not contain a symlink component") + self.path = path + self.manifest_sha256 = manifest_sha256 + self._lock = asyncio.Lock() + + async def write_success( + self, requests: Sequence[SynthesisRequest], evidence: Mapping[str, Any] + ) -> None: + batch = len(requests) + if batch not in (1, 16, 32, 64): + raise ValueError("functional-v1 receipt requires an exact fleet batch") + ids = [request.request_id for request in requests] + payload = { + "schema_version": 1, + "cohort_id": cohort_id(ids, batch), + "public_request_ids": ids, + "real_weights": evidence["real_weights"], + "execution_contract": evidence["execution_contract"], + "actual_requests": batch, + "predictor_plan_slices": list(evidence["predictor_plan_slices"]), + "stage1_plan_slices": list(evidence["stage1_plan_slices"]), + "decode_steps": evidence["decode_steps"], + "predictor_plan_invocations": evidence["predictor_plan_invocations"], + "predictor_passes_per_lane": evidence["predictor_passes_per_lane"], + "padding_count": evidence["padding_count"], + "predictor_fallback_count": evidence["predictor_fallback_count"], + "stage1_fallback_count": evidence["stage1_fallback_count"], + "code_h2d_copies": evidence["code_h2d_copies"], + "code_d2h_copies": evidence["code_d2h_copies"], + "manifest_sha256": self.manifest_sha256, + } + encoded = (json.dumps(payload, sort_keys=True, separators=(",", ":")) + "\n").encode() + async with self._lock: + flags = ( + os.O_APPEND + | os.O_CREAT + | os.O_WRONLY + | getattr(os, "O_CLOEXEC", 0) + | getattr(os, "O_NOFOLLOW", 0) + ) + descriptor = os.open(self.path, flags, 0o640) + try: + if not stat.S_ISREG(os.fstat(descriptor).st_mode): + raise OSError("receipt destination is not a regular file") + offset = 0 + while offset < len(encoded): + written = os.write(descriptor, encoded[offset:]) + if written <= 0: + raise OSError(f"short receipt write at byte {offset}") + offset += written + os.fsync(descriptor) + finally: + os.close(descriptor) diff --git a/python/nemotron_native_backend/server.py b/python/nemotron_native_backend/server.py new file mode 100644 index 0000000..418f29a --- /dev/null +++ b/python/nemotron_native_backend/server.py @@ -0,0 +1,354 @@ +from __future__ import annotations + +import asyncio +import contextlib +import os +import stat +import uuid +from collections import deque +from pathlib import Path +from typing import Any + +from .backend import BackendError, SynthesisBackend +from .batcher import ExactBatcher, RequestAborted +from .lifecycle import record +from .manifest import SHA256_RE, WorkerManifest +from .models import SynthesisRequest +from .protocol import ( + ProtocolError, + message, + read_frame, + validate_envelope, + write_frame, +) + +COMPLETED_REQUEST_ID_WINDOW = 4096 + + +class CompletedRequestIdWindow: + def __init__(self, capacity: int = COMPLETED_REQUEST_ID_WINDOW) -> None: + if type(capacity) is not int or capacity < 1: + raise ValueError("completed request ID capacity must be a positive integer") + self._capacity = capacity + self._ordered: deque[str] = deque() + self._ids: set[str] = set() + + def __contains__(self, request_id: str) -> bool: + return request_id in self._ids + + def add(self, request_id: str) -> None: + if request_id in self._ids: + return + if len(self._ordered) == self._capacity: + self._ids.remove(self._ordered.popleft()) + self._ordered.append(request_id) + self._ids.add(request_id) + + +class WorkerServer: + def __init__( + self, + *, + socket_path: Path, + manifest: WorkerManifest, + backend: SynthesisBackend, + batcher: ExactBatcher, + max_frame_bytes: int, + ) -> None: + self.socket_path = socket_path + self.manifest = manifest + self.backend = backend + self.batcher = batcher + if type(max_frame_bytes) is not int or not (1024 <= max_frame_bytes <= 64 * 1024 * 1024): + raise ValueError("max_frame_bytes must be an integer in [1KiB,64MiB]") + self.max_frame_bytes = max_frame_bytes + self.worker_instance_id = str(uuid.uuid4()) + self.completed_request_id_capacity = COMPLETED_REQUEST_ID_WINDOW + self._server: asyncio.AbstractServer | None = None + self._client_lock = asyncio.Lock() + self._client_active = False + self._client_writer: asyncio.StreamWriter | None = None + self._stopping = False + self._listener_closed = False + self._clients_closed = False + self._socket_unlinked = False + self._delivery_tasks: set[asyncio.Task[None]] = set() + self._batcher_shutdown = False + self._shutdown_complete = False + self._shutdown_lock = asyncio.Lock() + self._socket_creation_attempted = False + + async def start(self) -> None: + if self.backend.identity.enabled_buckets != self.manifest.enabled_buckets: + raise RuntimeError( + "backend/manifest enabled bucket mismatch: " + f"{self.backend.identity.enabled_buckets} != {self.manifest.enabled_buckets}" + ) + if self.socket_path.exists() or self.socket_path.is_symlink(): + raise RuntimeError(f"refusing to replace existing socket path: {self.socket_path}") + if not self.socket_path.parent.is_dir(): + raise RuntimeError(f"socket parent does not exist: {self.socket_path.parent}") + self.batcher.start() + self._socket_creation_attempted = True + try: + self._server = await asyncio.start_unix_server( + self._handle_client, path=str(self.socket_path) + ) + os.chmod(self.socket_path, 0o660) + except BaseException as error: + try: + await self.shutdown() + except BaseException as cleanup_error: + error.add_note(f"startup rollback failed: {cleanup_error}") + raise + + async def serve_forever(self) -> None: + if self._server is None: + raise RuntimeError("server is not started") + async with self._server: + await self._server.serve_forever() + + async def _handle_client( + self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> None: + async with self._client_lock: + if self._stopping or self._client_active: + writer.close() + await writer.wait_closed() + return + self._client_active = True + self._client_writer = writer + send_lock = asyncio.Lock() + request_ids: set[str] = set() + completed_ids = CompletedRequestIdWindow(self.completed_request_id_capacity) + + async def send(value: dict[str, Any]) -> None: + async with send_lock: + await write_frame(writer, value, self.max_frame_bytes) + + try: + hello = validate_envelope( + await read_frame(reader, self.max_frame_bytes), from_client=True + ) + self._validate_hello(hello) + ready = message( + "worker_ready", + worker_instance_id=self.worker_instance_id, + **self.manifest.worker_ready_fields(), + cuda_device_name=self.backend.identity.cuda_device_name, + cuda_compute_capability=list(self.backend.identity.cuda_compute_capability), + ) + await send(ready) + while True: + incoming = validate_envelope( + await read_frame(reader, self.max_frame_bytes), from_client=True + ) + kind = incoming["kind"] + if kind == "submit": + request = SynthesisRequest.from_message(incoming) + if request.request_id in request_ids or request.request_id in completed_ids: + await send( + self._failed( + request.request_id, + BackendError( + "duplicate_request_id", + "request ID was already used on this connection", + ), + ) + ) + continue + try: + future = self.batcher.submit(request) + except BackendError as error: + await send(self._failed(request.request_id, error)) + continue + request_ids.add(request.request_id) + await send(message("accepted", request_id=request.request_id)) + task = asyncio.create_task( + self._deliver( + request.request_id, + future, + send, + request_ids, + completed_ids, + ), + name=f"deliver-{request.request_id}", + ) + self._delivery_tasks.add(task) + task.add_done_callback(self._delivery_done) + elif kind == "abort": + self._validate_exact_keys(incoming, {"request_id"}) + request_id = incoming["request_id"] + if not isinstance(request_id, str) or not self.batcher.abort(request_id): + await send( + self._failed( + str(request_id), + BackendError("unknown_request", "request is not pending"), + ) + ) + elif kind == "ping": + self._validate_exact_keys(incoming, {"nonce"}) + nonce = self._validate_ping_nonce(incoming["nonce"]) + await send( + message( + "pong", + nonce=nonce, + queue_depth=self.batcher.queue_depth, + active_batches=self.batcher.active_batches, + ) + ) + else: + raise ProtocolError(f"unexpected post-handshake kind {kind!r}") + except (asyncio.IncompleteReadError, ConnectionError, BrokenPipeError): + pass + finally: + for request_id in tuple(request_ids): + self.batcher.abort(request_id) + writer.close() + with contextlib.suppress(Exception): + await writer.wait_closed() + async with self._client_lock: + if self._client_writer is writer: + self._client_writer = None + self._client_active = False + + async def _deliver( + self, + request_id: str, + future: asyncio.Future[Any], + send: Any, + request_ids: set[str], + completed_ids: CompletedRequestIdWindow, + ) -> None: + try: + result = await future + await send( + message( + "audio", + request_id=request_id, + sequence=0, + is_final=True, + encoding="pcm_f32_le", + sample_rate_hz=result.sample_rate_hz, + channels=result.channels, + samples=result.samples, + data=result.pcm_f32_le, + ) + ) + except asyncio.CancelledError: + raise + except RequestAborted: + with contextlib.suppress(ConnectionError, BrokenPipeError): + await send(message("aborted", request_id=request_id)) + except BackendError as error: + with contextlib.suppress(ConnectionError, BrokenPipeError): + await send(self._failed(request_id, error)) + except Exception as error: + with contextlib.suppress(ConnectionError, BrokenPipeError): + await send( + self._failed( + request_id, + BackendError("backend_internal", str(error), retryable=False), + ) + ) + finally: + request_ids.discard(request_id) + completed_ids.add(request_id) + + def _validate_hello(self, value: dict[str, Any]) -> None: + self._validate_exact_keys(value, {"expected_manifest_sha256", "client_instance_id"}) + expected = value["expected_manifest_sha256"] + if not isinstance(expected, str) or SHA256_RE.fullmatch(expected) is None: + raise ProtocolError("expected_manifest_sha256 is malformed") + if expected != self.manifest.sha256: + raise ProtocolError("client/worker manifest mismatch") + client_id = value["client_instance_id"] + if type(client_id) is not str or not client_id or len(client_id.encode("utf-8")) > 1024: + raise ProtocolError("client_instance_id must contain 1..1024 UTF-8 bytes") + + @staticmethod + def _validate_ping_nonce(value: Any) -> int: + if type(value) is not int or not (0 <= value <= 0xFFFFFFFFFFFFFFFF): + raise ProtocolError("ping nonce must be a uint64 integer") + return value + + def _delivery_done(self, task: asyncio.Task[None]) -> None: + self._delivery_tasks.discard(task) + if not task.cancelled(): + task.exception() + + @staticmethod + def _validate_exact_keys(value: dict[str, Any], fields: set[str]) -> None: + expected = {"protocol_version", "kind", *fields} + if set(value) != expected: + raise ProtocolError( + f"{value.get('kind')} fields differ: missing={sorted(expected - set(value))}, " + f"extra={sorted(set(value) - expected)}" + ) + + @staticmethod + def _failed(request_id: str, error: BackendError) -> dict[str, Any]: + return message( + "failed", + request_id=request_id, + code=error.code, + message=str(error), + retryable=error.retryable, + ) + + def begin_shutdown(self) -> None: + self._stopping = True + if self._server is not None and not self._listener_closed: + self._server.close() + self._listener_closed = True + record("listener_closed") + if self._client_writer is not None: + self._client_writer.close() + + async def shutdown(self) -> None: + async with self._shutdown_lock: + if self._shutdown_complete: + return + errors: list[BaseException] = [] + self.begin_shutdown() + if self._server is not None: + try: + await self._server.wait_closed() + except BaseException as error: + errors.append(error) + if not self._clients_closed: + self._clients_closed = True + record("clients_closed") + for task in tuple(self._delivery_tasks): + task.cancel() + if self._delivery_tasks: + await asyncio.gather(*self._delivery_tasks, return_exceptions=True) + if not self._batcher_shutdown: + self._batcher_shutdown = True + try: + await self.batcher.shutdown() + except BaseException as error: + errors.append(error) + socket_stat = None + if self._socket_creation_attempted: + try: + socket_stat = self.socket_path.lstat() + except FileNotFoundError: + pass + if socket_stat is not None: + if stat.S_ISSOCK(socket_stat.st_mode): + try: + self.socket_path.unlink() + except BaseException as error: + errors.append(error) + else: + errors.append( + RuntimeError(f"refusing to unlink non-socket path: {self.socket_path}") + ) + if not self.socket_path.exists() and not self.socket_path.is_symlink(): + if not self._socket_unlinked: + self._socket_unlinked = True + record("socket_unlinked") + self._shutdown_complete = True + if errors: + raise RuntimeError("worker shutdown failed: " + "; ".join(map(str, errors))) diff --git a/python/nemotron_native_backend/streaming_batcher.py b/python/nemotron_native_backend/streaming_batcher.py new file mode 100644 index 0000000..d316644 --- /dev/null +++ b/python/nemotron_native_backend/streaming_batcher.py @@ -0,0 +1,534 @@ +from __future__ import annotations + +import asyncio +from collections import defaultdict, deque +from collections.abc import Awaitable, Callable, Sequence +from dataclasses import dataclass, field +from typing import Any, Protocol + +from .backend import BackendError, BackendIdentity +from .batcher import QueueFull, RequestAborted, exact_partitions +from .lifecycle import record +from .streaming_models import ( + StreamingAudioChunk, + StreamingResult, + StreamingSynthesisRequest, +) + +StreamChunkEmitter = Callable[[StreamingAudioChunk], Awaitable[bool]] + + +PARTIAL_HOLD_COHORT_DEADLINES = 5.0 + + +PARTIAL_IDLE_SETTLE_FRACTION = 1.0 / 20.0 + + +class StreamingSynthesisBackend(Protocol): + @property + def identity(self) -> BackendIdentity: ... + + async def stream_batch( + self, + requests: Sequence[StreamingSynthesisRequest], + emit_chunk: StreamChunkEmitter, + ) -> list[StreamingResult]: ... + + async def shutdown(self) -> None: ... + + +@dataclass(slots=True) +class StreamingHandle: + request_id: str + result: asyncio.Future[StreamingResult] + _chunks: asyncio.Queue[Any] + + async def next_chunk(self) -> StreamingAudioChunk: + value = await self._chunks.get() + if isinstance(value, BaseException): + raise value + if not isinstance(value, StreamingAudioChunk): + raise RuntimeError("streaming batcher queue ABI drift") + return value + + +@dataclass(slots=True) +class _Pending: + request: StreamingSynthesisRequest + handle: StreamingHandle + enqueued_at: float = 0.0 + state: str = "queued" + aborted: bool = False + forced_error: BaseException | None = None + next_sequence: int = 0 + chunk_count: int = 0 + total_samples: int = 0 + final_emitted: bool = False + client_terminal_enqueued: bool = False + terminal_error_enqueued: bool = False + abort_event: asyncio.Event = field(default_factory=asyncio.Event) + + +class StreamingExactBatcher: + def __init__( + self, + backend: StreamingSynthesisBackend, + *, + batch_window_ms: float, + cohort_timeout_ms: float = 1000.0, + expected_cohort_size: int | None = None, + max_queue_depth: int, + max_inflight_batches: int = 1, + chunk_queue_depth: int = 4, + cohort_hold: bool = True, + partial_hold_ms: float | None = None, + partial_idle_settle_ms: float | None = None, + ) -> None: + self.backend = backend + self.batch_window_seconds = float(batch_window_ms) / 1000.0 + self.cohort_timeout_seconds = float(cohort_timeout_ms) / 1000.0 + self.expected_cohort_size = expected_cohort_size + self.partial_hold_seconds = ( + PARTIAL_HOLD_COHORT_DEADLINES * self.cohort_timeout_seconds + if partial_hold_ms is None + else float(partial_hold_ms) / 1000.0 + ) + if self.partial_hold_seconds <= 0.0: + raise ValueError("partial_hold_ms must be positive") + self.partial_idle_settle_seconds = ( + PARTIAL_IDLE_SETTLE_FRACTION * self.cohort_timeout_seconds + if partial_idle_settle_ms is None + else float(partial_idle_settle_ms) / 1000.0 + ) + if self.partial_idle_settle_seconds < 0.0: + raise ValueError("partial_idle_settle_ms must be nonnegative") + if ( + expected_cohort_size is not None + and expected_cohort_size not in backend.identity.enabled_buckets + ): + raise ValueError("expected cohort size has no enabled exact backend plan") + self.max_queue_depth = int(max_queue_depth) + self.max_inflight_batches = int(max_inflight_batches) + if self.max_inflight_batches not in (1, 2): + raise ValueError("max_inflight_batches must be one or two") + self.chunk_queue_depth = int(chunk_queue_depth) + if self.chunk_queue_depth < 1: + raise ValueError("chunk_queue_depth must be positive") + self._queue: deque[_Pending] = deque() + self._pending: dict[str, _Pending] = {} + self._wake = asyncio.Event() + self._runner: asyncio.Task[None] | None = None + self._dispatch_tasks: set[asyncio.Task[None]] = set() + self._dispatch_slots = asyncio.Semaphore(self.max_inflight_batches) + self._closed = False + self._active_batches = 0 + + self._committed_batches = 0 + + self._last_completion_at = 0.0 + + self.cohort_hold = bool(cohort_hold) + self._holding = ( + self.cohort_hold + and self.expected_cohort_size is not None + and not self.batch_window_seconds + ) + + def start(self) -> None: + if self._runner is not None: + raise RuntimeError("streaming batcher already started") + self._last_completion_at = asyncio.get_running_loop().time() + self._runner = asyncio.create_task(self._run(), name="nemotron-streaming-exact-batcher") + + def submit(self, request: StreamingSynthesisRequest) -> StreamingHandle: + if self._closed: + raise BackendError("worker_stopping", "worker is stopping", retryable=True) + if request.request_id in self._pending: + raise BackendError("duplicate_request_id", "request ID is already pending") + if len(self._pending) >= self.max_queue_depth: + raise QueueFull() + loop = asyncio.get_running_loop() + handle = StreamingHandle( + request_id=request.request_id, + result=loop.create_future(), + _chunks=asyncio.Queue(maxsize=self.chunk_queue_depth), + ) + pending = _Pending(request=request, handle=handle, enqueued_at=loop.time()) + self._pending[request.request_id] = pending + self._queue.append(pending) + self._wake.set() + return handle + + def abort(self, request_id: str) -> bool: + pending = self._pending.get(request_id) + if pending is None: + return False + pending.aborted = True + pending.abort_event.set() + if pending.state == "queued": + self._finish_queued_error(pending, RequestAborted()) + elif pending.state in {"scheduled", "active"}: + error = RequestAborted() + pending.forced_error = error + self._force_terminal_error(pending, error) + self._wake.set() + return True + + @property + def queue_depth(self) -> int: + return sum(1 for item in self._queue if not item.aborted) + + @property + def active_batches(self) -> int: + return self._active_batches + + async def _run(self) -> None: + try: + if self._holding: + await self._run_holding() + return + while not self._closed: + await self._wake.wait() + self._wake.clear() + deadline = asyncio.get_running_loop().time() + self.cohort_timeout_seconds + if self.batch_window_seconds: + await asyncio.sleep(self.batch_window_seconds) + while self.expected_cohort_size is not None: + counts: dict[tuple[Any, ...], int] = defaultdict(int) + for item in self._queue: + if not item.aborted: + counts[item.request.compatibility_key] += 1 + if max(counts.values(), default=0) >= self.expected_cohort_size: + break + remaining = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + break + self._wake.clear() + try: + await asyncio.wait_for(self._wake.wait(), timeout=remaining) + except TimeoutError: + break + snapshot: list[_Pending] = [] + while self._queue: + item = self._queue.popleft() + if not item.aborted: + snapshot.append(item) + groups: dict[tuple[Any, ...], list[_Pending]] = defaultdict(list) + for item in snapshot: + groups[item.request.compatibility_key].append(item) + for compatible in groups.values(): + self._launch_compatible(compatible) + if self._queue: + self._wake.set() + except asyncio.CancelledError: + raise + except Exception as error: + for item in tuple(self._pending.values()): + self._finish_queued_error(item, error) + raise + + async def _run_holding(self) -> None: + loop = asyncio.get_running_loop() + while not self._closed: + self._purge_aborted() + groups = self._grouped_queue() + if not groups: + self._wake.clear() + await self._wake.wait() + continue + plan, wait_seconds = self._plan_cohorts(groups, loop.time()) + if plan: + for cohort in plan: + self._launch_cohort(cohort) + continue + self._wake.clear() + if wait_seconds is None: + await self._wake.wait() + continue + try: + await asyncio.wait_for(self._wake.wait(), timeout=wait_seconds) + except TimeoutError: + continue + + def _purge_aborted(self) -> None: + if any(item.aborted for item in self._queue): + self._queue = deque(item for item in self._queue if not item.aborted) + + def _grouped_queue(self) -> dict[tuple[Any, ...], list[_Pending]]: + groups: dict[tuple[Any, ...], list[_Pending]] = defaultdict(list) + for item in self._queue: + groups[item.request.compatibility_key].append(item) + return groups + + def _largest_bucket_within(self, depth: int) -> int | None: + buckets = [ + int(value) for value in self.backend.identity.enabled_buckets if int(value) <= depth + ] + return max(buckets) if buckets else None + + def _plan_cohorts( + self, groups: dict[tuple[Any, ...], list[_Pending]], now: float + ) -> tuple[list[list[_Pending]], float | None]: + expected = self.expected_cohort_size + assert expected is not None + plan: list[list[_Pending]] = [] + committed = self._committed_batches + waits: list[float] = [] + stalled_for = now - self._last_completion_at + backend_is_wedged = stalled_for >= self.partial_hold_seconds + + settled = stalled_for >= self.partial_idle_settle_seconds + for items in groups.values(): + overdue = now - items[0].enqueued_at - self.cohort_timeout_seconds + if len(items) < expected and overdue >= 0.0: + if self._largest_bucket_within(len(items)) is None: + self._reject_unrepresentable(items) + continue + if committed >= self.max_inflight_batches: + continue + if len(items) >= expected: + plan.append(items[:expected]) + committed += 1 + continue + if overdue < 0.0: + waits.append(-overdue) + continue + + if committed == 0 and not settled and not backend_is_wedged: + waits.append(self.partial_idle_settle_seconds - stalled_for) + continue + if (committed == 0 and settled) or backend_is_wedged: + size = self._largest_bucket_within(len(items)) + if size is None: + self._reject_unrepresentable(items) + continue + plan.append(items[:size]) + committed += 1 + continue + waits.append(self.partial_hold_seconds - stalled_for) + if plan: + return plan, None + return plan, min(waits) if waits else None + + def _reject_unrepresentable(self, items: list[_Pending]) -> None: + error = BackendError( + "unsupported_exact_batch", + f"{len(items)} request(s) cannot be represented by enabled exact plans " + f"{list(self.backend.identity.enabled_buckets)}", + retryable=True, + ) + chosen = {id(item) for item in items} + self._queue = deque(item for item in self._queue if id(item) not in chosen) + for item in items: + self._finish_queued_error(item, error) + + def _launch_cohort(self, cohort: list[_Pending]) -> None: + chosen = {id(item) for item in cohort} + self._queue = deque(item for item in self._queue if id(item) not in chosen) + for item in cohort: + item.state = "scheduled" + self._start_dispatch(cohort) + + def _start_dispatch(self, cohort: list[_Pending]) -> None: + self._committed_batches += 1 + task = asyncio.create_task( + self._dispatch(cohort), + name=f"nemotron-streaming-dispatch-b{len(cohort)}", + ) + self._dispatch_tasks.add(task) + task.add_done_callback(self._dispatch_done) + + def _launch_compatible(self, items: list[_Pending]) -> None: + partitions, remainder = exact_partitions(len(items), self.backend.identity.enabled_buckets) + cursor = 0 + for batch in partitions: + cohort = items[cursor : cursor + batch] + cursor += batch + for item in cohort: + item.state = "scheduled" + self._start_dispatch(cohort) + if remainder: + error = BackendError( + "unsupported_exact_batch", + f"{remainder} request(s) cannot be represented by enabled exact plans " + f"{list(self.backend.identity.enabled_buckets)}", + retryable=True, + ) + for item in items[cursor:]: + self._finish_queued_error(item, error) + + def _dispatch_done(self, task: asyncio.Task[None]) -> None: + self._dispatch_tasks.discard(task) + self._committed_batches -= 1 + self._last_completion_at = asyncio.get_running_loop().time() + if self._holding: + self._wake.set() + if task.cancelled(): + return + error = task.exception() + if error is not None and self._runner is not None and not self._runner.done(): + self._runner.cancel(error) + + async def _dispatch(self, cohort: list[_Pending]) -> None: + async with self._dispatch_slots: + for item in cohort: + item.state = "active" + self._active_batches += 1 + by_id = {item.request.request_id: item for item in cohort} + + async def emit(chunk: StreamingAudioChunk) -> bool: + chunk.validate() + item = by_id.get(chunk.request_id) + if item is None: + raise BackendError( + "backend_contract", "streaming chunk request ID differs from cohort" + ) + if item.final_emitted: + raise BackendError( + "backend_contract", "streaming backend emitted after final chunk" + ) + if chunk.sequence != item.next_sequence: + raise BackendError( + "backend_contract", + f"streaming sequence {chunk.sequence} != {item.next_sequence}", + ) + item.next_sequence += 1 + item.chunk_count += 1 + item.total_samples += chunk.samples + item.final_emitted = chunk.is_final + return await self._enqueue_audio_or_abort(item, chunk) + + try: + if all(item.aborted or item.forced_error is not None for item in cohort): + for item in cohort: + await self._finish_active_error(item, item.forced_error or RequestAborted()) + return + results = await self.backend.stream_batch( + tuple(item.request for item in cohort), emit + ) + if len(results) != len(cohort): + raise BackendError("backend_contract", "streaming backend result count drift") + result_by_id = {result.request_id: result for result in results} + if len(result_by_id) != len(results) or set(result_by_id) != set(by_id): + raise BackendError("backend_contract", "streaming backend result IDs differ") + for item in cohort: + result = result_by_id[item.request.request_id] + result.validate() + if not item.final_emitted: + raise BackendError( + "backend_contract", "streaming backend completed without final audio" + ) + if ( + result.chunk_count != item.chunk_count + or result.total_samples != item.total_samples + ): + raise BackendError("backend_contract", "streaming result accounting drift") + if item.aborted: + await self._finish_active_error(item, item.forced_error or RequestAborted()) + else: + self._pending.pop(item.request.request_id, None) + item.state = "done" + if not item.handle.result.done(): + item.handle.result.set_result(result) + except asyncio.CancelledError as error: + for item in cohort: + await self._finish_active_error(item, error) + raise + except Exception as error: + for item in cohort: + await self._finish_active_error(item, error) + finally: + self._active_batches -= 1 + + def _finish_queued_error(self, item: _Pending, error: BaseException) -> None: + self._pending.pop(item.request.request_id, None) + item.state = "done" + if not item.handle.result.done(): + item.handle.result.set_exception(error) + self._force_terminal_error(item, error) + + def _force_terminal_error(self, item: _Pending, error: BaseException) -> None: + if item.terminal_error_enqueued: + return + while True: + try: + item.handle._chunks.get_nowait() + except asyncio.QueueEmpty: + break + item.handle._chunks.put_nowait(error) + item.terminal_error_enqueued = True + item.client_terminal_enqueued = False + + async def _enqueue_audio_or_abort(self, item: _Pending, chunk: StreamingAudioChunk) -> bool: + if item.aborted: + return False + try: + item.handle._chunks.put_nowait(chunk) + except asyncio.QueueFull: + put_task = asyncio.create_task(item.handle._chunks.put(chunk)) + abort_task = asyncio.create_task(item.abort_event.wait()) + try: + done, _ = await asyncio.wait( + {put_task, abort_task}, + return_when=asyncio.FIRST_COMPLETED, + ) + if abort_task in done or item.aborted: + put_task.cancel() + await asyncio.gather(put_task, return_exceptions=True) + self._force_terminal_error(item, item.forced_error or RequestAborted()) + return False + abort_task.cancel() + await asyncio.gather(abort_task, return_exceptions=True) + put_task.result() + finally: + for task in (put_task, abort_task): + if not task.done(): + task.cancel() + if item.aborted: + self._force_terminal_error(item, item.forced_error or RequestAborted()) + return False + if chunk.is_final: + item.client_terminal_enqueued = True + return True + + async def _finish_active_error(self, item: _Pending, error: BaseException) -> None: + self._pending.pop(item.request.request_id, None) + item.state = "done" + if not item.handle.result.done(): + item.handle.result.set_exception(error) + self._force_terminal_error(item, error) + + async def shutdown(self) -> None: + record("batcher_drain_begin") + self._closed = True + self._wake.set() + if self._runner is not None: + self._runner.cancel() + try: + await self._runner + except asyncio.CancelledError: + pass + stopping = BackendError("worker_stopping", "worker stopped", retryable=True) + for item in tuple(self._pending.values()): + item.forced_error = stopping + item.aborted = True + item.abort_event.set() + if item.state == "queued": + self._finish_queued_error(item, stopping) + elif item.state in {"scheduled", "active"}: + self._force_terminal_error(item, stopping) + if self._dispatch_tasks: + await asyncio.gather(*tuple(self._dispatch_tasks), return_exceptions=True) + for item in tuple(self._pending.values()): + await self._finish_active_error(item, stopping) + self._pending.clear() + record("batcher_drain_end") + await self.backend.shutdown() + + +__all__ = [ + "StreamChunkEmitter", + "StreamingExactBatcher", + "StreamingHandle", + "StreamingSynthesisBackend", +] diff --git a/python/nemotron_native_backend/streaming_contract.py b/python/nemotron_native_backend/streaming_contract.py new file mode 100644 index 0000000..8e27f54 --- /dev/null +++ b/python/nemotron_native_backend/streaming_contract.py @@ -0,0 +1,475 @@ +from __future__ import annotations + +import os +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +AUDIO_CODEBOOKS = 16 +AUDIO_CODEC_TOKEN_MAX = 1023 +AUDIO_MODEL_TOKEN_MAX = 1031 +AUDIO_SENTINEL = -1 +STREAMING_MAX_TOKENS = 512 +MAX_GENERATED_RAW_ROWS = STREAMING_MAX_TOKENS - 1 +STREAMING_SPEECH_DELAY = 5 +SPEECH_DELAY_ENV = "NEMOTRON_STREAMING_SPEECH_DELAY" + + +def resolved_speech_delay() -> int: + raw = os.environ.get(SPEECH_DELAY_ENV, "") + if raw == "": + return STREAMING_SPEECH_DELAY + if raw not in {"3", "4", "5"}: + raise StreamingContractError(f"{SPEECH_DELAY_ENV} must be 3, 4, or 5; observed {raw!r}") + return int(raw) + + +MR2_MATCHED_STARTUP_CHUNKS = (6, 6) +MR2_MATCHED_STEADY_CHUNK = 8 +OPTIMIZED_H100_STARTUP_CHUNKS = (10, 14, 16) +OPTIMIZED_H100_STEADY_CHUNK = 12 +LOW_TTFA_H100_STARTUP_CHUNKS = (3, 5) +LOW_TTFA_H100_F1_STARTUP_CHUNKS = (1, 5) + + +LOW_TTFA_H100_R4_STARTUP_CHUNKS = (1, 2, 2, 3) + + +LOW_TTFA_H100_STEADY_CHUNK = 8 +STREAMING_TAP_CAPACITY = 512 +SUPPORTED_EXACT_BATCHES = (1, 16, 32, 64) +ZERO_CODEC_ROW = (0,) * AUDIO_CODEBOOKS + + +class StreamingContractError(RuntimeError): + pass + + +def _row(value: Any, *, generated: bool) -> tuple[int, ...]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise StreamingContractError("audio-code row is not a sequence") + result = tuple(int(item) for item in value) + if len(result) != AUDIO_CODEBOOKS: + raise StreamingContractError( + f"audio-code row has {len(result)} values, expected {AUDIO_CODEBOOKS}" + ) + if generated: + if any(item < 0 or item > AUDIO_MODEL_TOKEN_MAX for item in result): + raise StreamingContractError( + f"generated audio codes must be in [0,{AUDIO_MODEL_TOKEN_MAX}]" + ) + elif result != (AUDIO_SENTINEL,) * AUDIO_CODEBOOKS: + raise StreamingContractError( + "first audio-code row is not the exact [-1]*16 context sentinel" + ) + return result + + +def classify_generated_row(value: Any) -> tuple[tuple[int, ...], str]: + row = _row(value, generated=True) + if not any(row): + return row, "zero" + if max(row) > AUDIO_CODEC_TOKEN_MAX: + return row, "control" + return row, "codec" + + +@dataclass(frozen=True, slots=True) +class ObservedAudioRow: + generated_index: int + codes: tuple[int, ...] + classification: str + + +class CumulativeAudioCodeCursor: + def __init__(self, *, speech_delay: int = STREAMING_SPEECH_DELAY) -> None: + self.speech_delay = int(speech_delay) + if self.speech_delay < 0: + raise ValueError("speech_delay must be nonnegative") + self._rows: tuple[tuple[int, ...], ...] = () + + @property + def observed_rows(self) -> int: + return len(self._rows) + + @property + def generated_rows(self) -> int: + return max(0, len(self._rows) - 1) + + def consume( + self, rows: Sequence[Sequence[int]] + ) -> tuple[tuple[ObservedAudioRow, ...], tuple[tuple[int, ...], ...]]: + if not isinstance(rows, Sequence) or isinstance(rows, (str, bytes)): + raise StreamingContractError("audio-code snapshot is not a sequence") + if not rows: + raise StreamingContractError("audio-code snapshot is empty") + normalized = (_row(rows[0], generated=False),) + tuple( + _row(value, generated=True) for value in rows[1:] + ) + if len(normalized) < len(self._rows): + raise StreamingContractError( + f"cumulative audio-code count regressed from {len(self._rows)} to {len(normalized)}" + ) + if normalized[: len(self._rows)] != self._rows: + raise StreamingContractError( + "cumulative audio codes changed an already observed prefix" + ) + first = max(len(self._rows), 1) + self._rows = normalized + observed: list[ObservedAudioRow] = [] + codec: list[tuple[int, ...]] = [] + for row_index in range(first, len(normalized)): + codes, classification = classify_generated_row(normalized[row_index]) + generated_index = row_index - 1 + observed.append(ObservedAudioRow(generated_index, codes, classification)) + if classification == "codec" and generated_index >= self.speech_delay: + codec.append(codes) + return tuple(observed), tuple(codec) + + def finish(self) -> int: + if not self._rows: + raise StreamingContractError("stream ended without the context sentinel") + if len(self._rows) == 1: + raise StreamingContractError("stream ended without a generated audio-code row") + return len(self._rows) - 1 + + +@dataclass(frozen=True, slots=True) +class ChunkSchedule: + profile: str + startup: tuple[int, ...] + steady: int + + def validate(self) -> None: + if not self.profile or not self.startup: + raise ValueError("chunk schedule profile/startup must be nonempty") + if any(not 1 <= value <= 16 for value in (*self.startup, self.steady)): + raise ValueError("chunk schedule frames must be in [1,16]") + + +MR2_MATCHED_CHUNK_SCHEDULE = ChunkSchedule( + "mr2_matched", MR2_MATCHED_STARTUP_CHUNKS, MR2_MATCHED_STEADY_CHUNK +) +OPTIMIZED_H100_CHUNK_SCHEDULE = ChunkSchedule( + "optimized_h100", + OPTIMIZED_H100_STARTUP_CHUNKS, + OPTIMIZED_H100_STEADY_CHUNK, +) +LOW_TTFA_H100_CHUNK_SCHEDULE = ChunkSchedule( + "low_ttfa_h100", + LOW_TTFA_H100_STARTUP_CHUNKS, + LOW_TTFA_H100_STEADY_CHUNK, +) +LOW_TTFA_H100_F1_CHUNK_SCHEDULE = ChunkSchedule( + "low_ttfa_h100_f1", + LOW_TTFA_H100_F1_STARTUP_CHUNKS, + LOW_TTFA_H100_STEADY_CHUNK, +) +LOW_TTFA_H100_R4_CHUNK_SCHEDULE = ChunkSchedule( + "low_ttfa_h100_r4", + LOW_TTFA_H100_R4_STARTUP_CHUNKS, + LOW_TTFA_H100_STEADY_CHUNK, +) +CHUNK_SCHEDULES = { + value.profile: value + for value in ( + MR2_MATCHED_CHUNK_SCHEDULE, + OPTIMIZED_H100_CHUNK_SCHEDULE, + LOW_TTFA_H100_CHUNK_SCHEDULE, + LOW_TTFA_H100_F1_CHUNK_SCHEDULE, + LOW_TTFA_H100_R4_CHUNK_SCHEDULE, + ) +} + + +SERVICE_TO_STAGE1_CHUNK_PROFILE = { + "mr2_matched": "mr2_matched", + "optimized_h100": "optimized", + "low_ttfa_h100": "low_ttfa", + "low_ttfa_h100_f1": "low_ttfa_f1", + "low_ttfa_h100_r4": "low_ttfa_r4", +} + + +def target_chunk_frames( + chunk_index: int, schedule: ChunkSchedule = MR2_MATCHED_CHUNK_SCHEDULE +) -> int: + value = int(chunk_index) + if value < 0: + raise ValueError("chunk index must be nonnegative") + schedule.validate() + if value < len(schedule.startup): + return schedule.startup[value] + return schedule.steady + + +@dataclass(frozen=True, slots=True) +class CohortCodeChunk: + chunk_index: int + model_frames: int + codes: tuple[tuple[tuple[int, ...], ...], ...] + true_lengths: tuple[int, ...] + final_by_lane: tuple[bool, ...] + + +class ExactCohortChunker: + def __init__( + self, + request_ids: Sequence[str], + *, + schedule: ChunkSchedule = MR2_MATCHED_CHUNK_SCHEDULE, + ) -> None: + self.request_ids = tuple(str(value) for value in request_ids) + if len(self.request_ids) not in SUPPORTED_EXACT_BATCHES: + raise ValueError(f"unsupported exact streaming cohort B{len(self.request_ids)}") + if any(not value for value in self.request_ids): + raise ValueError("request IDs must be nonempty") + if len(set(self.request_ids)) != len(self.request_ids): + raise ValueError("request IDs must be unique") + self._lane = {value: index for index, value in enumerate(self.request_ids)} + schedule.validate() + self.schedule = schedule + self._pending: list[list[tuple[int, ...]]] = [[] for _ in self.request_ids] + self._ended = [False for _ in self.request_ids] + self._final_reported = [False for _ in self.request_ids] + self._chunk_index = 0 + + @property + def chunk_index(self) -> int: + return self._chunk_index + + @property + def complete(self) -> bool: + return all(self._ended) and not any(self._pending) and all(self._final_reported) + + def take_terminal_lanes(self) -> tuple[int, ...]: + result: list[int] = [] + for lane, (ended, pending, reported) in enumerate( + zip(self._ended, self._pending, self._final_reported, strict=True) + ): + if ended and not pending and not reported: + self._final_reported[lane] = True + result.append(lane) + return tuple(result) + + def push( + self, + request_id: str, + frames: Sequence[Sequence[int]] = (), + *, + ended: bool = False, + ) -> None: + try: + lane = self._lane[request_id] + except KeyError as error: + raise StreamingContractError( + f"request {request_id!r} is not in this exact cohort" + ) from error + if self._ended[lane]: + raise StreamingContractError( + f"request {request_id!r} received data after terminal state" + ) + for value in frames: + row, classification = classify_generated_row(value) + if classification != "codec": + raise StreamingContractError( + "cohort chunker accepts only already-filtered codec rows" + ) + self._pending[lane].append(row) + if ended: + self._ended[lane] = True + + def pop_ready(self) -> CohortCodeChunk | None: + if self.complete: + return None + target = target_chunk_frames(self._chunk_index, self.schedule) + if any( + len(pending) < target and not ended + for pending, ended in zip(self._pending, self._ended, strict=True) + ): + return None + all_short_terminal = all(self._ended) and all( + len(pending) < target for pending in self._pending + ) + model_frames = ( + max((len(pending) for pending in self._pending), default=0) + if all_short_terminal + else target + ) + if model_frames <= 0: + return None + codes: list[tuple[tuple[int, ...], ...]] = [] + true_lengths: list[int] = [] + final_by_lane: list[bool] = [] + for lane, (pending, ended) in enumerate(zip(self._pending, self._ended, strict=True)): + true_length = min(len(pending), model_frames) + prefix = tuple(pending[:true_length]) + del pending[:true_length] + codes.append(prefix + (ZERO_CODEC_ROW,) * (model_frames - true_length)) + true_lengths.append(true_length) + final = ended and not pending and not self._final_reported[lane] + final_by_lane.append(final) + if final: + self._final_reported[lane] = True + result = CohortCodeChunk( + chunk_index=self._chunk_index, + model_frames=model_frames, + codes=tuple(codes), + true_lengths=tuple(true_lengths), + final_by_lane=tuple(final_by_lane), + ) + self._chunk_index += 1 + return result + + +@dataclass(frozen=True, slots=True) +class PublishedAudioChunk: + request_id: str + sequence: int + payload: Any + samples: int + is_final: bool + + +class OneChunkLookaheadPublisher: + def __init__( + self, request_ids: Sequence[str], *, eager: bool = False, completion_marker: bool = False + ) -> None: + self.request_ids = tuple(str(value) for value in request_ids) + if not self.request_ids or len(set(self.request_ids)) != len(self.request_ids): + raise ValueError("publisher request IDs must be nonempty and unique") + self._lane = {value: index for index, value in enumerate(self.request_ids)} + self._pending: list[tuple[Any, int] | None] = [None] * len(self.request_ids) + self._sequence = [0] * len(self.request_ids) + self._terminal = [False] * len(self.request_ids) + self.eager = bool(eager) + self.completion_marker = bool(completion_marker) + if self.completion_marker and not self.eager: + raise ValueError("completion markers exist only for eager publication") + + self._published_any = [False] * len(self.request_ids) + + def _index(self, request_id: str) -> int: + try: + return self._lane[request_id] + except KeyError as error: + raise StreamingContractError( + f"request {request_id!r} is not in this publisher" + ) from error + + def _publish(self, lane: int, *, final: bool) -> PublishedAudioChunk: + pending = self._pending[lane] + if pending is None: + raise StreamingContractError("stream ended without a nonempty decoded audio chunk") + payload, samples = pending + self._pending[lane] = None + result = PublishedAudioChunk( + request_id=self.request_ids[lane], + sequence=self._sequence[lane], + payload=payload, + samples=samples, + is_final=final, + ) + self._sequence[lane] += 1 + self._published_any[lane] = True + if final: + self._terminal[lane] = True + return result + + def decoded( + self, + request_id: str, + payload: Any, + samples: int, + *, + terminal: bool = False, + continuation_expected: bool = True, + ) -> tuple[PublishedAudioChunk, ...]: + lane = self._index(request_id) + if self._terminal[lane]: + raise StreamingContractError("decoded audio arrived after final publication") + count = int(samples) + if count <= 0: + raise StreamingContractError("decoded audio chunk must contain samples") + emitted: list[PublishedAudioChunk] = [] + if self._pending[lane] is not None: + emitted.append(self._publish(lane, final=False)) + self._pending[lane] = (payload, count) + if terminal: + emitted.append(self._publish(lane, final=True)) + elif self.eager and continuation_expected: + emitted.append(self._publish(lane, final=False)) + return tuple(emitted) + + def observe_codec_continuation(self, request_id: str) -> tuple[PublishedAudioChunk, ...]: + lane = self._index(request_id) + if self._terminal[lane] or self._pending[lane] is None: + return () + return (self._publish(lane, final=False),) + + def finish(self, request_id: str) -> PublishedAudioChunk: + lane = self._index(request_id) + if self._terminal[lane]: + raise StreamingContractError("request was already published final") + if self._pending[lane] is None and self.eager and self._published_any[lane]: + if self.completion_marker: + marker = PublishedAudioChunk( + request_id=self.request_ids[lane], + sequence=self._sequence[lane], + payload=b"", + samples=0, + is_final=True, + ) + self._sequence[lane] += 1 + self._terminal[lane] = True + return marker + raise StreamingContractError( + "eager lane drained before its final chunk: continuation was promised " + "for a chunk that turned out to be last" + ) + return self._publish(lane, final=True) + + +def minimum_streaming_kv_tokens( + *, + batch: int, + prompt_tokens: int, + max_new_tokens: int = STREAMING_MAX_TOKENS, +) -> int: + values = (int(batch), int(prompt_tokens), int(max_new_tokens)) + if any(value <= 0 for value in values): + raise ValueError("batch, prompt_tokens, and max_new_tokens must be positive") + return values[0] * (values[1] + values[2]) + + +__all__ = [ + "AUDIO_CODEBOOKS", + "AUDIO_CODEC_TOKEN_MAX", + "AUDIO_MODEL_TOKEN_MAX", + "AUDIO_SENTINEL", + "CHUNK_SCHEDULES", + "ChunkSchedule", + "CohortCodeChunk", + "CumulativeAudioCodeCursor", + "ExactCohortChunker", + "MR2_MATCHED_CHUNK_SCHEDULE", + "MR2_MATCHED_STARTUP_CHUNKS", + "MR2_MATCHED_STEADY_CHUNK", + "OPTIMIZED_H100_CHUNK_SCHEDULE", + "OPTIMIZED_H100_STARTUP_CHUNKS", + "OPTIMIZED_H100_STEADY_CHUNK", + "ObservedAudioRow", + "OneChunkLookaheadPublisher", + "PublishedAudioChunk", + "STREAMING_MAX_TOKENS", + "STREAMING_SPEECH_DELAY", + "SPEECH_DELAY_ENV", + "STREAMING_TAP_CAPACITY", + "SUPPORTED_EXACT_BATCHES", + "StreamingContractError", + "ZERO_CODEC_ROW", + "classify_generated_row", + "minimum_streaming_kv_tokens", + "resolved_speech_delay", + "target_chunk_frames", +] diff --git a/python/nemotron_native_backend/streaming_models.py b/python/nemotron_native_backend/streaming_models.py new file mode 100644 index 0000000..ce21446 --- /dev/null +++ b/python/nemotron_native_backend/streaming_models.py @@ -0,0 +1,130 @@ +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any + +from .models import ( + SAMPLE_RATE_HZ, + TEXT_EOS_TOKEN_ID, + RequestError, + _strict_submit_fields, +) +from .streaming_contract import STREAMING_MAX_TOKENS + +STREAMING_OUTPUT_CONTRACT = "pcm_f32_le_22050_mono_variable_streaming_v2" +STREAMING_FLEET_EXECUTION_CONTRACT = ( + "trtllm_stage0_natural512_trt_predictor_fleet_gpu_bridge_trt_stage1_stateful_fleet_v2" +) + + +@dataclass(frozen=True, slots=True) +class StreamingSynthesisRequest: + request_id: str + session_id: str | None + normalized_text: str + token_ids: tuple[int, ...] + speaker_id: str + max_tokens: int + temperature: float + top_k: int + sample_rate_hz: int + seed: int + + @classmethod + def from_message(cls, value: dict[str, Any]) -> StreamingSynthesisRequest: + request = cls(**_strict_submit_fields(value)) + request.validate() + return request + + def validate(self) -> None: + if not self.request_id or len(self.request_id.encode("utf-8")) > 1024: + raise RequestError("request_id must contain 1..1024 UTF-8 bytes") + if not self.normalized_text.strip(): + raise RequestError("normalized_text must not be empty") + if len(self.normalized_text.encode("utf-8")) > 1_000_000: + raise RequestError("normalized_text is too large") + if len(self.token_ids) > 4096: + raise RequestError("token_ids exceeds 4096") + if any(item < 0 or item >= 131_075 for item in self.token_ids): + raise RequestError("token_ids contains an out-of-vocabulary ID") + if self.token_ids[-1] != TEXT_EOS_TOKEN_ID: + raise RequestError("streaming v2 requires the configured text EOS token") + if self.speaker_id != "english_sample": + raise RequestError("streaming v2 supports only speaker_id=english_sample") + if self.max_tokens != STREAMING_MAX_TOKENS: + raise RequestError("streaming v2 requires max_tokens=512") + if not math.isfinite(self.temperature) or self.temperature != 0.7: + raise RequestError("streaming v2 requires temperature=0.7") + if self.top_k != 80: + raise RequestError("streaming v2 requires top_k=80") + if self.sample_rate_hz != SAMPLE_RATE_HZ: + raise RequestError("streaming v2 requires sample_rate_hz=22050") + if self.seed < 0 or self.seed > 0xFFFFFFFFFFFFFFFF: + raise RequestError("seed must fit uint64") + + @property + def compatibility_key(self) -> tuple[Any, ...]: + return ( + self.speaker_id, + self.max_tokens, + self.temperature, + self.top_k, + self.sample_rate_hz, + ) + + +@dataclass(frozen=True, slots=True) +class StreamingAudioChunk: + request_id: str + sequence: int + pcm_f32_le: bytes + samples: int + is_final: bool + sample_rate_hz: int = SAMPLE_RATE_HZ + channels: int = 1 + + def validate(self) -> None: + if type(self.request_id) is not str or not self.request_id: + raise RuntimeError("streaming audio chunk request_id is empty or not a string") + if type(self.sequence) is not int or self.sequence < 0: + raise RuntimeError("streaming audio chunk sequence is negative") + + if type(self.samples) is not int or self.samples < 0: + raise RuntimeError("streaming audio chunk sample count is negative") + if self.samples == 0 and not self.is_final: + raise RuntimeError("streaming audio chunk must be nonempty") + if type(self.is_final) is not bool: + raise RuntimeError("streaming audio chunk final flag is not a boolean") + if type(self.sample_rate_hz) is not int or type(self.channels) is not int: + raise RuntimeError("streaming audio chunk rate and channels must be integers") + if self.sample_rate_hz != SAMPLE_RATE_HZ or self.channels != 1: + raise RuntimeError("streaming audio chunk is not 22050 Hz mono") + if type(self.pcm_f32_le) is not bytes: + raise RuntimeError("streaming PCM-F32LE payload must be bytes") + if len(self.pcm_f32_le) != self.samples * 4: + raise RuntimeError("streaming PCM-F32LE byte count differs from samples") + + +@dataclass(frozen=True, slots=True) +class StreamingResult: + request_id: str + chunk_count: int + total_samples: int + + def validate(self) -> None: + if type(self.request_id) is not str or not self.request_id: + raise RuntimeError("streaming result request_id is empty or not a string") + if type(self.chunk_count) is not int or type(self.total_samples) is not int: + raise RuntimeError("streaming result counts must be integers") + if self.chunk_count <= 0 or self.total_samples <= 0: + raise RuntimeError("streaming result must contain nonempty audio") + + +__all__ = [ + "STREAMING_OUTPUT_CONTRACT", + "STREAMING_FLEET_EXECUTION_CONTRACT", + "StreamingAudioChunk", + "StreamingResult", + "StreamingSynthesisRequest", +] diff --git a/python/nemotron_native_backend/streaming_server.py b/python/nemotron_native_backend/streaming_server.py new file mode 100644 index 0000000..6145313 --- /dev/null +++ b/python/nemotron_native_backend/streaming_server.py @@ -0,0 +1,305 @@ +from __future__ import annotations + +import asyncio +import contextlib +from pathlib import Path +from typing import Any + +from . import r31 +from .backend import BackendError +from .batcher import RequestAborted +from .manifest import WorkerManifest +from .protocol import ( + TWO_PHASE_OUTPUT_CONTRACT, + ProtocolError, + message, + read_frame, + validate_envelope, + write_frame, +) +from .server import CompletedRequestIdWindow, WorkerServer +from .streaming_batcher import StreamingExactBatcher, StreamingHandle +from .streaming_models import ( + STREAMING_FLEET_EXECUTION_CONTRACT, + STREAMING_OUTPUT_CONTRACT, + StreamingSynthesisRequest, +) + + +class StreamingWorkerServer(WorkerServer): + def __init__( + self, + *, + socket_path: Path, + manifest: WorkerManifest, + backend: Any, + batcher: StreamingExactBatcher, + max_frame_bytes: int, + ) -> None: + super().__init__( + socket_path=socket_path, + manifest=manifest, + backend=backend, + batcher=batcher, + max_frame_bytes=max_frame_bytes, + ) + if manifest.raw.get("manifest_status") != ("functional_v2_fleet_streaming_frozen"): + raise RuntimeError("streaming worker requires the frozen v2 manifest") + if manifest.worker_ready_fields()["output_contract"] not in ( + STREAMING_OUTPUT_CONTRACT, + TWO_PHASE_OUTPUT_CONTRACT, + ): + raise RuntimeError("streaming worker output contract drift") + identity = backend.identity + if ( + not identity.real_weights + or identity.execution_contract != STREAMING_FLEET_EXECUTION_CONTRACT + ): + raise RuntimeError( + "streaming fleet manifest requires real weights and the v2 streaming contract" + ) + + async def _handle_client( + self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> None: + async with self._client_lock: + if self._stopping or self._client_active: + writer.close() + await writer.wait_closed() + return + self._client_active = True + self._client_writer = writer + send_lock = asyncio.Lock() + request_ids: set[str] = set() + completed_ids = CompletedRequestIdWindow(self.completed_request_id_capacity) + + async def send(value: dict[str, Any]) -> None: + async with send_lock: + await write_frame(writer, value, self.max_frame_bytes) + + try: + hello = validate_envelope( + await read_frame(reader, self.max_frame_bytes), from_client=True + ) + self._validate_hello(hello) + await send( + message( + "worker_ready", + worker_instance_id=self.worker_instance_id, + **self.manifest.worker_ready_fields(), + cuda_device_name=self.backend.identity.cuda_device_name, + cuda_compute_capability=list(self.backend.identity.cuda_compute_capability), + ) + ) + while True: + incoming = validate_envelope( + await read_frame(reader, self.max_frame_bytes), from_client=True + ) + kind = incoming["kind"] + if kind == "submit": + request = StreamingSynthesisRequest.from_message(incoming) + if request.request_id in request_ids or request.request_id in completed_ids: + await send( + self._failed( + request.request_id, + BackendError( + "duplicate_request_id", + "request ID was already used on this connection", + ), + ) + ) + continue + try: + handle = self.batcher.submit(request) + except BackendError as error: + await send(self._failed(request.request_id, error)) + continue + request_ids.add(request.request_id) + lookup = getattr(self.batcher, "r31_dispatch_event", None) + dispatched = lookup(request.request_id) if lookup is not None else None + if dispatched is None: + await send(message("accepted", request_id=request.request_id)) + task = asyncio.create_task( + self._deliver_stream(handle, send, request_ids, completed_ids, dispatched), + name=f"deliver-stream-{request.request_id}", + ) + self._delivery_tasks.add(task) + task.add_done_callback(self._delivery_done) + elif kind == "abort": + self._validate_exact_keys(incoming, {"request_id"}) + request_id = incoming["request_id"] + if not isinstance(request_id, str) or not self.batcher.abort(request_id): + await send( + self._failed( + str(request_id), + BackendError("unknown_request", "request is not pending"), + ) + ) + elif kind == "ping": + self._validate_exact_keys(incoming, {"nonce"}) + nonce = self._validate_ping_nonce(incoming["nonce"]) + await send( + message( + "pong", + nonce=nonce, + queue_depth=self.batcher.queue_depth, + active_batches=self.batcher.active_batches, + ) + ) + else: + raise ProtocolError(f"unexpected post-handshake kind {kind!r}") + except (asyncio.IncompleteReadError, ConnectionError, BrokenPipeError): + pass + finally: + for request_id in tuple(request_ids): + self.batcher.abort(request_id) + writer.close() + with contextlib.suppress(Exception): + await writer.wait_closed() + async with self._client_lock: + if self._client_writer is writer: + self._client_writer = None + self._client_active = False + + async def _deliver_stream( + self, + handle: StreamingHandle, + send: Any, + request_ids: set[str], + completed_ids: CompletedRequestIdWindow, + dispatched: asyncio.Event | None = None, + ) -> None: + request_id = handle.request_id + handle.result.add_done_callback(self._drain_result_exception) + completed = False + chunk_count = 0 + total_samples = 0 + try: + if dispatched is not None: + waiter = asyncio.ensure_future(dispatched.wait()) + try: + await asyncio.wait((waiter, handle.result), return_when=asyncio.FIRST_COMPLETED) + finally: + waiter.cancel() + if not dispatched.is_set(): + await handle.result + raise BackendError( + "backend_contract", "request finished before its cohort launched" + ) + await send(message("accepted", request_id=request_id)) + output_contract = self.manifest.worker_ready_fields()["output_contract"] + if output_contract not in (STREAMING_OUTPUT_CONTRACT, TWO_PHASE_OUTPUT_CONTRACT): + raise BackendError("backend_contract", "unsupported streaming output contract") + two_phase = output_contract == TWO_PHASE_OUTPUT_CONTRACT + marker_allowed = two_phase and r31.enabled("eager_publication") + final_chunk = None + audio_records = 0 + while True: + chunk = await handle.next_chunk() + chunk.validate() + marker = marker_allowed and chunk.is_final and chunk.samples == 0 + if two_phase and chunk.samples == 0 and not marker: + raise BackendError("backend_contract", "empty PCM is not a completion message") + if chunk.request_id != request_id or chunk.sequence != chunk_count: + raise BackendError( + "backend_contract", "streaming audio identity/sequence drift" + ) + chunk_count += 1 + total_samples += chunk.samples + if chunk.is_final: + final_chunk = chunk + if not two_phase or marker: + break + await send( + message( + "audio", + request_id=request_id, + sequence=chunk.sequence, + is_final=False, + encoding="pcm_f32_le", + sample_rate_hz=chunk.sample_rate_hz, + channels=chunk.channels, + samples=chunk.samples, + data=chunk.pcm_f32_le, + ) + ) + audio_records += 1 + if chunk.is_final: + break + result = await handle.result + result.validate() + if ( + final_chunk is None + or result.request_id != request_id + or result.chunk_count != chunk_count + or result.total_samples != total_samples + ): + raise BackendError( + "backend_contract", + "streaming final/result accounting drift", + ) + if two_phase: + await send( + message( + "stream_complete_v1", + request_id=request_id, + next_sequence=audio_records, + chunk_count=audio_records, + total_samples=total_samples, + ) + ) + else: + await send( + message( + "audio", + request_id=request_id, + sequence=final_chunk.sequence, + is_final=True, + encoding="pcm_f32_le", + sample_rate_hz=final_chunk.sample_rate_hz, + channels=final_chunk.channels, + samples=final_chunk.samples, + data=final_chunk.pcm_f32_le, + ) + ) + completed = True + except StopAsyncIteration: + with contextlib.suppress(ConnectionError, BrokenPipeError): + await send( + self._failed( + request_id, + BackendError( + "backend_contract", + "stream ended without a final audio chunk", + ), + ) + ) + except asyncio.CancelledError: + raise + except RequestAborted: + with contextlib.suppress(ConnectionError, BrokenPipeError): + await send(message("aborted", request_id=request_id)) + except BackendError as error: + with contextlib.suppress(ConnectionError, BrokenPipeError): + await send(self._failed(request_id, error)) + except Exception as error: + with contextlib.suppress(ConnectionError, BrokenPipeError): + await send( + self._failed( + request_id, + BackendError("backend_internal", str(error), retryable=False), + ) + ) + finally: + if not completed: + self.batcher.abort(request_id) + request_ids.discard(request_id) + completed_ids.add(request_id) + + @staticmethod + def _drain_result_exception(future: asyncio.Future[Any]) -> None: + if not future.cancelled(): + future.exception() + + +__all__ = ["StreamingWorkerServer"] diff --git a/python/nemotron_native_backend/worker.py b/python/nemotron_native_backend/worker.py new file mode 100644 index 0000000..0e45da0 --- /dev/null +++ b/python/nemotron_native_backend/worker.py @@ -0,0 +1,632 @@ +from __future__ import annotations + +import asyncio +import hashlib +import json +import os +import threading +import time +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from types import SimpleNamespace + +from . import protocol, r31 +from .backend import BackendError, BackendIdentity +from .bundle import BundleError +from .execution.composition import create_for_qualification +from .execution.engine import _CONSTRUCTION_LOCK +from .hardware import validate_native_device +from .host.cpu_pools import configure, snapshot +from .host.diagnostic import Trace +from .host.exact_admission import install as install_exact_admission +from .host.serialization import install as install_serialization +from .host.transport import SlabServerMixin +from .server import WorkerServer +from .streaming_batcher import StreamingExactBatcher +from .streaming_contract import OneChunkLookaheadPublisher +from .streaming_models import StreamingAudioChunk, StreamingResult +from .streaming_server import StreamingWorkerServer + +CONTRACT = "native_full_prefill_full_stage0_full_stage1_streaming_v1" + + +class NativeManifest: + def __init__(self, plan): + self.plan = plan + self.sha256 = plan.manifest_sha256 + self.enabled_buckets = tuple(sorted({1, plan.physical_width})) + compatibility = plan.contract["transport_compatibility"] + self.raw = dict( + execution_contract=compatibility["execution_contract"], + model_sha256=plan.contract["model"]["checkpoint_sha256"], + output_contract=compatibility["output_contract"], + installed_transport_versions=dict(compatibility["installed_transport_versions"]), + ) + if self.raw["execution_contract"] != CONTRACT: + raise BundleError("Native backend contract mismatch") + + def worker_ready_fields(self): + return dict( + manifest_sha256=self.sha256, + model_sha256=self.raw["model_sha256"], + enabled_buckets=list(self.enabled_buckets), + output_contract=self.raw["output_contract"], + **self.raw["installed_transport_versions"], + ) + + +class NativeServer(StreamingWorkerServer): + def __init__(self, **kwargs): + + if kwargs["backend"].identity.execution_contract != CONTRACT: + raise RuntimeError("native backend contract mismatch") + WorkerServer.__init__(self, **kwargs) + + +class TimedBatcher(StreamingExactBatcher): + def submit(self, request): + observed = time.monotonic_ns() + handle = super().submit(request) + self.backend.arrivals[request.request_id] = observed + return handle + + +class BaseNativeBackend: + def __init__(self, manifest, batch, receipt_dir, *, engine_factory, delivery_group): + import torch + + hardware = validate_native_device(torch.cuda, receipt_dir) + self.engine_factory = engine_factory + self.delivery_group = delivery_group + self.manifest = manifest + self.batch = batch + self.receipt_dir = Path(receipt_dir) + self.receipt_dir.mkdir(parents=True, exist_ok=True) + self.executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="native-cohort") + self._shutdown_task = None + self.engines = {} + self.epoch = 0 + self.arrivals = {} + self.owner = asyncio.Lock() + self._identity = BackendIdentity( + hardware["detected"]["name"], + tuple(hardware["detected"]["compute_capability"]), + manifest.enabled_buckets, + True, + CONTRACT, + ) + + @property + def identity(self): + return self._identity + + async def prepare(self): + await asyncio.get_running_loop().run_in_executor(self.executor, self._prepare) + + async def shutdown(self): + if self._shutdown_task is None: + self._shutdown_task = asyncio.create_task(self._shutdown()) + await asyncio.shield(self._shutdown_task) + + async def _shutdown(self): + def close(): + error = None + inflight = getattr(self, "inflight_engine", None) + if inflight is not None: + try: + inflight.close() + except BaseException as exc: + error = exc + self.engines.pop(self.batch, None) + for engine in self.engines.values(): + try: + engine.close() + except BaseException as exc: + if error is None: + error = exc + if error is not None: + raise error + + try: + await asyncio.get_running_loop().run_in_executor(self.executor, close) + finally: + self.executor.shutdown(wait=True) + + async def _stream_batch_impl(self, requests, emit_chunk): + group_size = self.delivery_group + deferred_audit = False + async with self.owner: + self.epoch += 1 + epoch = self.epoch + batch = len(requests) + audit_enabled = deferred_audit and batch == 512 + if any(r.seed != requests[0].seed for r in requests): + raise BackendError("seed_contract", "native cohort requires a shared base seed") + loop = asyncio.get_running_loop() + request_ids = [r.request_id for r in requests] + lanes = {request_id: lane for (lane, request_id) in enumerate(request_ids)} + if len(lanes) != batch: + raise RuntimeError("request identities must be unique") + eager = r31.enabled("eager_publication") + if ( + eager + and self.manifest.raw["output_contract"] + != "pcm_f32_le_22050_mono_variable_streaming_v3" + ): + raise BundleError("Eager publication requires the streaming_v3 completion message") + publisher = OneChunkLookaheadPublisher( + request_ids, eager=eager, completion_marker=eager + ) + counts = [0] * batch + samples = [0] * batch + first_public = [None] * batch + start = time.perf_counter() + start_ns = time.monotonic_ns() + captured = {} + handoffs = [] + abandoned = threading.Event() + pending_future = [] + published_since_yield = 0 + fairness_yields = 0 + + async def publish(chunk): + nonlocal published_since_yield, fairness_yields + lane = lanes[chunk.request_id] + audio = StreamingAudioChunk( + chunk.request_id, chunk.sequence, chunk.payload, chunk.samples, chunk.is_final + ) + await emit_chunk(audio) + counts[lane] += 1 + samples[lane] += audio.samples + if first_public[lane] is None: + first_public[lane] = (time.perf_counter() - start) * 1000 + if audit_enabled: + audit = getattr(self.engines[batch], "_r11_audit", None) + if audit is None: + raise RuntimeError("deferred audit requires installed engine") + await audit.published(lane, chunk.payload) + published_since_yield += 1 + if group_size and published_since_yield == group_size: + published_since_yield = 0 + fairness_yields += 1 + await asyncio.sleep(0) + + async def consume(events): + for event in events: + kind = event["type"] + if kind == "chunk": + chunks = publisher.decoded( + request_ids[event["lane"]], event["pcm_f32_le"], event["samples"] + ) + for chunk in chunks: + await publish(chunk) + elif kind == "continuation": + for lane in event["lanes"]: + for chunk in publisher.observe_codec_continuation(request_ids[lane]): + await publish(chunk) + elif kind == "finish": + await publish(publisher.finish(request_ids[event["lane"]])) + elif kind == "receipt": + captured.update(event["receipt"]) + else: + raise RuntimeError("unexpected native event: " + kind) + + def handoff(events): + if abandoned.is_set(): + raise RuntimeError("cohort delivery abandoned") + begin = time.monotonic_ns() + future = asyncio.run_coroutine_threadsafe(consume(events), loop) + pending_future.append(future) + try: + if abandoned.is_set(): + future.cancel() + future.result() + finally: + pending_future.remove(future) + handoffs.append( + { + "kind": events[0]["type"], + "events": len(events), + "start_ns": begin, + "end_ns": time.monotonic_ns(), + } + ) + + def produce(): + if batch not in self.engines: + self.engines[batch] = self.engine_factory(batch) + if audit_enabled: + raise BundleError("Deferred audit is not part of the imported host contract") + events = self.engines[batch].stream( + [r.normalized_text for r in requests], + token_ids=[r.token_ids for r in requests], + speaker=requests[0].speaker_id, + seed=requests[0].seed, + epoch=epoch, + ) + buffered = [] + try: + for event in events: + if event["type"] == "chunk": + buffered.append(event) + if len(buffered) > batch: + raise RuntimeError("bounded callback batch exceeded") + if event.get("batch_end", True): + handoff(tuple(buffered)) + buffered.clear() + else: + if buffered: + raise RuntimeError("chunk group missing terminal marker") + handoff((event,)) + finally: + events.close() + + producer = loop.run_in_executor(self.executor, produce) + try: + await asyncio.shield(producer) + except BaseException as error: + abandoned.set() + for future in tuple(pending_future): + future.cancel() + while not producer.done(): + try: + await asyncio.shield(producer) + except asyncio.CancelledError: + continue + except BaseException: + break + try: + producer.result() + except BaseException: + pass + (self.receipt_dir / f"cohort-{epoch:05d}-FAILED.json").write_text( + json.dumps(dict(epoch=epoch, batch=batch, error=repr(error)), indent=2) + "\n" + ) + raise + captured.update( + execution_contract=CONTRACT, + epoch=epoch, + batch=batch, + backend_started_monotonic_ns=start_ns, + backend_wall_ms=(time.perf_counter() - start) * 1000, + first_public_enqueue_ms=first_public, + r8_handoffs=handoffs, + r11_delivery=dict( + group_size=group_size, + fairness_yields=fairness_yields, + deferred_audit=audit_enabled, + ), + requests=[ + dict( + request_id=r.request_id, + normalized_text=r.normalized_text, + normalized_text_sha256=hashlib.sha256( + r.normalized_text.encode() + ).hexdigest(), + token_ids=list(r.token_ids), + seed=r.seed, + emitted_samples=samples[i], + submit_monotonic_ns=self.arrivals.pop(r.request_id, None), + published_chunks=counts[i], + ) + for (i, r) in enumerate(requests) + ], + ) + captured["r31"] = r31.receipt() + (self.receipt_dir / f"cohort-{epoch:05d}.json").write_text( + json.dumps(captured, indent=2) + "\n" + ) + r31.cohort_boundary() + return [ + StreamingResult(r.request_id, counts[i], samples[i]) + for (i, r) in enumerate(requests) + ] + + +class NativeBackend(BaseNativeBackend): + async def stream_batch(self, requests, emit_chunk): + if self.epoch >= 1 and len(requests) != self.batch: + arrival_values = list(self.arrivals.values()) + row = dict( + status="REJECTED_NONWIDE_DISPATCH_BEFORE_NEURAL_EXECUTION", + expected=self.batch, + selected=len(requests), + epoch=self.epoch, + outstanding_arrivals=len(arrival_values), + observed_monotonic_ns=time.monotonic_ns(), + oldest_submit_ns=min(arrival_values) if arrival_values else None, + latest_submit_ns=max(arrival_values) if arrival_values else None, + startup_request_only_allows_b1=True, + ) + path = self.receipt_dir / "route-rejection.json" + if not path.exists(): + path.write_text(json.dumps(row, indent=2) + "\n") + raise RuntimeError( + "Native wide HTTP refuses post-readiness width " + str(len(requests)) + ) + if self.epoch == 1 and len(requests) == self.batch: + values = [self.arrivals[r.request_id] for r in requests] + row = dict( + status="FIRST_EXCLUDED_WIDE_WAVE_ADMITTED", + physical_batch=self.batch, + request_count=len(requests), + arrival_span_ms=(max(values) - min(values)) / 1000000.0, + latest_to_dispatch_ms=(time.monotonic_ns() - max(values)) / 1000000.0, + cohort_timeout_ms=10000, + readiness_timeout_ms=1000, + epoch_before_dispatch=self.epoch, + ) + (self.receipt_dir / "first-wide-dispatch.json").write_text( + json.dumps(row, indent=2) + "\n" + ) + return await self._stream_batch_impl(requests, emit_chunk) + + def _prepare(self): + mode = self.manifest.plan.contract["host"].get("batching_mode", "fixed_cohort") + widths = {1, self.batch} if mode == "fixed_cohort" else {self.batch} + for width in sorted(widths): + engine = self.engine_factory(width) + engine.retain_codes = False + self.engines[width] = engine + events = engine.stream( + ["Hello from the native three megakernel runner."] * width, seed=0, epoch=1 + ) + counts = {"chunk": 0, "continuation": 0, "finish": 0, "receipt": 0} + try: + for event in events: + counts[event["type"]] += 1 + finally: + events.close() + if counts["finish"] != width or counts["receipt"] != 1: + raise RuntimeError("Wide HTTP startup failed publication lifecycle") + receipt = dict( + status="PASS_EXECUTION_INTEGRITY", + batch=width, + events=counts, + native=engine.last_receipt, + source_contract=self.manifest.sha256, + scope="Prepared before HTTP; fresh width-specific execution qualification " + "required by launch seal", + cpu_pools=snapshot(self), + ) + (self.receipt_dir / f"http-preparation-b{width}.json").write_text( + json.dumps(receipt, indent=2) + "\n" + ) + if mode != "fixed_cohort": + from .execution.inflight_engine import InflightEngine + + if mode not in ("fixed_request", "inflight"): + raise BundleError("Unsupported request batching mode") + if not r31.enabled("eager_publication"): + raise BundleError( + "In-flight execution requires the retained eager publication path" + ) + engine = self.engines[self.batch] + library = engine.native_artifacts.get("stage0_inflight") + if library is None: + raise BundleError("Missing qualified source-bound in-flight producer") + self.inflight_engine = InflightEngine.from_engine( + engine, + stage0_library=library, + policy=self.manifest.plan.contract["host"].get("inflight_policy"), + ) + (self.receipt_dir / "inflight-binding.json").write_text( + json.dumps( + dict( + mode=mode, + physical_batch=self.batch, + request_stable_rng=True, + contraction=False, + stage0_library_sha256=hashlib.sha256(library.read_bytes()).hexdigest(), + ), + indent=2, + ) + + "\n" + ) + + +class NativeSlabServer(SlabServerMixin, NativeServer): + def __init__(self, *, host_options, slab_audit_enabled=False, trace=None, **kwargs): + if dict(host_options) != {"f5": True}: + raise BundleError("Only the selected f5 host transport is ported") + self.native_host_options = dict(host_options) + self.native_dma = None + self.slab_audit_enabled = bool(slab_audit_enabled) + self.native_trace = trace if trace is not None else Trace() + super().__init__(**kwargs) + + +@dataclass(frozen=True) +class ReadinessBinding: + factory: Callable[[], object] + source_sha256: str + physical_width: int = 1 + + +@dataclass(frozen=True) +class QualificationWorker: + backend: NativeBackend + batcher: object + server: NativeSlabServer + readiness_source_sha256: str + + +async def prepare_for_qualification( + plan, + artifact_root, + model_root, + *, + source_pins, + model_identity, + readiness: ReadinessBinding, + socket_path, + receipts, + qualification_only, + slab_audit_enabled=False, + trace=None, + engine_factory_override=None, +): + if qualification_only is not True or plan.contract["enabled"] is not False: + raise BundleError("Production native package serving is not enabled") + if ( + readiness.physical_width != 1 + or len(readiness.source_sha256) != 64 + or any(c not in "0123456789abcdef" for c in readiness.source_sha256) + ): + raise BundleError("A separately pinned physical-B1 readiness factory is required") + from .execution.readiness import ReadinessFactory + + if type(readiness.factory) is not ReadinessFactory: + raise BundleError("Readiness requires the package source-verified factory") + readiness.factory.verify_identity(readiness.source_sha256) + host = plan.contract["host"] + max_active_requests = None + configured_cap = os.environ.get("NEMOTRON_MAX_ACTIVE_REQUESTS") + if configured_cap is not None: + try: + max_active_requests = json.loads(configured_cap) + except (ValueError, TypeError): + raise BundleError("NEMOTRON_MAX_ACTIVE_REQUESTS requires a JSON integer") from None + if ( + host.get("batching_mode", "fixed_cohort") != "inflight" + or type(max_active_requests) is not int + or not 1 <= max_active_requests <= plan.physical_width + ): + raise BundleError( + "NEMOTRON_MAX_ACTIVE_REQUESTS requires inflight mode and integer 1..physical width" + ) + if ( + host["backend_executor_threads"] != 1 + or host["backend_delivery_group"] != 64 + or host["deferred_audit"] + ): + raise BundleError("Changed imported worker ownership or delivery contract") + configure( + intraop_threads=host["torch_intra_threads"], interop_threads=host["torch_interop_threads"] + ) + install_serialization(protocol) + + def engine_factory(width): + if width == 1: + with _CONSTRUCTION_LOCK: + previous_grid = os.environ.get("NEMOTRON_STAGE0_MEGAKERNEL_GRID") + try: + value = readiness.factory() + finally: + if previous_grid is None: + os.environ.pop("NEMOTRON_STAGE0_MEGAKERNEL_GRID", None) + else: + os.environ["NEMOTRON_STAGE0_MEGAKERNEL_GRID"] = previous_grid + if value.batch != 1 or value.steps != 511 or value.active_lanes_config: + value.close() + raise BundleError("Readiness factory returned an incompatible native engine") + value.retain_codes = False + return value + if width != plan.physical_width: + raise BundleError("Unregistered physical neural width") + return create_for_qualification( + plan, + artifact_root, + model_root, + source_pins=source_pins, + model_identity=model_identity, + qualification_only=True, + ) + + if engine_factory_override is not None: + from .execution.program import ProgramFactory + + if type(engine_factory_override) is not ProgramFactory: + raise BundleError("Serving requires the source-verified program factory") + engine_factory_override.verify(plan, readiness) + engine_factory = engine_factory_override + + manifest = NativeManifest(plan) + backend = NativeBackend( + manifest, + plan.physical_width, + receipts, + engine_factory=engine_factory, + delivery_group=host["backend_delivery_group"], + ) + try: + await backend.prepare() + (backend.receipt_dir / "cpu-pools.json").write_text( + json.dumps(snapshot(backend, asyncio.get_running_loop()), indent=2) + "\n" + ) + mode = host.get("batching_mode", "fixed_cohort") + if mode == "fixed_cohort": + admission = SimpleNamespace(TimedBatcher=TimedBatcher, BackendError=BackendError) + if plan.physical_width > 1: + install_exact_admission( + admission, + incomplete_timeout_seconds=plan.contract["publication"][ + "incomplete_deadline_ms" + ] + / 1000.0, + ) + batcher = admission.TimedBatcher( + backend, + batch_window_ms=0, + cohort_timeout_ms=1000, + expected_cohort_size=plan.physical_width, + max_queue_depth=4 * plan.physical_width, + max_inflight_batches=2, + chunk_queue_depth=4, + cohort_hold=True, + ) + else: + from .inflight_batcher import InflightBatcher + + if max_active_requests is None: + batcher = InflightBatcher(backend, mode=mode) + else: + batcher = InflightBatcher( + backend, mode=mode, max_active_requests=max_active_requests + ) + with (backend.receipt_dir / "admission-cap.json").open("x") as stream: + json.dump( + dict( + schema="nemotron.active-admission-cap.binding.v1", + status="CONFIGURED_NOT_STARTED", + configuration="NEMOTRON_MAX_ACTIVE_REQUESTS", + mode=mode, + physical_batch=backend.batch, + max_active_requests=batcher.max_active_requests, + physical_slot_capacity=len(batcher._slots), + physical_epoch_capacity=len(batcher._epochs), + max_queue_depth=batcher.max_queue_depth, + chunk_queue_depth=batcher.chunk_queue_depth, + scope="engine_owners_including_paused_and_deferred", + excludes_retired_delivery_owners=True, + queue_wait_in_request_latency=True, + factory="nemotron_native_backend.worker.prepare_for_qualification", + backend=f"{type(backend).__module__}.{type(backend).__qualname__}", + controller=f"{type(batcher).__module__}.{type(batcher).__qualname__}", + ), + stream, + indent=2, + ) + stream.write("\n") + server = NativeSlabServer( + socket_path=Path(socket_path), + manifest=manifest, + backend=backend, + batcher=batcher, + max_frame_bytes=67108864, + host_options=host["transport_options"], + slab_audit_enabled=slab_audit_enabled, + trace=trace, + ) + if mode != "fixed_cohort": + batcher.on_fatal = server.begin_shutdown + r31.after_prepare() + (backend.receipt_dir / "r31-options.json").write_text( + json.dumps(r31.receipt(), indent=2) + "\n" + ) + return QualificationWorker(backend, batcher, server, readiness.source_sha256) + except BaseException: + await backend.shutdown() + raise diff --git a/scripts/check_license_headers.py b/scripts/check_license_headers.py index 1291484..29df458 100644 --- a/scripts/check_license_headers.py +++ b/scripts/check_license_headers.py @@ -40,6 +40,18 @@ "src/asr/encoder/rel_pos_attention.cpp", "src/asr/encoder/rel_pos_attention.h", } +IMPORTED_SOURCE_PREFIXES = ( + "deploy/nemotron_native/kernels/", + "native/tts-normalizer/", + "python/nemotron_native_backend/", +) +IMPORTED_SOURCE_FILES = { + "deploy/nemotron_native/convert_model.py", + "deploy/nemotron_native/serve.py", + "scripts/nemotron_native/build_kernels.py", + "scripts/nemotron_native/build_package.py", + "scripts/nemotron_native/runtime_dependencies.py", +} def repository_root() -> Path: @@ -74,6 +86,9 @@ def tracked_source_files(root: Path) -> list[Path]: def is_source_file(path: Path) -> bool: + relative = path.as_posix() + if relative.startswith(IMPORTED_SOURCE_PREFIXES) or relative in IMPORTED_SOURCE_FILES: + return False name = path.name return ( path.suffix.lower() in SOURCE_SUFFIXES diff --git a/scripts/nemotron_benchmark.py b/scripts/nemotron_benchmark.py new file mode 100644 index 0000000..b66f71d --- /dev/null +++ b/scripts/nemotron_benchmark.py @@ -0,0 +1,537 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import hashlib +import http.client +import json +import math +import os +import platform +import statistics +import tempfile +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timezone +from pathlib import Path +from urllib.parse import urlsplit + + +def distribution(values): + values = sorted(values) + return { + "count": len(values), + "avg": statistics.fmean(values) if values else None, + **{ + f"p{percent}": ( + values[max(0, math.ceil(len(values) * percent / 100) - 1)] if values else None + ) + for percent in (50, 95, 99) + }, + } + + +def reconstruct(arrivals, frame_samples, startup, steady, sample_rate): + chunks, inter_chunk, inter_frame = [], [], [] + completed = 0 + previous_arrival = None + buffered = 0.0 + deficits = [] + gaps = [] + prior_samples = 0 + for arrival in arrivals: + stamp, cumulative = arrival["elapsed_ms"], arrival["cumulative_samples"] + if cumulative == prior_samples: + continue + if previous_arrival is not None: + gap = (stamp - previous_arrival) / 1000 + gaps.append(gap * 1000) + if gap > buffered: + deficits.append((gap - buffered) * 1000) + buffered = max(0.0, buffered - gap) + buffered += (cumulative - prior_samples) / sample_rate + previous_arrival, prior_samples = stamp, cumulative + while True: + frames = startup[len(chunks)] if len(chunks) < len(startup) else steady + boundary = completed + frames * frame_samples + if cumulative < boundary: + break + if chunks: + interval = stamp - chunks[-1]["elapsed_ms"] + inter_chunk.append(interval) + inter_frame.append(interval / frames) + chunks.append({"elapsed_ms": stamp, "frames": frames, "cumulative_samples": boundary}) + completed = boundary + return { + "logical_chunk_completions": chunks, + "inter_chunk_ms": inter_chunk, + "inter_frame_ms": inter_frame, + "network_inter_arrival_ms": gaps, + "terminal_partial_samples": prior_samples - completed, + "playback_deficit_events": len(deficits), + "max_playback_deficit_ms": max(deficits, default=0.0), + "total_playback_deficit_ms": sum(deficits), + } + + +def request(args, case, index, trial_start, gate=None): + result = { + "request_index": index, + "case_id": case["id"], + "seed": args.seed, + "status": "FAILED", + "http_status": None, + "error": None, + "ttfb_ms": None, + "ttfa_ms": None, + "headers_ms": None, + "completion_ms": None, + "connect_ms": None, + "bytes": 0, + "samples": 0, + "audio_seconds": 0.0, + "arrivals": [], + "sample_rate": args.sample_rate, + "sample_rate_source": "configured_fallback", + } + connection = response = transport = None + digest = hashlib.sha256() + started = time.perf_counter() + deadline = started + args.timeout + measured = started + result["worker_start_ms"] = (started - trial_start) * 1000 + body = { + "model": args.model, + "input": case["text"], + "voice": args.voice, + "language_code": args.language, + "seed": args.seed, + "sample_rate": args.sample_rate, + "response_format": "pcm", + "stream_format": "audio", + } + if args.apply_normalize is not None: + body["apply_normalize"] = args.apply_normalize == "true" + + def remaining(): + seconds = deadline - time.perf_counter() + if seconds <= 0: + raise TimeoutError("request exceeded timeout") + if transport is not None and transport.fileno() >= 0: + transport.settimeout(seconds) + return seconds + + def observe(data): + if not data: + return + stamp = (time.perf_counter() - measured) * 1000 + digest.update(data) + result["bytes"] += len(data) + result["samples"] = result["bytes"] // 2 + result["arrivals"].append( + { + "elapsed_ms": stamp, + "bytes": len(data), + "cumulative_bytes": result["bytes"], + "cumulative_samples": result["samples"], + } + ) + if result["ttfb_ms"] is None: + result["ttfb_ms"] = stamp + if result["ttfa_ms"] is None and result["samples"]: + result["ttfa_ms"] = stamp + + try: + if gate is not None: + gate.wait(timeout=args.timeout) + started = measured = time.perf_counter() + deadline = started + args.timeout + result["start_ms"] = (started - trial_start) * 1000 + parsed = urlsplit(args.url) + connection_class = ( + http.client.HTTPSConnection if parsed.scheme == "https" else http.client.HTTPConnection + ) + connection = connection_class(parsed.hostname, parsed.port, timeout=remaining()) + connection.connect() + transport = connection.sock + result["connect_ms"] = (time.perf_counter() - started) * 1000 + if args.exclude_connect: + measured = time.perf_counter() + result["stopwatch_start_ms"] = (measured - trial_start) * 1000 + headers = {"Content-Type": "application/json", "Accept": "audio/pcm", "Connection": "close"} + if args.api_key_env: + headers["Authorization"] = "Bearer " + os.environ[args.api_key_env] + remaining() + connection.request("POST", parsed.path or "/", json.dumps(body), headers) + remaining() + response = connection.getresponse() + result["http_status"] = response.status + result["headers_ms"] = (time.perf_counter() - measured) * 1000 + if response.status != 200: + remaining() + detail = response.read(8192).decode("utf-8", errors="replace") + raise RuntimeError(f"HTTP {response.status}: {detail}") + if response.getheader("Content-Type", "").split(";", 1)[0].strip().lower() != "audio/pcm": + raise RuntimeError("response Content-Type is not audio/pcm") + rate = response.getheader("x-audio-sample-rate") + if rate is not None: + if int(rate) != args.sample_rate: + raise RuntimeError( + f"server sample rate {rate} differs from requested {args.sample_rate}" + ) + result["sample_rate_source"] = "response_header" + result["transfer_encoding"] = response.getheader("Transfer-Encoding") + result["content_length"] = response.getheader("Content-Length") + while True: + remaining() + try: + data = response.read1(args.read_bytes) + except http.client.IncompleteRead as error: + observe(error.partial) + raise + if not data: + break + observe(data) + if result["bytes"] == 0 or result["bytes"] % 2: + raise RuntimeError("response contains empty or sample-misaligned PCM16") + if ( + result["content_length"] is not None + and int(result["content_length"]) != result["bytes"] + ): + raise RuntimeError("response ended before Content-Length bytes arrived") + result["status"] = "COMPLETE" + except Exception as error: + result["error"] = {"type": type(error).__name__, "message": str(error)} + finally: + result["completion_ms"] = (time.perf_counter() - measured) * 1000 + result["finish_ms"] = (time.perf_counter() - trial_start) * 1000 + if response is not None: + response.close() + if connection is not None: + connection.close() + result["audio_seconds"] = result["samples"] / args.sample_rate + result["pcm_sha256"] = digest.hexdigest() + result.update( + reconstruct( + result["arrivals"], + args.codec_frame_samples, + args.startup_chunk_frames, + args.steady_chunk_frames, + args.sample_rate, + ) + ) + return result + + +def trial_summary(records, wall_seconds): + complete = [record for record in records if record["status"] == "COMPLETE"] + audio_seconds = sum(record["audio_seconds"] for record in records) + complete_seconds = sum(record["audio_seconds"] for record in complete) + return { + "requests": len(records), + "completed": len(complete), + "failures": len(records) - len(complete), + "completion_fraction": len(complete) / len(records) if records else 0, + "wall_seconds_including_drain": wall_seconds, + "received_audio_seconds": audio_seconds, + "completed_audio_seconds": complete_seconds, + "rtfx": audio_seconds / wall_seconds if wall_seconds else 0, + "completed_rtfx": complete_seconds / wall_seconds if wall_seconds else 0, + "requests_per_second": len(records) / wall_seconds if wall_seconds else 0, + "completed_requests_per_second": len(complete) / wall_seconds if wall_seconds else 0, + "ttfa_ms": distribution([record["ttfa_ms"] for record in complete]), + "observed_ttfa_ms": distribution( + [record["ttfa_ms"] for record in records if record["ttfa_ms"] is not None] + ), + "completion_ms": distribution([record["completion_ms"] for record in complete]), + "observed_request_elapsed_ms": distribution( + [record["completion_ms"] for record in records] + ), + "inter_frame_ms": distribution( + [gap for record in complete for gap in record["inter_frame_ms"]] + ), + "inter_chunk_ms": distribution( + [gap for record in complete for gap in record["inter_chunk_ms"]] + ), + "playback_deficit_responses": sum( + record["playback_deficit_events"] > 0 for record in records + ), + "playback_deficit_events": sum(record["playback_deficit_events"] for record in records), + "max_playback_deficit_ms": max( + (record["max_playback_deficit_ms"] for record in records), default=0 + ), + } + + +def run_trial(args, corpus, count, mode, executor, output): + started = time.perf_counter() + output.update(started_utc=datetime.now(timezone.utc).isoformat(), mode=mode, requests=[]) + futures = [] + try: + if mode == "fixed-cohort": + for begin in range(0, count, args.concurrency): + indices = range(begin, min(count, begin + args.concurrency)) + gate = threading.Barrier(len(indices)) + wave = [] + for index in indices: + future = executor.submit( + request, args, corpus[index % len(corpus)], index, started, gate + ) + futures.append(future) + wave.append(future) + for future in wave: + future.result() + else: + for index in range(count): + futures.append( + executor.submit(request, args, corpus[index % len(corpus)], index, started) + ) + for future in futures: + future.result() + finally: + for future in futures: + output["requests"].append(future.result()) + output["summary"] = trial_summary(output["requests"], time.perf_counter() - started) + output["planned_requests"] = count + output["not_submitted"] = count - len(futures) + + +def summarize(trials): + summaries = [trial["summary"] for trial in trials if "summary" in trial] + result = { + "trials": len(summaries), + "aggregation": "equal mean of trial statistics; quantiles are not pooled", + } + if not summaries: + return result + for key in ( + "requests", + "completed", + "failures", + "playback_deficit_responses", + "playback_deficit_events", + ): + result[key] = sum(value[key] for value in summaries) + for key in ( + "rtfx", + "completed_rtfx", + "requests_per_second", + "completed_requests_per_second", + "completion_fraction", + "max_playback_deficit_ms", + ): + result[key] = statistics.fmean(value[key] for value in summaries) + for key in ( + "ttfa_ms", + "observed_ttfa_ms", + "completion_ms", + "observed_request_elapsed_ms", + "inter_frame_ms", + "inter_chunk_ms", + ): + present = [value[key] for value in summaries if value[key]["count"]] + result[key] = { + "contributing_trials": len(present), + "observations": sum(value["count"] for value in present), + **{ + stat: statistics.fmean(value[stat] for value in present) if present else None + for stat in ("avg", "p50", "p95", "p99") + }, + } + return result + + +def save(path, report): + descriptor, temporary = tempfile.mkstemp(prefix=path.name + ".", dir=path.parent) + try: + with os.fdopen(descriptor, "w") as stream: + json.dump(report, stream, indent=2, allow_nan=False) + stream.write("\n") + os.replace(temporary, path) + finally: + if os.path.exists(temporary): + os.unlink(temporary) + + +def parse_args(argv=None): + parser = argparse.ArgumentParser( + description="Portable HTTP PCM benchmark; timings differ from historical clients" + ) + parser.add_argument("--url", default="http://127.0.0.1:8080/v1/audio/speech") + parser.add_argument("--mode", choices=("closed-loop", "fixed-cohort"), default="closed-loop") + parser.add_argument( + "--width", + type=int, + help="declared physical backend width, recorded but not inferred from HTTP", + ) + parser.add_argument("--concurrency", type=int, default=32) + parser.add_argument("--num-requests", type=int, default=128, help="measured requests per trial") + parser.add_argument("--warmup-waves", type=int, default=2) + parser.add_argument("--trials", type=int, default=3) + corpus = parser.add_mutually_exclusive_group() + corpus.add_argument("--corpus-tsv", type=Path) + corpus.add_argument( + "--text", default="The quick brown fox jumps over the lazy dog near the riverbank." + ) + parser.add_argument("--corpus-sha256") + parser.add_argument("--model", default="nemotron-tts") + parser.add_argument("--voice", default="english_sample") + parser.add_argument("--language", default="en-US") + parser.add_argument("--apply-normalize", choices=("true", "false")) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--sample-rate", type=int, default=22050) + parser.add_argument("--timeout", type=float, default=300) + parser.add_argument("--read-bytes", type=int, default=65536) + parser.add_argument("--codec-frame-samples", type=int, default=1764) + parser.add_argument("--startup-chunk-frames", default="2,3,4,5") + parser.add_argument("--steady-chunk-frames", type=int, default=8) + parser.add_argument( + "--exclude-connect", + action="store_true", + help="exclude TCP/TLS connect from request stopwatch, not trial wall time", + ) + parser.add_argument( + "--api-key-env", + help="read a bearer token from this environment variable; token is not recorded", + ) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args(argv) + try: + args.startup_chunk_frames = [int(value) for value in args.startup_chunk_frames.split(",")] + except ValueError: + parser.error("startup-chunk-frames must be comma-separated positive integers") + for key in ( + "concurrency", + "num_requests", + "trials", + "sample_rate", + "read_bytes", + "codec_frame_samples", + "steady_chunk_frames", + ): + if getattr(args, key) <= 0: + parser.error(key.replace("_", "-") + " must be positive") + if args.width is not None and args.width <= 0: + parser.error("width must be positive") + if args.warmup_waves < 0 or any(value <= 0 for value in args.startup_chunk_frames): + parser.error("warmup-waves must be nonnegative and chunk frame counts positive") + if not math.isfinite(args.timeout) or args.timeout <= 0 or not 0 <= args.seed <= 2147483647: + parser.error("timeout must be positive and seed must be in 0..2147483647") + if args.num_requests < args.concurrency: + parser.error("num-requests must be at least concurrency") + parsed = urlsplit(args.url) + if ( + parsed.scheme not in ("http", "https") + or not parsed.hostname + or parsed.username + or parsed.password + or parsed.query + or parsed.fragment + ): + parser.error("url must be an HTTP(S) endpoint without credentials, query, or fragment") + if args.api_key_env and not os.environ.get(args.api_key_env): + parser.error("api-key-env names a missing or empty environment variable") + if args.corpus_sha256 and args.corpus_tsv is None: + parser.error("corpus-sha256 requires corpus-tsv") + return args + + +def main(argv=None): + args = parse_args(argv) + if args.corpus_tsv: + raw = args.corpus_tsv.read_bytes() + digest = hashlib.sha256(raw).hexdigest() + if args.corpus_sha256 and digest != args.corpus_sha256: + raise ValueError("corpus SHA256 does not match the requested pin") + corpus = [] + for number, line in enumerate(raw.decode("utf-8").splitlines(), 1): + if not line.strip(): + continue + identifier, separator, text = line.partition("\t") + if not separator or not identifier or not text.strip(): + raise ValueError(f"invalid corpus row {number}: expected idtext") + corpus.append({"id": identifier, "text": text}) + else: + digest = None + corpus = [{"id": "text", "text": args.text}] + if not corpus or any(not value["text"].strip() for value in corpus): + raise ValueError("corpus must contain nonempty text") + config = { + key: str(value) if isinstance(value, Path) else value for key, value in vars(args).items() + } + report = { + "schema": "nemo-speech.portable-http-benchmark.v1", + "status": "RUNNING", + "client_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + "python": platform.python_version(), + "configuration": config, + "corpus_sha256": digest, + "corpus": corpus, + "warmup": [], + "trials": [], + "definitions": { + "comparability": "distinct portable client; not interchangeable with measure_public_api.py or tts_perf_client", + "ttfa": "request stopwatch to first complete PCM16 sample returned by HTTPResponse.read1", + "request_stopwatch": ( + "after TCP/TLS connect" if args.exclude_connect else "before TCP/TLS connect" + ), + "connection_pooling": False, + "request_queue": "stopwatch excludes executor queue and fixed-cohort barrier wait", + "timeout": "overall request budget including connect, enforced at socket-operation/read boundaries", + "fixed_cohort": "barrier before connect; next wave starts only after every current request terminates", + "closed_loop": "bounded thread pool immediately refills completed requests; no wave barrier", + "rtfx": "sum observed sample-derived audio seconds, including failed partial responses, divided by trial wall time through drain", + "completed_rtfx": "sum audio seconds from successful responses divided by the same trial wall time", + "latency_population": "ttfa/completion/ITL use completed requests; observed_ttfa/request_elapsed also retain failures", + "quantiles": "nearest rank within trial; equal mean of trial quantiles in summary", + "itl": "sample-threshold reconstruction, not server codec events; same read may complete multiple boundaries with zero interval", + "inter_frame_ms": "interval between complete logical chunks divided by current chunk frame count", + "partial_chunk": "retained in sample totals; excluded from reconstructed inter-chunk/frame latency", + "playback": "playback starts at first aligned sample, pauses on underrun; deficits measured before adding new samples; no post-EOF silence", + "memory": "not measured; collect GPU telemetry separately", + }, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("x") as stream: + json.dump(report, stream, indent=2) + stream.write("\n") + executor = ThreadPoolExecutor(max_workers=args.concurrency, thread_name_prefix="tts-bench") + try: + if args.warmup_waves: + warmup = {"phase": "warmup"} + report["warmup"].append(warmup) + run_trial( + args, corpus, args.warmup_waves * args.concurrency, "fixed-cohort", executor, warmup + ) + save(args.output, report) + for index in range(args.trials): + measured = {"phase": "measured", "trial_index": index} + report["trials"].append(measured) + run_trial(args, corpus, args.num_requests, args.mode, executor, measured) + report["summary"] = summarize(report["trials"]) + save(args.output, report) + print(json.dumps({"trial": index, **measured["summary"]}), flush=True) + warmup_failures = sum(value["summary"]["failures"] for value in report["warmup"]) + report["warmup_failures"] = warmup_failures + report["status"] = "FAIL" if warmup_failures or report["summary"]["failures"] else "PASS" + except BaseException as error: + report["status"] = "INTERRUPTED" if isinstance(error, KeyboardInterrupt) else "ERROR" + report["error"] = {"type": type(error).__name__, "message": str(error)} + raise + finally: + executor.shutdown(wait=True) + report["summary"] = summarize(report["trials"]) + save(args.output, report) + print( + json.dumps( + {"status": report["status"], "output": str(args.output), "summary": report["summary"]} + ), + flush=True, + ) + return 0 if report["status"] == "PASS" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/nemotron_native/build_kernels.py b/scripts/nemotron_native/build_kernels.py new file mode 100644 index 0000000..e25e1b6 --- /dev/null +++ b/scripts/nemotron_native/build_kernels.py @@ -0,0 +1,300 @@ +import argparse +import concurrent.futures +import hashlib +import json +import os +import re +import shutil +import signal +import subprocess +import sys +import time +from pathlib import Path + +REPOSITORY = Path(__file__).resolve().parents[2] +SOURCE_ROOT = REPOSITORY / "deploy/nemotron_native/kernels" +SCHEMA = "nemotron-native-kernel-recipes.v2" +SUPPORTED_WIDTHS = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096) + + +def digest(path): + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def confined(root, value): + if not isinstance(value, str) or not value or "\\" in value: + raise ValueError(f"Invalid relative path: {value!r}") + relative = Path(value) + if relative.is_absolute() or ".." in relative.parts: + raise ValueError(f"Unsafe relative path: {value!r}") + target = root / relative + if not target.resolve().is_relative_to(root.resolve()): + raise ValueError(f"Path escapes root: {value!r}") + return target + + +def write_json(path, value): + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(value, indent=2) + "\n") + temporary.replace(path) + + +def load_manifest(root=SOURCE_ROOT): + path = root / "recipes.json" + manifest = json.loads(path.read_text()) + if manifest.get("schema") != SCHEMA or manifest.get("architecture") != "sm_90a": + raise ValueError("Unsupported kernel recipe schema or architecture") + files = manifest["files"] + if not files: + raise ValueError("Empty source manifest") + for relative, expected in files.items(): + path = confined(root, relative) + if not re.fullmatch(r"[0-9a-f]{64}", expected): + raise ValueError(f"Invalid source digest: {relative}") + if not path.is_file() or digest(path) != expected: + raise ValueError(f"Missing or changed kernel source: {relative}") + recipes = manifest["recipes"] + if {r["width"] for r in recipes.values()} != set(SUPPORTED_WIDTHS): + raise ValueError("Incomplete physical-width catalog") + destinations = set() + for name, recipe in recipes.items(): + for relative in [recipe["source"], *recipe["dependencies"]]: + if relative not in files: + raise ValueError(f"Unpinned source for {name}: {relative}") + for relative in recipe["include_dirs"]: + if not confined(root, relative).is_dir(): + raise ValueError(f"Missing include directory: {relative}") + destination = recipe["output"] + confined(Path("/output"), destination) + if destination in destinations: + raise ValueError(f"Duplicate library destination: {destination}") + destinations.add(destination) + flags = recipe["flags"] + if not flags or any(not isinstance(x, str) for x in flags): + raise ValueError(f"Invalid compiler arguments: {name}") + if any(x.startswith(("sha256:", "-I/")) or x in ("-o", "--output-file") for x in flags): + raise ValueError(f"Nonportable compiler arguments: {name}") + alias = recipe.get("alias_of") + if alias is not None: + if alias not in recipes or recipes[alias].get("alias_of"): + raise ValueError(f"Invalid alias: {name}") + target = recipes[alias] + if any(recipe[key] != target[key] for key in ("source", "flags", "include_dirs")): + raise ValueError(f"Alias changes compiled implementation: {name}") + for relative in recipe["dependencies"]: + path = confined(root, relative) + for include in re.findall( + r'^\s*#\s*include\s*"([^"\n]+)"', path.read_text(), re.MULTILINE + ): + candidates = [ + path.parent / include, + *[root / d / include for d in recipe["include_dirs"]], + ] + matched = next((p for p in candidates if p.is_file()), None) + if matched is None or not matched.resolve().is_relative_to(root.resolve()): + raise ValueError(f"Unresolved local include in {relative}: {include}") + if str(matched.resolve().relative_to(root.resolve())) not in files: + raise ValueError(f"Unpinned local include in {relative}: {include}") + return manifest + + +def select_recipes(manifest, widths): + selected = {k: v for k, v in manifest["recipes"].items() if v["width"] in widths} + required = dict(selected) + for recipe in selected.values(): + alias = recipe.get("alias_of") + if alias: + required[alias] = manifest["recipes"][alias] + return selected, required + + +def command(recipe, output, nvcc, root=SOURCE_ROOT): + source = confined(root, recipe["source"]) + destination = confined(output, recipe["output"]) + return [ + nvcc, + str(source), + "-o", + str(destination), + *recipe["flags"], + *["-I" + str(confined(root, p)) for p in recipe["include_dirs"]], + ] + + +def parse_widths(value): + try: + widths = sorted({int(x.removeprefix("b")) for x in value.split(",")}) + except ValueError as error: + raise argparse.ArgumentTypeError("Widths must be comma-separated integers") from error + if not widths or any(x not in SUPPORTED_WIDTHS for x in widths): + raise argparse.ArgumentTypeError("Unsupported physical width") + return widths + + +def compile_one(name, recipe, output, nvcc, timeout, source_digest): + destination = confined(output, recipe["output"]) + destination.parent.mkdir(parents=True, exist_ok=True) + log = confined(output, "logs/" + name.replace("/", "-") + ".log") + log.parent.mkdir(parents=True, exist_ok=True) + argv = command(recipe, output, nvcc) + start = time.monotonic() + result = { + "recipe": name, + "path": recipe["output"], + "source_manifest_sha256": source_digest, + "command": argv, + "log": str(log.relative_to(output)), + "status": "FAILED", + } + try: + with log.open("w") as handle: + process = subprocess.Popen( + argv, stdout=handle, stderr=subprocess.STDOUT, start_new_session=True + ) + try: + returncode = process.wait(timeout=timeout) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGTERM) + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGKILL) + process.wait() + raise + result["returncode"] = returncode + if returncode != 0: + return result + if not destination.is_file() or destination.read_bytes()[:4] != b"\x7fELF": + result["error"] = "Compiler did not create an ELF library" + return result + result.update( + status="BUILT_UNQUALIFIED", sha256=digest(destination), bytes=destination.stat().st_size + ) + except (OSError, subprocess.TimeoutExpired) as error: + result["error"] = str(error) + finally: + result["seconds"] = time.monotonic() - start + write_json(output / "receipts" / (name.replace("/", "-") + ".json"), result) + return result + + +def main(argv=None): + parser = argparse.ArgumentParser() + parser.add_argument("--output", type=Path) + parser.add_argument("--widths", type=parse_widths, default=list(SUPPORTED_WIDTHS)) + parser.add_argument("--nvcc", default=os.environ.get("NVCC", "nvcc")) + parser.add_argument("--jobs", type=int, default=1) + parser.add_argument("--timeout", type=int, default=3600) + mode = parser.add_mutually_exclusive_group() + mode.add_argument("--plan-only", action="store_true") + mode.add_argument("--validate", action="store_true") + args = parser.parse_args(argv) + args.widths = sorted({1, *args.widths}) + if not 1 <= args.jobs <= 8 or args.timeout <= 0: + parser.error("jobs must be 1..8 and timeout must be positive") + if args.output is None and not args.validate: + parser.error("--output is required") + manifest = load_manifest() + selected, required = select_recipes(manifest, args.widths) + source_digest = digest(SOURCE_ROOT / "recipes.json") + if args.validate: + print( + json.dumps( + { + "status": "PASS_SOURCE_CLOSURE", + "recipes": len(selected), + "source_files": len(manifest["files"]), + "source_manifest_sha256": source_digest, + } + ) + ) + return 0 + output = args.output.resolve() + if output.is_relative_to(SOURCE_ROOT) or SOURCE_ROOT.is_relative_to(output): + raise ValueError("Build output cannot overlap kernel sources") + if (output / "BUILD.json").exists(): + raise ValueError("Build result already exists; use a fresh output directory") + output.mkdir(parents=True, exist_ok=True) + plan = { + "schema": "nemotron-native-kernel-build.v1", + "status": "PLANNED", + "architecture": manifest["architecture"], + "widths": args.widths, + "source_manifest_sha256": source_digest, + "builder_sha256": digest(Path(__file__)), + "recipes": { + k: {**v, "command": command(v, output, args.nvcc)} for k, v in required.items() + }, + "library_map": { + v["output"]: {"recipe": k, "path": v["output"]} for k, v in selected.items() + }, + } + write_json(output / "PLAN.json", plan) + if args.plan_only: + print( + json.dumps( + {"status": "PLANNED", "recipes": len(selected), "plan": str(output / "PLAN.json")} + ) + ) + return 0 + version = subprocess.run( + [args.nvcc, "--version"], capture_output=True, text=True, timeout=30, check=True + ) + plan["nvcc_version"] = version.stdout.strip() + pending = {k: v for k, v in required.items() if not v.get("alias_of")} + results = {} + with concurrent.futures.ThreadPoolExecutor(max_workers=args.jobs) as pool: + futures = { + pool.submit(compile_one, k, v, output, args.nvcc, args.timeout, source_digest): k + for k, v in pending.items() + } + for future in concurrent.futures.as_completed(futures): + name = futures[future] + results[name] = future.result() + print( + json.dumps( + { + "recipe": name, + "status": results[name]["status"], + "seconds": results[name]["seconds"], + } + ), + flush=True, + ) + for name, recipe in selected.items(): + if not recipe.get("alias_of"): + continue + original = results[recipe["alias_of"]] + result = { + **original, + "recipe": name, + "path": recipe["output"], + "alias_of": recipe["alias_of"], + "seconds": 0.0, + } + if original["status"] == "BUILT_UNQUALIFIED": + destination = confined(output, recipe["output"]) + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(confined(output, original["path"]), destination) + results[name] = result + write_json(output / "receipts" / (name.replace("/", "-") + ".json"), result) + library_map = {selected[k]["output"]: results[k] for k in selected} + status = ( + "BUILT_UNQUALIFIED" + if all(v["status"] == "BUILT_UNQUALIFIED" for v in library_map.values()) + else "FAILED" + ) + write_json( + output / "BUILD.json", + {**plan, "status": status, "library_map": library_map, "results": results}, + ) + return 0 if status == "BUILT_UNQUALIFIED" else 1 + + +if __name__ == "__main__": + try: + sys.exit(main()) + except (OSError, ValueError, KeyError, subprocess.SubprocessError) as error: + print(str(error), file=sys.stderr) + sys.exit(1) diff --git a/scripts/nemotron_native/build_package.py b/scripts/nemotron_native/build_package.py new file mode 100644 index 0000000..aeb6484 --- /dev/null +++ b/scripts/nemotron_native/build_package.py @@ -0,0 +1,423 @@ +import argparse +import copy +import json +import os +import shutil +import struct +import sys +import tempfile +from pathlib import Path + +from build_kernels import ( + REPOSITORY, + SOURCE_ROOT, + SUPPORTED_WIDTHS, + confined, + digest, + load_manifest, + parse_widths, + write_json, +) + +CATALOG = REPOSITORY / "deploy/nemotron_native/programs.json" +BATCHING_MODES = ("fixed_cohort", "fixed_request", "inflight") +INFLIGHT_PROFILES = ("retained", "playback-safe") + + +def playback_profiles(): + path = REPOSITORY / "deploy/nemotron_native/profiles/playback-safe.json" + value = json.loads(path.read_text()) + policy = {"playback_reserve_guard": True, "codec_deadline_priority": True} + expected = { + str(width): { + "physical_width": width, + "policy": {**policy, **({"retire_batch": True} if width == 1024 else {})}, + "environment": { + "NEMOTRON_INITIAL_AUDIO_HOLD_MS": hold, + "NEMOTRON_MAX_ACTIVE_REQUESTS": cap, + }, + } + for width, hold, cap in ((1024, "40", None), (2048, "300", "1024")) + } + if value != {"schema": "nemotron.playback-profile.v1", "profiles": expected}: + raise ValueError("Unsupported playback deployment profile") + return expected + + +def profile_policy(width, batching_mode, profile): + if profile not in INFLIGHT_PROFILES: + raise ValueError("Unsupported in-flight profile") + if profile == "retained": + return {} + if batching_mode != "inflight": + raise ValueError("Playback deployment profile requires inflight batching") + return playback_profiles().get(str(width), {}).get("policy", {}) + + +def source_seal(): + root = REPOSITORY / "python/nemotron_native_backend" + return { + "package_source_pins": { + str(path.relative_to(root)): digest(path) for path in sorted(root.rglob("*.py")) + } + } + + +def load_catalog(manifest): + sys.path.insert(0, str(REPOSITORY / "python")) + try: + from nemotron_native_backend.execution.inflight_policy import InflightPolicy + finally: + sys.path.pop(0) + + catalog = json.loads(CATALOG.read_text()) + if ( + catalog.get("schema") != "nemotron.native-program-catalog.v1" + or catalog.get("architecture") != manifest["architecture"] + or set(catalog["programs"]) != {str(width) for width in SUPPORTED_WIDTHS} + ): + raise ValueError("Unsupported native program catalog") + recipes = {recipe["output"]: recipe for recipe in manifest["recipes"].values()} + for key, program in catalog["programs"].items(): + width = int(key) + host = program["serving"]["host"] + if "inflight_policy" in host or host.get("batching_mode", "fixed_cohort") != "fixed_cohort": + raise ValueError("Program catalog must preserve the fixed-cohort serving defaults") + InflightPolicy.parse(program.get("inflight_policy")) + if ( + program["neural"]["physical_width"] != width + or program["serving"]["physical_width"] != width + ): + raise ValueError("Inconsistent physical width") + for role, path in program["libraries"].items(): + recipe = recipes[path] + if recipe["width"] != width or recipe["role"] != role: + raise ValueError("Program changes kernel binding: " + path) + return catalog + + +def validate_build(build_root, manifest, catalog, widths): + build = json.loads((build_root / "BUILD.json").read_text()) + if ( + build.get("schema") != "nemotron-native-kernel-build.v1" + or build.get("status") != "BUILT_UNQUALIFIED" + or build.get("source_manifest_sha256") != digest(SOURCE_ROOT / "recipes.json") + or build.get("architecture") != manifest["architecture"] + or not set(widths) <= set(build["widths"]) + or not build.get("nvcc_version") + ): + raise ValueError("Missing, failed, or stale native kernel build") + for width in widths: + for relative in catalog["programs"][str(width)]["libraries"].values(): + row = build["library_map"][relative] + recipe = manifest["recipes"][row["recipe"]] + path = confined(build_root, relative) + if ( + row["status"] != "BUILT_UNQUALIFIED" + or row["returncode"] != 0 + or row["path"] != relative + or row["source_manifest_sha256"] != build["source_manifest_sha256"] + or recipe["output"] != relative + or row.get("bytes") != path.stat().st_size + or digest(path) != row["sha256"] + ): + raise ValueError("Invalid compiled library: " + relative) + compiled = manifest["recipes"][recipe.get("alias_of", row["recipe"])] + argv = row.get("command") + if ( + not isinstance(argv, list) + or len(argv) != 4 + len(compiled["flags"]) + len(compiled["include_dirs"]) + or any(not isinstance(arg, str) or not arg or "\x00" in arg for arg in argv) + or argv[2] != "-o" + ): + raise ValueError("Invalid compiler command: " + relative) + roots = [] + for actual, expected in ((argv[1], compiled["source"]), (argv[3], compiled["output"])): + absolute = Path(actual) + suffix = Path(expected) + if ( + not absolute.is_absolute() + or ".." in absolute.parts + or absolute.parts[-len(suffix.parts) :] != suffix.parts + or len(absolute.parts) <= len(suffix.parts) + ): + raise ValueError("Compiler path differs from recipe: " + relative) + roots.append(absolute.parents[len(suffix.parts) - 1]) + source_root, output_root = roots + expected_argv = [ + argv[0], + str(source_root / compiled["source"]), + "-o", + str(output_root / compiled["output"]), + *compiled["flags"], + *["-I" + str(source_root / entry) for entry in compiled["include_dirs"]], + ] + if argv != expected_argv: + raise ValueError("Compiler command differs from recipe: " + relative) + with path.open("rb") as stream: + header = stream.read(64) + if ( + len(header) != 64 + or header[:6] != b"\x7fELF\x02\x01" + or struct.unpack_from(" path.stat().st_size + ): + raise ValueError("Native library is not a complete x86-64 shared ELF: " + relative) + return build + + +def state_artifacts(root): + state = root / "state/bf16" + state.mkdir(parents=True, exist_ok=True) + scales = state / "scales.f32" + mapping = state / "map.i32" + scales.write_bytes(struct.pack(" last or last > max(available, default=-1): + raise ValueError("requested CPUs must be within the launcher CPU affinity") + selected.update(range(first, last + 1)) + if not selected or not selected.issubset(available): + raise ValueError("requested CPUs must be within the launcher CPU affinity") + return sorted(selected) + + +def cpu_configuration(args): + native = args.native_cpus + frontend = args.frontend_cpus + if frontend is not None and args.http_port is None: + raise ValueError("frontend-cpus requires http-port") + supported = sys.platform == "linux" and hasattr(os, "sched_getaffinity") + requested = native is not None or frontend is not None + if requested and not supported: + raise ValueError("CPU role controls require Linux CPU affinity support") + available = set(os.sched_getaffinity(0)) if supported else None + taskset = shutil.which("taskset") if requested else None + if requested and taskset is None: + raise ValueError("CPU role controls require taskset in PATH") + roles = {"worker": native} + if not getattr(args, "direct_worker", False): + roles["bridge"] = native + if args.http_port is not None: + roles["frontend"] = frontend + configuration = { + "available_cpus": sorted(available) if available is not None else None, + "roles": { + role: { + "requested_cpus": parse_cpus(value, available) if value is not None else None, + "effective_cpus": None, + } + for role, value in roles.items() + }, + } + return taskset, configuration + + +def cpu_command(command, role, taskset, configuration): + cpus = configuration["roles"][role]["requested_cpus"] + return [taskset, "--cpu-list", ",".join(map(str, cpus)), *command] if cpus else command + + +def record_cpu_affinity(configuration, role, process): + record = configuration["roles"][role] + record["pid"] = process.pid + if configuration["available_cpus"] is not None: + record["effective_cpus"] = sorted(os.sched_getaffinity(process.pid)) + if ( + record["requested_cpus"] is not None + and record["effective_cpus"] != record["requested_cpus"] + ): + raise RuntimeError(role + " CPU affinity differs from its requested CPUs") + + +def commands(args): + programs = args.package / "programs" + selected = programs / f"b{args.width}" + serving = selected / "serving.json" + plan = json.loads(serving.read_text()) + if ( + plan["physical_width"] != args.width + or plan["host"].get("batching_mode", "fixed_cohort") != args.mode + ): + raise ValueError("width or mode differs from the assembled package") + worker = [ + sys.executable, + "-B", + str(REPOSITORY / "deploy/nemotron_native/serve.py"), + "--qualification-only", + "--plan", + str(serving), + "--neural-plan", + str(selected / "neural.json"), + "--readiness", + str(programs / "b1/neural.json"), + "--readiness-artifacts", + str(programs / "b1/artifacts"), + "--artifacts", + str(selected / "artifacts"), + "--model", + str(args.model), + "--model-seal", + str(args.model_seal), + "--source-seal", + str(args.package / "source-seal.json"), + "--receipts", + str(args.receipts / "worker"), + "--socket", + str(args.native_worker_socket), + ] + bridge = [ + sys.executable, + "-B", + "-m", + "nemotron_cpp_bridge", + "--native-worker-socket", + str(args.native_worker_socket), + "--manifest", + str(serving), + "--model", + str(args.model), + "--model-seal", + str(args.model_seal), + "--listen", + str(args.listen), + "--width", + str(args.width), + "--mode", + args.mode, + ] + if args.exercise_lifecycle: + worker.append("--exercise-lifecycle") + if args.normalization_config: + bridge.extend(["--normalization-config", str(args.normalization_config.resolve())]) + return worker, bridge + + +def direct_worker_configuration(args): + from nemotron_cpp_bridge.native_preparation import FrontendPreparation + + serving = args.package / "programs" / f"b{args.width}" / "serving.json" + raw = serving.read_bytes() + plan = json.loads(raw) + if ( + type(plan["physical_width"]) is not int + or not 1 <= args.width <= 4096 + or args.width & (args.width - 1) + or plan["physical_width"] != args.width + or plan["host"].get("batching_mode", "fixed_cohort") != args.mode + ): + raise ValueError("direct worker width or mode differs from package") + transport = plan["transport_compatibility"] + if transport["output_contract"] != "pcm_f32_le_22050_mono_variable_streaming_v3": + raise ValueError("direct worker requires streaming-v3 output") + versions = transport["installed_transport_versions"] + if ( + not isinstance(versions, dict) + or set(versions) != {"trt_version", "trtllm_version", "trtllm_commit"} + or any( + not isinstance(value, str) or not 1 <= len(value) <= 1024 for value in versions.values() + ) + ): + raise ValueError("invalid direct worker transport versions") + if ( + not isinstance(plan["model"]["checkpoint_sha256"], str) + or re.fullmatch(r"[0-9a-f]{64}", plan["model"]["checkpoint_sha256"]) is None + ): + raise ValueError("invalid direct worker model identity") + seal = json.loads(args.model_seal.read_text()) + model = args.model.resolve(strict=True) + tokenizer = (model / "tokenizer.json").resolve(strict=True) + expected = seal["files"]["tokenizer.json"] + if ( + seal["model_identity"] != plan["model"]["checkpoint_sha256"] + or not tokenizer.is_relative_to(model) + or tokenizer.stat().st_size != expected["bytes"] + or hashlib.sha256(tokenizer.read_bytes()).hexdigest() != expected["sha256"] + ): + raise ValueError("tokenizer does not match the selected sealed model") + library = os.environ.get("NEMOTRON_NATIVE_PREPARATION_LIBRARY") + if not library: + raise ValueError("direct-worker requires the native preparation library") + preparation = FrontendPreparation(library, tokenizer, args.normalization_config) + return { + "schema": "nemotron.direct_worker.v1", + "width": args.width, + "mode": args.mode, + "manifest_sha256": hashlib.sha256(raw).hexdigest(), + "model_sha256": plan["model"]["checkpoint_sha256"], + "output_contract": transport["output_contract"], + "enabled_buckets": sorted({1, args.width}), + "installed_transport_versions": versions, + "text_normalization_enabled": preparation.enabled, + "text_preparation": preparation.metadata, + } + + +async def wait_ready(path, processes, timeout, stopping): + async def observe(): + while not path.exists(): + if stopping.is_set(): + raise InterruptedError("startup interrupted") + if any(process.returncode is not None for process in processes): + raise RuntimeError("a supervised process exited during startup") + await asyncio.sleep(0.1) + + await asyncio.wait_for(observe(), timeout) + + +async def wait_http(host, port, processes, timeout, stopping): + connect_host = "127.0.0.1" if host == "0.0.0.0" else "::1" if host == "::" else host + + async def observe(): + while True: + if stopping.is_set(): + raise InterruptedError("startup interrupted") + if any(process.returncode is not None for process in processes): + raise RuntimeError("a supervised process exited during HTTP startup") + try: + _, writer = await asyncio.wait_for(asyncio.open_connection(connect_host, port), 1) + writer.close() + await writer.wait_closed() + return + except (OSError, asyncio.TimeoutError): + await asyncio.sleep(0.1) + + await asyncio.wait_for(observe(), timeout) + + +async def run(args): + taskset, affinity = cpu_configuration(args) + worker_command, bridge_command = commands(args) + worker_command = cpu_command(worker_command, "worker", taskset, affinity) + direct = getattr(args, "direct_worker", False) + direct_metadata = direct_worker_configuration(args) if direct else None + if not direct: + bridge_command = cpu_command(bridge_command, "bridge", taskset, affinity) + if args.listen == args.native_worker_socket: + raise ValueError("worker and bridge sockets must differ") + for path in (args.listen, args.native_worker_socket): + if path.exists() or path.is_symlink(): + raise ValueError("refusing to replace existing socket: " + str(path)) + path.parent.mkdir(parents=True, exist_ok=True) + if args.receipts.exists() and any(args.receipts.iterdir()): + raise ValueError("receipts directory must be new or empty") + args.receipts.mkdir(parents=True, exist_ok=True) + stopping = asyncio.Event() + loop = asyncio.get_running_loop() + for sig in (signal.SIGTERM, signal.SIGINT): + loop.add_signal_handler(sig, stopping.set) + environment = os.environ.copy() + environment.pop("NEMOTRON_DIRECT_WORKER_CONFIG", None) + environment.update( + PYTHONPATH=str(REPOSITORY / "python") + + (os.pathsep + environment["PYTHONPATH"] if environment.get("PYTHONPATH") else ""), + PYTHONDONTWRITEBYTECODE="1", + PYTHONUNBUFFERED="1", + TOKENIZERS_PARALLELISM="false", + NVIDIA_TF32_OVERRIDE="0", + CUDA_MODULE_LOADING="LAZY", + ) + package = json.loads((args.package / "PACKAGE.json").read_text()) + profile = package.get("deployment_profiles", {}).get(str(args.width), {}) + for key, value in profile.get("environment", {}).items(): + if value is not None: + environment.setdefault(key, value) + frontend_environment = environment.copy() + if direct: + direct_path = args.receipts / "DIRECT_WORKER.json" + direct_path.write_text(json.dumps(direct_metadata, indent=2) + "\n") + direct_path.chmod(0o600) + frontend_environment["NEMOTRON_DIRECT_WORKER_CONFIG"] = str(direct_path) + processes = [] + logs = [] + tasks = [] + report = { + "status": "STARTING", + "width": args.width, + "mode": args.mode, + "cpu_affinity": affinity, + "transport": "direct_worker" if direct else "compatibility_bridge", + } + try: + log = (args.receipts / "worker.log").open("wb") + logs.append(log) + worker = await asyncio.create_subprocess_exec( + *worker_command, + env=environment, + stdout=log, + stderr=asyncio.subprocess.STDOUT, + ) + processes.append(worker) + await wait_ready( + args.receipts / "worker/NATIVE_READY.json", processes, args.startup_timeout, stopping + ) + record_cpu_affinity(affinity, "worker", worker) + if not direct: + log = (args.receipts / "bridge.log").open("wb") + logs.append(log) + bridge = await asyncio.create_subprocess_exec( + *bridge_command, + env=environment, + stdout=log, + stderr=asyncio.subprocess.STDOUT, + ) + processes.append(bridge) + await wait_ready(args.listen, processes, args.startup_timeout, stopping) + record_cpu_affinity(affinity, "bridge", bridge) + if args.http_port is not None: + frontend_command = [ + args.frontend_binary, + "serve", + "--tts-backend", + "nemotron", + "--native-socket", + str(args.native_worker_socket if direct else args.listen), + "--port", + str(args.http_port), + "--host", + args.http_host, + "--no-warmup", + ] + if args.http_threads is not None: + frontend_command.extend(["--http-threads", str(args.http_threads)]) + frontend_command = cpu_command(frontend_command, "frontend", taskset, affinity) + log = (args.receipts / "frontend.log").open("wb") + logs.append(log) + frontend = await asyncio.create_subprocess_exec( + *frontend_command, + env=frontend_environment, + stdout=log, + stderr=asyncio.subprocess.STDOUT, + ) + processes.append(frontend) + await wait_http( + args.http_host, args.http_port, processes, args.startup_timeout, stopping + ) + record_cpu_affinity(affinity, "frontend", frontend) + report["frontend"] = { + "command": frontend_command, + "pid": frontend.pid, + "host": args.http_host, + "port": args.http_port, + } + report.update( + status="READY", socket=str(args.native_worker_socket if direct else args.listen) + ) + (args.receipts / "LAUNCH.json").write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(report), flush=True) + stop_task = asyncio.create_task(stopping.wait()) + tasks = [stop_task, *(asyncio.create_task(process.wait()) for process in processes)] + done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + if stop_task not in done: + raise RuntimeError("a supervised process exited unexpectedly") + report["status"] = "STOPPED" + except InterruptedError: + report["status"] = "STOPPED" + except BaseException as error: + report.update(status="FAILED", error=str(error)) + raise + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + for process in reversed(processes): + if process.returncode is None: + with contextlib.suppress(ProcessLookupError): + process.terminate() + try: + await asyncio.wait_for(process.wait(), args.shutdown_timeout) + except asyncio.TimeoutError: + with contextlib.suppress(ProcessLookupError): + process.kill() + await asyncio.wait_for(process.wait(), 5) + for log in logs: + log.close() + report["exit_codes"] = [process.returncode for process in processes] + (args.receipts / "EXIT.json").write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + asyncio.run(run(parse_args())) diff --git a/scripts/nemotron_native/requirements.txt b/scripts/nemotron_native/requirements.txt new file mode 100644 index 0000000..ac1a84f --- /dev/null +++ b/scripts/nemotron_native/requirements.txt @@ -0,0 +1,4 @@ +msgpack==1.2.2 +numpy==2.3.5 +safetensors==0.8.0 +tokenizers==0.23.1 diff --git a/scripts/nemotron_native/runtime_dependencies.py b/scripts/nemotron_native/runtime_dependencies.py new file mode 100644 index 0000000..bb3a062 --- /dev/null +++ b/scripts/nemotron_native/runtime_dependencies.py @@ -0,0 +1,245 @@ +import argparse +import hashlib +import importlib +import importlib.metadata +import importlib.util +import json +import os +import re +import shutil +import subprocess +import sys +from pathlib import Path + +SCHEMA = "nemotron-native-runtime-dependencies.v1" +MODULES = ("tensorrt", "tensorrt_llm", "torch_tensorrt") +REQUIRED = ("msgpack", "numpy", "safetensors", "tokenizers", "torch") +BACKEND_MODULES = ( + "nemotron_native_backend.worker", + "nemotron_native_backend.execution.inflight_engine", + "nemotron_native_backend.execution.native_stage1", +) +LIBRARY_ROOTS = (Path("/usr/lib"), Path("/usr/local/lib"), Path("/opt")) +REMNANTS = (Path("/opt/tensorrt"), Path("/usr/src/tensorrt")) +SYSTEM_LIBRARIES = Path("/usr/lib/x86_64-linux-gnu") + + +def canonical_name(name): + return re.sub(r"[-_.]+", "-", name).lower().split(":", 1)[0] + + +def is_tensorrt_package(name): + value = canonical_name(name) + return ( + value == "polygraphy" + or value.startswith(("tensorrt", "torch-tensorrt", "nvidia-tensorrt")) + or value.startswith(("libnvinfer", "libnvonnxparsers", "python3-libnvinfer")) + ) + + +def is_tensorrt_library(name): + return ( + Path(name) + .name.lower() + .startswith( + ( + "libnvinfer", + "libnvonnxparser", + "libtensorrt", + "libtorchtrt", + "do_not_link_against_nvinfer", + ) + ) + ) + + +def run(*args): + return subprocess.run(args, check=True, capture_output=True, text=True).stdout + + +def debian_packages(): + output = run("dpkg-query", "-W", "-f=${binary:Package}\t${db:Status-Status}\n") + return sorted( + name + for line in output.splitlines() + for name, status in [line.split("\t", 1)] + if status == "installed" and is_tensorrt_package(name) + ) + + +def python_packages(): + return sorted( + {d.metadata["Name"] for d in importlib.metadata.distributions() if d.metadata["Name"]} + ) + + +def validate_purge_plan(output, requested): + selected = {canonical_name(name) for name in requested} + proposed = { + canonical_name(line.split()[1]) + for line in output.splitlines() + if line.startswith(("Remv ", "Purg ")) + } + if proposed != selected: + raise RuntimeError(f"TensorRT purge changes unexpected packages: {proposed ^ selected}") + + +def remove_builder_resource_links(root): + removed = [] + for path in sorted(root.glob("do_not_link_against_nvinfer_builder_resource_*")): + if not path.is_symlink() or path.exists(): + raise RuntimeError(f"Expected an orphaned TensorRT builder resource link: {path}") + if not path.readlink().name.startswith("libnvinfer_builder_resource_"): + raise RuntimeError(f"Unexpected TensorRT resource link target: {path}") + path.unlink() + removed.append(str(path)) + return removed + + +def remove_tensorrt(): + python = [name for name in python_packages() if is_tensorrt_package(name)] + debian = debian_packages() + if python: + subprocess.run([sys.executable, "-m", "pip", "uninstall", "--yes", *python], check=True) + if debian: + simulated = run("apt-get", "--simulate", "--no-auto-remove", "purge", *debian) + validate_purge_plan(simulated, debian) + subprocess.run(["apt-get", "--yes", "--no-auto-remove", "purge", *debian], check=True) + removed = remove_builder_resource_links(SYSTEM_LIBRARIES) + for path in REMNANTS: + if path.is_symlink() or path.is_file(): + path.unlink() + removed.append(str(path)) + elif path.exists(): + shutil.rmtree(path) + removed.append(str(path)) + subprocess.run(["ldconfig"], check=True) + return {"python_packages": python, "debian_packages": debian, "remnants": removed} + + +def find_tensorrt_libraries(roots): + found = set() + for root in roots: + for directory, subdirectories, files in os.walk(root, followlinks=False): + subdirectories[:] = [name for name in subdirectories if name != "__pycache__"] + for name in files: + if is_tensorrt_library(name): + found.add(str(Path(directory) / name)) + return sorted(found) + + +def needed_libraries(path): + output = run("readelf", "--dynamic", str(path)) + return re.findall(r"\(NEEDED\).*?\[([^\]]+)\]", output) + + +def audit_elf(path): + path = Path(path) + needed = needed_libraries(path) + with path.open("rb") as handle: + digest = hashlib.file_digest(handle, "sha256").hexdigest() + return { + "path": str(path), + "sha256": digest, + "needed": needed, + "tensorrt_needed": [name for name in needed if is_tensorrt_library(name)], + } + + +def import_runtime(include_backend, required=REQUIRED): + modules = [*required, *(BACKEND_MODULES if include_backend else ())] + imported = {} + for name in modules: + module = importlib.import_module(name) + imported[name] = getattr(module, "__version__", None) + torch = sys.modules["torch"] + if torch.cuda.is_initialized(): + raise RuntimeError("The CPU dependency audit unexpectedly initialized CUDA") + return imported + + +def check(*, include_backend=False, library_roots=LIBRARY_ROOTS, elf_roots=(), required=REQUIRED): + imported = import_runtime(include_backend, required) + python = [name for name in python_packages() if is_tensorrt_package(name)] + debian = debian_packages() + modules = [name for name in MODULES if importlib.util.find_spec(name) is not None] + loaded = [ + name + for name in sys.modules + if any(name == root or name.startswith(root + ".") for root in MODULES) + ] + libraries = find_tensorrt_libraries(library_roots) + loader = [ + line.strip() + for line in run("ldconfig", "-p").splitlines() + if line.strip() and is_tensorrt_library(line.split()[0]) + ] + bad_environment = { + name: value + for name in ("PATH", "LD_LIBRARY_PATH", "LIBRARY_PATH", "PYTHONPATH") + if any(part in (value := os.environ.get(name, "")).lower() for part in MODULES) + } + elf = [] + roots = [Path(sys.modules["torch"].__file__).parent / "lib", *map(Path, elf_roots)] + for root in roots: + if not root.is_dir(): + raise ValueError(f"ELF audit directory does not exist: {root}") + for path in sorted(root.rglob("*.so*")): + if path.is_file() and not path.is_symlink(): + with path.open("rb") as handle: + if handle.read(4) != b"\x7fELF": + continue + elf.append(audit_elf(path)) + findings = { + "python_packages": python, + "debian_packages": debian, + "importable_modules": modules, + "loaded_modules": loaded, + "library_files": libraries, + "loader_cache": loader, + "search_paths": bad_environment, + "elf_dependencies": [row for row in elf if row["tensorrt_needed"]], + } + return { + "schema": SCHEMA, + "status": "PASS" if not any(findings.values()) else "FAIL", + "scope": "CPU imports and dependency absence; no GPU inference qualification", + "python": sys.version, + "imports": imported, + "torch_cuda_version": sys.modules["torch"].version.cuda, + "cuda_initialized": False, + "findings": findings, + "elf": elf, + } + + +def main(argv=None): + parser = argparse.ArgumentParser() + parser.add_argument("--remove-tensorrt", action="store_true") + parser.add_argument("--include-backend", action="store_true") + parser.add_argument("--elf-root", action="append", type=Path, default=[]) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args(argv) + receipt = {"schema": SCHEMA, "status": "FAIL"} + try: + removed = remove_tensorrt() if args.remove_tensorrt else None + receipt = check( + include_backend=args.include_backend, + elf_roots=args.elf_root, + required=("numpy", "torch") if args.remove_tensorrt else REQUIRED, + ) + if removed is not None: + receipt["removed"] = removed + except Exception as error: + receipt["error"] = f"{type(error).__name__}: {error}" + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(receipt, indent=2) + "\n") + summary = {"status": receipt["status"], "receipt": str(args.output)} + if receipt["status"] != "PASS": + summary.update(error=receipt.get("error"), findings=receipt.get("findings")) + print(json.dumps(summary)) + return 0 if receipt["status"] == "PASS" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/nemotron_native/tests/test_bridge.py b/scripts/nemotron_native/tests/test_bridge.py new file mode 100644 index 0000000..630c181 --- /dev/null +++ b/scripts/nemotron_native/tests/test_bridge.py @@ -0,0 +1,792 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import asyncio +import contextlib +import hashlib +import json +import math +import os +import struct +import sys +import tempfile +import threading +import unittest +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from types import SimpleNamespace +from unittest import mock + +ROOT = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(ROOT / "python")) + +from nemotron_cpp_bridge.protocol import ( + control_frame, + float_audio_frame, + pcm16, + pcm16_scalar, + read_control, +) +from nemotron_cpp_bridge.server import ( + Bridge, + Request, + Session, + initialize_preparation_worker, + parse_args, + preparation_cpus, +) +from nemotron_cpp_bridge.slabs import SlabReader +from nemotron_native_backend.protocol import message, read_frame, write_frame + + +class PreparationAffinityTests(unittest.TestCase): + def test_default_does_not_change_affinity(self): + with mock.patch.dict(os.environ, {}, clear=True): + self.assertIsNone(preparation_cpus()) + with mock.patch.object(os, "sched_setaffinity", create=True) as setter: + initialize_preparation_worker(None) + setter.assert_not_called() + + def test_cpu_list_validation(self): + with mock.patch.object(os, "sched_getaffinity", return_value={1, 2, 4, 5}, create=True): + for invalid in ("", "-1", "1,,2", "2-1", "0", "1-4", "9999999999", "1.5"): + with self.subTest(value=invalid), mock.patch.dict( + os.environ, {"NEMOTRON_BRIDGE_PREPARATION_CPUS": invalid} + ), self.assertRaises(ValueError): + preparation_cpus() + with mock.patch.dict(os.environ, {"NEMOTRON_BRIDGE_PREPARATION_CPUS": "1-2, 4,5"}): + self.assertEqual(preparation_cpus(), frozenset({1, 2, 4, 5})) + + @unittest.skipUnless(hasattr(os, "sched_getaffinity"), "Linux affinity required") + def test_only_executor_workers_receive_affinity(self): + original = os.sched_getaffinity(0) + chosen = frozenset({min(original)}) + with ThreadPoolExecutor( + max_workers=1, initializer=initialize_preparation_worker, initargs=(chosen,) + ) as executor: + self.assertEqual(executor.submit(os.sched_getaffinity, 0).result(), chosen) + self.assertEqual(os.sched_getaffinity(0), original) + + +class Tokenizer: + def encode_batch(self, texts, *, add_special_tokens): + if add_special_tokens: + raise AssertionError("target text must not add special tokens") + if len(texts) != 1: + raise AssertionError("text preparation must not wait for a request batch") + return [SimpleNamespace(ids=[ord(char) for char in texts[0]])] + + +class Preparer: + enabled = False + + def prepare(self, text, language, apply_normalize=True): + return text.strip() + + def close(self): + pass + + +async def receive(reader): + size = struct.unpack(">I", await reader.readexactly(4))[0] + payload = await reader.readexactly(size) + if payload[:1] == b"J": + return json.loads(payload[1:]) + request_id, sequence = struct.unpack(">QI", payload[1:13]) + return { + "kind": "audio", + "request_id": request_id, + "sequence": sequence, + "data": payload[13:], + "encoding": payload[:1], + } + + +async def wait_until(predicate): + async with asyncio.timeout(3): + while not predicate(): + await asyncio.sleep(0.001) + + +class ProtocolTests(unittest.IsolatedAsyncioTestCase): + def test_float_audio_frame_preserves_extent_and_header(self): + pcm = struct.pack("<4f", -1, 0, 0.5, 1) + encoded = float_audio_frame(123, 7, pcm) + self.assertEqual(struct.unpack(">I", encoded[:4])[0], 13 + len(pcm)) + self.assertEqual(encoded[4:17], b"F" + struct.pack(">QI", 123, 7)) + self.assertEqual(encoded[17:], pcm) + for malformed in (b"", b"x", b"abc", b"abcde"): + with self.assertRaises(ValueError): + float_audio_frame(123, 7, malformed) + + async def test_valid_text_budget_survives_json_escaping(self): + text = "\x01" * 1_000_000 + encoded = control_frame("submit", request_id=1, text=text) + self.assertGreater(len(encoded), 2 * 1024 * 1024) + reader = asyncio.StreamReader() + reader.feed_data(encoded) + reader.feed_eof() + self.assertEqual((await read_control(reader))["text"], text) + + +class PCMTests(unittest.TestCase): + def test_matches_asymmetric_half_away_rounding_and_saturation(self): + values = [-2, -1, -0.5, -0.5 / 32768, -0.0, 0, 0.5 / 32767, 0.5, 1, 2] + payload = struct.pack("<" + "f" * len(values), *values) + expected = [-32768, -32768, -16384, -1, 0, 0, 1, 16384, 32767, 32767] + self.assertEqual(list(struct.unpack("<10h", pcm16_scalar(payload))), expected) + self.assertEqual(pcm16(payload), pcm16_scalar(payload)) + + def test_rejects_nan_infinity_and_misaligned_input(self): + for payload in ( + b"abc", + *(struct.pack(" #include #include +#include #include #include #include @@ -404,6 +405,15 @@ struct Server::Impl { } else { server = std::make_unique(); } +#if defined(NEMO_SPEECH_REGISTRY_TTS) + if (const char* value = std::getenv("NEMOTRON_HTTP_TCP_NODELAY"); + value && std::strcmp(value, "1") == 0) { + const auto capabilities = this->models.capabilities(); + if (std::find(capabilities.begin(), capabilities.end(), "tts") != capabilities.end() && + this->models.tts()->backend_name() == "nemotron") + server->set_tcp_nodelay(true); + } +#endif server->set_payload_max_length(this->config.max_upload_bytes); server->set_read_timeout(this->config.read_timeout_seconds); server->set_write_timeout(this->config.write_timeout_seconds); @@ -712,8 +722,21 @@ struct Server::Impl { auto synthesizer = this->models.tts(); tts::SynthesisRequest synthesis; synthesis.text = body.string_or("input"); - synthesis.language_code = body.string_or("language"); + synthesis.language_code = + body.string_or("language", body.string_or("language_code")); synthesis.voice_name = openai_voice(*synthesizer, body.string_or("voice")); + if (const auto* normalize = body.find("apply_normalize")) { + if (!normalize->is_bool()) + throw std::invalid_argument("apply_normalize must be boolean"); + synthesis.apply_normalize = normalize->boolean() ? 1 : 0; + } + if (const auto* seed = body.find("seed")) { + if (!seed->is_number() || seed->number() != std::floor(seed->number()) || + seed->number() < 0 || seed->number() > 2147483647) + throw std::invalid_argument( + "seed must be an integer between 0 and 2147483647"); + synthesis.options.seed = static_cast(seed->number()); + } if (const auto* sample_rate = body.find("sample_rate")) { if (!sample_rate->is_number()) throw std::invalid_argument("sample_rate must be a positive integer"); @@ -734,6 +757,32 @@ struct Server::Impl { if (format != "wav" && format != "pcm") throw std::invalid_argument("response_format must be wav or pcm"); + if (format == "pcm" && synthesizer->backend_name() == "nemotron") { + const auto prepared = synthesizer->prepare(synthesis); + response.set_header( + "x-audio-sample-rate", std::to_string(prepared.metadata.sample_rate)); + response.set_header("Cache-Control", "no-store"); + response.set_chunked_content_provider( + "audio/pcm", [synthesizer, prepared](size_t, httplib::DataSink& sink) { + try { + const auto result = synthesizer->synthesize( + prepared, [&](const auto&, const std::string& chunk) { + return sink.is_writable() && + sink.write(chunk.data(), chunk.size()); + }); + if (result.cancelled) + return false; + sink.done(); + return true; + } + catch (const std::exception& error) { + std::fprintf(stderr, "Nemotron stream failed: %s\n", error.what()); + return false; + } + }); + return; + } + uint64_t generation = 0; if (this->config.preempt_tts) { generation = this->tts_preemption.claim(); @@ -785,6 +834,8 @@ struct Server::Impl { << " e2e_ttfa_ms=" << stats.e2e_ttfa_ms << " e2e_rtfx=" << stats.e2e_rtfx << '\n'; } + response.set_header( + "x-audio-sample-rate", std::to_string(result.metadata.sample_rate)); if (format == "pcm") response.set_content(std::move(pcm), "audio/pcm"); else diff --git a/src/server/riva_server.cc b/src/server/riva_server.cc index 92ba72b..a925b1f 100644 --- a/src/server/riva_server.cc +++ b/src/server/riva_server.cc @@ -131,6 +131,8 @@ asr_configured(const AsrServerConfig& cfg) { #if defined(NEMO_SPEECH_BUILD_TTS) bool tts_configured(const TtsServerConfig& cfg) { + if (cfg.server.backend == "nemotron") + return !cfg.server.native_socket.empty(); return !cfg.server.runtime.magpie_model.empty() && !cfg.server.runtime.codec_model.empty() && !cfg.server.tokenizer_model_dir.empty(); } @@ -380,6 +382,9 @@ main(int argc, char** argv) { << "\n"; try { nemo_speech::tts::SynthesizerConfig synthesizer_config; + synthesizer_config.backend = cfg.tts.server.backend; + synthesizer_config.native_socket = cfg.tts.server.native_socket; + synthesizer_config.native_timeout_seconds = cfg.tts.server.native_timeout_seconds; synthesizer_config.runtime = std::move(cfg.tts.server.runtime); synthesizer_config.tokenizer_model_dir = std::move(cfg.tts.server.tokenizer_model_dir); synthesizer_config.text_normalizer_model_dir = std::move(cfg.tts.server.tn_model_dir); diff --git a/src/tts/c_api.cpp b/src/tts/c_api.cpp index 6f7e9c3..86a1ae4 100644 --- a/src/tts/c_api.cpp +++ b/src/tts/c_api.cpp @@ -217,6 +217,13 @@ synthesizer_config_of(const nemo_speech_tts_synthesizer_config* c) { cfg.tokenizer_model_dir = tokenizer_dir_of(c); if (c && HAS_FIELD(c, nemo_speech_tts_synthesizer_config, model) && c->model) { const auto* model = c->model; + if (HAS_FIELD(model, nemo_speech_tts_model_config, model_backend) && model->model_backend) + cfg.backend = model->model_backend; + if (HAS_FIELD(model, nemo_speech_tts_model_config, native_socket)) + cfg.native_socket = str_or_empty(model->native_socket); + if (HAS_FIELD(model, nemo_speech_tts_model_config, native_timeout_seconds) && + model->native_timeout_seconds > 0) + cfg.native_timeout_seconds = model->native_timeout_seconds; if (HAS_FIELD(model, nemo_speech_tts_model_config, text_normalizer_model_dir)) { cfg.text_normalizer_model_dir = str_or_empty(model->text_normalizer_model_dir); } diff --git a/src/tts/magpietts/CMakeLists.txt b/src/tts/magpietts/CMakeLists.txt index 640ecdc..e44b189 100644 --- a/src/tts/magpietts/CMakeLists.txt +++ b/src/tts/magpietts/CMakeLists.txt @@ -13,10 +13,16 @@ add_library(nemo_speech_tts SHARED config.cpp ../c_api.cpp ../synthesizer.cpp + ../nemotron/native_runtime.cpp + ../nemotron/native_worker_protocol.cpp magpietts.cpp runtime.cpp ) +find_package(Threads REQUIRED) +target_link_libraries(nemo_speech_tts PRIVATE Threads::Threads ${CMAKE_DL_LIBS}) +target_include_directories(nemo_speech_tts PRIVATE "${NEMO_SPEECH_LLAMA_CPP_DIR}/vendor") + # Windows: this DLL relies on Linux default symbol visibility to export its C++ API. # MSVC/clang-cl export nothing from a DLL unless annotated, so direct C++ # consumers (the TTS C ABI lib, the gRPC TTS service, and smoke tests) would diff --git a/src/tts/magpietts/config.cpp b/src/tts/magpietts/config.cpp index ed22b6d..c6bd9e4 100644 --- a/src/tts/magpietts/config.cpp +++ b/src/tts/magpietts/config.cpp @@ -319,6 +319,9 @@ struct MagpieStreamParamsConfig { void MagpieTtsServerConfig::Register(common::ParameterParser& parser) { + parser.Register("backend", &backend, "TTS model backend: magpie or nemotron"); + parser.Register("native-socket", &native_socket, "Nemotron controller bridge Unix socket"); + parser.Register("native-timeout-seconds", &native_timeout_seconds, "Native request timeout"); register_runtime_config(parser, runtime); parser.Register( "tokenizer-model-dir", &tokenizer_model_dir, "Extracted Magpie .nemo directory"); diff --git a/src/tts/magpietts/config.h b/src/tts/magpietts/config.h index 0926071..db5fee1 100644 --- a/src/tts/magpietts/config.h +++ b/src/tts/magpietts/config.h @@ -14,6 +14,9 @@ class ParameterParser; namespace nemo_speech::tts { struct MagpieTtsServerConfig { + std::string backend = "magpie"; + std::string native_socket; + int native_timeout_seconds = 300; MagpieRuntimeConfig runtime; std::string tokenizer_model_dir; std::string tn_model_dir; diff --git a/src/tts/nemotron/native_runtime.cpp b/src/tts/nemotron/native_runtime.cpp new file mode 100644 index 0000000..4956e09 --- /dev/null +++ b/src/tts/nemotron/native_runtime.cpp @@ -0,0 +1,929 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#include "native_runtime.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "json.h" +#include "pcm_conversion.h" + +#ifndef _WIN32 +#include +#include +#include +#include +#include +#include + +#include + +#include "native_worker_protocol.h" +#endif + +namespace nemo_speech::tts::nemotron { +namespace { + +using Clock = std::chrono::steady_clock; +using Value = json::Value; +constexpr size_t kMaxFrameBytes = 16 * 1024 * 1024; +constexpr size_t kMaxQueuedBytes = 4 * 1024 * 1024; +constexpr size_t kMaxQueuedChunks = 256; +constexpr size_t kMaxPendingRequests = 1024; +constexpr uint64_t kMaxJsonInteger = 9007199254740991ULL; + +uint64_t +integer(const Value& value, const std::string& key, uint64_t maximum = kMaxJsonInteger) { + const auto& field = value.at(key); + if (!field.is_number() || !std::isfinite(field.number()) || field.number() < 0 || + std::floor(field.number()) != field.number() || field.number() > maximum) + throw std::runtime_error("invalid native worker integer: " + key); + return static_cast(field.number()); +} + +Value +control(const char* kind, uint64_t request_id) { + Value value(Value::Object{}); + value["kind"] = kind; + value["request_id"] = static_cast(request_id); + return value; +} + +uint64_t +read_be(const char* bytes, size_t count) { + uint64_t value = 0; + for (size_t i = 0; i < count; ++i) value = (value << 8) | static_cast(bytes[i]); + return value; +} + +double +milliseconds(Clock::duration duration) { + return std::chrono::duration(duration).count(); +} + +#ifndef _WIN32 +class TextPreparation { + public: + explicit TextPreparation(const Value& metadata) { + if (!metadata.is_object() || metadata.object().size() != 3) + throw std::runtime_error("invalid native text preparation metadata"); + const auto library = metadata.at("library").string(); + const auto tokenizer = metadata.at("tokenizer").string(); + const auto config = metadata.at("normalization_config").string(); + for (const auto* path : {&library, &tokenizer}) { + if (path->empty() || path->front() != '/' || path->find('\0') != std::string::npos) + throw std::runtime_error("native preparation paths must be absolute"); + } + if (config.empty() || config.size() > 1024 * 1024 || config.find('\0') != std::string::npos) + throw std::runtime_error("invalid native preparation configuration"); + library_ = ::dlopen(library.c_str(), RTLD_NOW | RTLD_LOCAL); + if (!library_) + throw std::runtime_error("cannot load native text preparation library"); + try { + auto create = symbol("nemotron_text_create"); + prepare_ = symbol("nemotron_text_prepare"); + free_ = symbol("nemotron_text_free"); + destroy_ = symbol("nemotron_text_destroy"); + char* error = nullptr; + const int status = create(config.c_str(), tokenizer.c_str(), &handle_, &error); + std::unique_ptr owned_error(error, free_); + check(status, error); + if (!handle_) + throw std::runtime_error("native preparation returned an empty handle"); + } + catch (...) { + release(); + throw; + } + } + + ~TextPreparation() { release(); } + TextPreparation(const TextPreparation&) = delete; + TextPreparation& operator=(const TextPreparation&) = delete; + + Value prepare(const std::string& input, bool normalize) const { + char* output = nullptr; + char* error = nullptr; + const int status = prepare_( + handle_, reinterpret_cast(input.data()), input.size(), normalize, + &output, &error); + std::unique_ptr owned_output(output, free_); + std::unique_ptr owned_error(error, free_); + check(status, error); + if (!output || ::strnlen(output, kMaxFrameBytes) >= kMaxFrameBytes) + throw std::runtime_error("native preparation returned an invalid result"); + auto result = Value::parse(output); + const auto& text = result.at("text").string(); + const auto& tokens = result.at("tokens").array(); + if (text.empty() || text.size() > kMaxFrameBytes || tokens.empty() || tokens.size() > 512) + throw std::runtime_error("native preparation returned invalid text or token count"); + for (const auto& token : tokens) { + if (!token.is_number() || !std::isfinite(token.number()) || token.number() < 0 || + token.number() >= 131075 || token.number() != std::floor(token.number())) + throw std::runtime_error("native preparation returned an invalid token ID"); + } + if (tokens.back().number() != 131073) + throw std::runtime_error("native preparation omitted the text EOS token"); + return result; + } + + private: + using Create = int (*)(const char*, const char*, void**, char**); + using Prepare = int (*)(void*, const uint8_t*, size_t, uint8_t, char**, char**); + using Free = void (*)(char*); + using Destroy = void (*)(void*); + + template + Function symbol(const char* name) { + auto address = ::dlsym(library_, name); + if (!address) + throw std::runtime_error(std::string("native preparation missing symbol: ") + name); + return reinterpret_cast(address); + } + + static void check(int status, const char* error) { + if (status == 0) + return; + const std::string message = error ? error : "native text preparation failed"; + if (status == 1) + throw std::invalid_argument(message); + throw std::runtime_error(message); + } + + void release() { + if (handle_ && destroy_) + destroy_(handle_); + handle_ = nullptr; + if (library_) + ::dlclose(library_); + library_ = nullptr; + } + + void* library_ = nullptr; + void* handle_ = nullptr; + Prepare prepare_ = nullptr; + Free free_ = nullptr; + Destroy destroy_ = nullptr; +}; +#endif + +} // namespace + +#ifndef _WIN32 + +struct NativeRuntime::Impl { + using Message = worker_protocol::Message; + + struct Request { + std::mutex mutex; + std::condition_variable changed; + std::deque chunks; + size_t queued_bytes = 0; + uint64_t received_samples = 0; + uint64_t received_chunks = 0; + bool accepted = false; + bool complete = false; + bool discard = false; + std::string error; + std::string processed_text; + size_t token_count = 0; + }; + + explicit Impl(const std::string& socket_path, int seconds) : timeout_seconds(seconds) { + if (seconds <= 0) + throw std::invalid_argument("native_timeout_seconds must be positive"); + if (socket_path.empty()) + throw std::invalid_argument("Nemotron backend requires native_socket"); + if (const char* config = std::getenv("NEMOTRON_DIRECT_WORKER_CONFIG")) + configure_direct(config); + sockaddr_un address{}; + address.sun_family = AF_UNIX; + if (socket_path.size() >= sizeof(address.sun_path)) + throw std::invalid_argument("native_socket path is too long"); + std::memcpy(address.sun_path, socket_path.c_str(), socket_path.size() + 1); + fd = ::socket(AF_UNIX, SOCK_STREAM, 0); + if (fd < 0) + throw std::runtime_error("cannot create native worker socket"); + try { + if (::fcntl(fd, F_SETFD, FD_CLOEXEC) < 0 || ::fcntl(fd, F_SETFL, O_NONBLOCK) < 0) + throw std::runtime_error("cannot configure native worker socket"); +#ifdef SO_NOSIGPIPE + int enabled = 1; + if (::setsockopt(fd, SOL_SOCKET, SO_NOSIGPIPE, &enabled, sizeof(enabled)) < 0) + throw std::runtime_error("cannot configure native worker signal handling"); +#endif + const auto deadline = Clock::now() + std::chrono::seconds(timeout_seconds); + if (::connect(fd, reinterpret_cast(&address), sizeof(address)) < 0) { + if (errno != EINPROGRESS && errno != EAGAIN) + throw std::runtime_error("cannot connect to native worker: " + socket_path); + wait_socket(POLLOUT, deadline); + int error = 0; + socklen_t size = sizeof(error); + if (::getsockopt(fd, SOL_SOCKET, SO_ERROR, &error, &size) < 0 || error != 0) + throw std::runtime_error("cannot connect to native worker: " + socket_path); + } + if (direct) { + send_direct( + {{"kind", "client_hello"}, + {"expected_manifest_sha256", direct_config.at("manifest_sha256")}, + {"client_instance_id", "nemo-speech-cpp-direct-" + request_prefix}}, + deadline); + validate_direct_ready(worker_protocol::unpack(read_frame(deadline))); + } else { + Value hello(Value::Object{}); + hello["kind"] = "hello"; + hello["protocol_version"] = 1; + send(hello, deadline); + const auto frame = read_frame(deadline); + if (frame.front() != 'J') + throw std::runtime_error("native worker did not send ready metadata"); + const auto ready = Value::parse(frame.substr(1)); + if (ready.string_or("kind") != "ready") + throw std::runtime_error( + "native worker handshake failed: " + + ready.string_or("message", "not ready")); + name = ready.at("model_name").string(); + if (const auto* version = ready.find("protocol_version"); + version && integer(ready, "protocol_version") != 1) + throw std::runtime_error("native worker protocol version is unsupported"); + if (name != "nemotron-tts" || integer(ready, "sample_rate") != 22050) + throw std::runtime_error("native worker model or sample rate is unsupported"); + for (const auto& voice : ready.at("voices").array()) + voices.push_back(voice.string()); + if (voices != std::vector{"english_sample"}) + throw std::runtime_error("native worker voices are unsupported"); + width = static_cast(integer(ready, "batch_size", 4096)); + if (width < 1) + throw std::runtime_error("native worker batch_size must be positive"); + execution_mode = ready.at("mode").string(); + if (execution_mode.empty()) + throw std::runtime_error("native worker mode is empty"); + normalizes = ready.at("text_normalization_enabled").boolean(); + if (const auto* preparation = ready.find("text_preparation")) + text_preparation = std::make_unique(*preparation); + } + reader = std::thread([this] { read_loop(); }); + } + catch (...) { + ::close(fd); + fd = -1; + throw; + } + } + + ~Impl() { + fail_connection("native worker connection closed"); + if (reader.joinable()) + reader.join(); + if (fd >= 0) + ::close(fd); + } + + static std::string hex_id(uint64_t value) { + std::string result(16, '0'); + constexpr char digits[] = "0123456789abcdef"; + for (size_t index = 0; index < result.size(); ++index) { + result[15 - index] = digits[value & 15]; + value >>= 4; + } + return result; + } + + void configure_direct(const std::string& path) { + direct_config = worker_protocol::load_config(path); + const std::set required = { + "schema", + "width", + "mode", + "manifest_sha256", + "model_sha256", + "output_contract", + "enabled_buckets", + "installed_transport_versions", + "text_normalization_enabled", + "text_preparation"}; + std::set found; + for (auto it = direct_config.begin(); it != direct_config.end(); ++it) + found.insert(it.key()); + if (found != required || + worker_protocol::string(direct_config, "schema") != "nemotron.direct_worker.v1") + throw std::runtime_error("invalid direct native worker configuration schema"); + width = static_cast(worker_protocol::integer(direct_config, "width", 4096)); + if (width < 1 || (width & (width - 1)) != 0) + throw std::runtime_error("invalid direct native worker width"); + execution_mode = worker_protocol::string(direct_config, "mode"); + if (execution_mode != "fixed_cohort" && execution_mode != "fixed_request" && + execution_mode != "inflight") + throw std::runtime_error("invalid direct native worker mode"); + for (const char* key : {"manifest_sha256", "model_sha256"}) { + const auto digest = worker_protocol::string(direct_config, key, 64); + if (digest.size() != 64 || + digest.find_first_not_of("0123456789abcdef") != std::string::npos) + throw std::runtime_error("invalid direct native worker identity hash"); + } + if (worker_protocol::string(direct_config, "output_contract") != + "pcm_f32_le_22050_mono_variable_streaming_v3") + throw std::runtime_error("unsupported direct native output contract"); + const Message expected_buckets = + width == 1 ? Message::array({1}) : Message::array({1, width}); + if (direct_config.at("enabled_buckets") != expected_buckets) + throw std::runtime_error("invalid direct native worker buckets"); + const auto& versions = direct_config.at("installed_transport_versions"); + const std::set version_keys = { + "trt_version", "trtllm_version", "trtllm_commit"}; + std::set actual_keys; + if (!versions.is_object()) + throw std::runtime_error("invalid direct native transport versions"); + for (auto it = versions.begin(); it != versions.end(); ++it) { + actual_keys.insert(it.key()); + if (!it.value().is_string() || it.value().get_ref().empty() || + it.value().get_ref().size() > 1024) + throw std::runtime_error("invalid direct native transport version"); + } + if (actual_keys != version_keys || + !direct_config.at("text_normalization_enabled").is_boolean()) + throw std::runtime_error("invalid direct native transport configuration"); + const auto metadata = Value::parse(direct_config.at("text_preparation").dump()); + text_preparation = std::make_unique(metadata); + const auto normalizer = Message::parse(metadata.at("normalization_config").string()); + normalizes = !normalizer.is_null() && normalizer.value("enabled", true); + if (normalizes != direct_config.at("text_normalization_enabled").get()) + throw std::runtime_error("direct native normalization configuration differs"); + name = "nemotron-tts"; + voices = {"english_sample"}; + std::random_device random; + request_prefix = hex_id((static_cast(random()) << 32) | random()); + direct = true; + } + + void validate_direct_ready(const Message& ready) { + if (worker_protocol::integer(ready, "protocol_version") != 1 || + worker_protocol::string(ready, "kind") != "worker_ready") + throw std::runtime_error("invalid direct native worker handshake"); + for (const char* key : + {"manifest_sha256", "model_sha256", "output_contract", "enabled_buckets"}) + if (ready.at(key) != direct_config.at(key)) + throw std::runtime_error( + std::string("direct native worker identity mismatch: ") + key); + for (auto it = direct_config.at("installed_transport_versions").begin(); + it != direct_config.at("installed_transport_versions").end(); ++it) + if (ready.at(it.key()) != it.value()) + throw std::runtime_error("direct native worker transport version mismatch"); + } + + uint64_t direct_id(const Message& value) const { + const auto id = worker_protocol::string(value, "request_id", 32); + if (id.size() != 32 || id.compare(0, 16, request_prefix) != 0) + throw std::runtime_error("direct native worker returned an unknown request identity"); + uint64_t result = 0; + for (size_t index = 16; index < id.size(); ++index) { + const auto digit = id[index]; + if (!((digit >= '0' && digit <= '9') || (digit >= 'a' && digit <= 'f'))) + throw std::runtime_error( + "direct native worker returned an invalid request identity"); + result = (result << 4) | + static_cast(digit <= '9' ? digit - '0' : digit - 'a' + 10); + } + if (result == 0) + throw std::runtime_error("direct native worker returned a zero request identity"); + return result; + } + + void send_direct(Message value, Clock::time_point deadline) { + value["protocol_version"] = 1; + send_payload(worker_protocol::pack(value), deadline); + } + + void wait_socket(short events, Clock::time_point deadline) { + while (true) { + int timeout = -1; + if (deadline != Clock::time_point::max()) { + const auto remaining = + std::chrono::duration_cast(deadline - Clock::now()) + .count(); + if (remaining <= 0) + throw std::runtime_error("native worker socket timed out"); + timeout = static_cast(std::min(remaining, 2147483647)); + } + pollfd descriptor{fd, events, 0}; + const int result = ::poll(&descriptor, 1, timeout); + if (result < 0 && errno == EINTR) + continue; + if (result == 0) + throw std::runtime_error("native worker socket timed out"); + if (result < 0 || descriptor.revents & (POLLERR | POLLNVAL)) + throw std::runtime_error("native worker socket failed"); + if (descriptor.revents & (events | POLLHUP)) + return; + } + } + + void read_exact(char* output, size_t bytes, Clock::time_point deadline) { + while (bytes > 0) { + wait_socket(POLLIN, deadline); + const auto received = ::recv(fd, output, bytes, 0); + if (received < 0 && (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK)) + continue; + if (received <= 0) + throw std::runtime_error("native worker disconnected"); + output += received; + bytes -= static_cast(received); + } + } + + std::string read_frame(Clock::time_point deadline) { + char header[4]; + read_exact(header, sizeof(header), deadline); + const auto length = read_be(header, sizeof(header)); + if (length < 2 || length > kMaxFrameBytes) + throw std::runtime_error("invalid native worker frame length"); + std::string frame(static_cast(length), '\0'); + read_exact(frame.data(), frame.size(), deadline); + return frame; + } + + void send(const Value& value, Clock::time_point deadline) { + send_payload("J" + value.dump(), deadline); + } + + void send_payload(const std::string& payload, Clock::time_point deadline) { + if (payload.size() > kMaxFrameBytes) + throw std::invalid_argument("native worker request is too large"); + std::string frame(4, '\0'); + const auto length = static_cast(payload.size()); + for (size_t i = 0; i < 4; ++i) frame[i] = static_cast(length >> (24 - 8 * i)); + frame += payload; + std::unique_lock lock(writes, std::defer_lock); + if (!lock.try_lock_until(deadline)) + throw std::runtime_error("native worker write timed out"); + size_t offset = 0; + while (offset < frame.size()) { + wait_socket(POLLOUT, deadline); +#ifdef MSG_NOSIGNAL + constexpr int flags = MSG_NOSIGNAL; +#else + constexpr int flags = 0; +#endif + const auto written = ::send(fd, frame.data() + offset, frame.size() - offset, flags); + if (written < 0 && (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK)) + continue; + if (written <= 0) + throw std::runtime_error("native worker write failed"); + offset += static_cast(written); + } + } + + void fail_connection(const std::string& reason) { + std::lock_guard lock(pending_mutex); + if (connection_error.empty()) + connection_error = reason; + ::shutdown(fd, SHUT_RDWR); + for (const auto& entry : pending) { + auto& request = *entry.second; + std::lock_guard request_lock(request.mutex); + request.error = connection_error; + request.changed.notify_all(); + } + pending.clear(); + } + + std::shared_ptr lookup(uint64_t id) { + std::lock_guard lock(pending_mutex); + const auto found = pending.find(id); + if (found == pending.end()) { + if (direct && retired.count(id)) + return {}; + throw std::runtime_error("native worker returned an unknown request ID"); + } + return found->second; + } + + void retire(uint64_t id) { + std::lock_guard lock(pending_mutex); + pending.erase(id); + if (direct && retired.insert(id).second) { + retired_order.push_back(id); + if (retired_order.size() > 4096) { + retired.erase(retired_order.front()); + retired_order.pop_front(); + } + } + } + + void receive_audio(const std::string& frame) { + const size_t sample_bytes = frame.front() == 'F' ? 4 : 2; + if (frame.size() < 13 + sample_bytes || (frame.size() - 13) % sample_bytes != 0) + throw std::runtime_error("native worker returned malformed audio"); + const uint64_t id = read_be(frame.data() + 1, 8); + const uint64_t sequence = read_be(frame.data() + 9, 4); + auto request = lookup(id); + std::lock_guard lock(request->mutex); + if (!request->accepted || request->complete || sequence != request->received_chunks) + throw std::runtime_error("native worker audio sequence is invalid"); + const size_t samples = (frame.size() - 13) / sample_bytes; + const size_t bytes = samples * 2; + ++request->received_chunks; + request->received_samples += samples; + if (request->discard) + return; + if (bytes > kMaxQueuedBytes - request->queued_bytes || + request->chunks.size() >= kMaxQueuedChunks) { + request->error = "native worker audio consumer exceeded its bounded queue"; + request->discard = true; + request->chunks.clear(); + request->queued_bytes = 0; + } else { + request->chunks.emplace_back( + sample_bytes == 4 ? pcm_f32_to_s16(frame.data() + 13, frame.size() - 13) + : frame.substr(13)); + request->queued_bytes += bytes; + } + request->changed.notify_all(); + } + + void receive_control(const Value& value) { + if (const auto* version = value.find("protocol_version"); + version && integer(value, "protocol_version") != 1) + throw std::runtime_error("native worker protocol version is unsupported"); + const uint64_t id = integer(value, "request_id"); + const auto kind = value.at("kind").string(); + auto request = lookup(id); + bool terminal = false; + { + std::lock_guard lock(request->mutex); + if (kind == "accepted") { + if (request->accepted || request->complete) + throw std::runtime_error("native worker accepted a request twice"); + request->processed_text = value.at("processed_text").string(); + request->token_count = static_cast(integer(value, "token_count", 4096)); + if (request->processed_text.empty() || request->token_count == 0) + throw std::runtime_error("native worker accepted empty text or tokens"); + request->accepted = true; + } else if (kind == "complete") { + if (!request->discard && + (!request->accepted || request->received_samples == 0 || + integer(value, "total_samples") != request->received_samples || + integer(value, "chunk_count") != request->received_chunks)) + throw std::runtime_error("native worker completion does not match audio"); + request->complete = true; + terminal = true; + } else if (kind == "failed") { + request->error = value.at("message").string(); + if (request->error.empty()) + request->error = "native worker synthesis failed"; + request->complete = true; + terminal = true; + } else if (kind == "aborted") { + request->error = "native worker synthesis aborted"; + request->complete = true; + terminal = true; + } else { + throw std::runtime_error("unsupported native worker message: " + kind); + } + request->changed.notify_all(); + } + if (terminal) + retire(id); + } + + void receive_direct_audio(const Message& value) { + const auto slab = slabs.read(value); + auto request = lookup(direct_id(value)); + bool discard = !request; + if (request) { + std::lock_guard lock(request->mutex); + if (!request->accepted || request->complete || + slab.sequence != request->received_chunks) + throw std::runtime_error("direct native audio sequence is invalid"); + ++request->received_chunks; + request->received_samples += slab.samples; + discard = request->discard; + } + auto pcm = discard ? std::string{} : pcm_f32_to_s16(slab.data, slab.bytes); + send_direct( + {{"kind", "audio_ack"}, + {"arena", value.at("arena")}, + {"slot", value.at("slot")}, + {"generation", value.at("generation")}}, + Clock::now() + std::chrono::seconds(timeout_seconds)); + if (!request) + return; + std::lock_guard lock(request->mutex); + if (request->discard) + return; + if (pcm.size() > kMaxQueuedBytes - request->queued_bytes || + request->chunks.size() >= kMaxQueuedChunks) { + request->error = "native worker audio consumer exceeded its bounded queue"; + request->discard = true; + request->chunks.clear(); + request->queued_bytes = 0; + } else { + request->queued_bytes += pcm.size(); + request->chunks.emplace_back(std::move(pcm)); + } + request->changed.notify_all(); + } + + void receive_direct(const Message& value) { + if (worker_protocol::integer(value, "protocol_version") != 1) + throw std::runtime_error("direct native worker protocol version changed"); + const auto kind = worker_protocol::string(value, "kind", 64); + if (kind == "audio_slab") { + receive_direct_audio(value); + return; + } + const auto id = direct_id(value); + auto request = lookup(id); + if (!request) + return; + bool terminal = false; + { + std::lock_guard lock(request->mutex); + if (kind == "accepted") { + if (request->accepted || request->complete) + throw std::runtime_error("direct native request was accepted twice"); + request->accepted = true; + } else if (kind == "stream_complete_v1") { + if (!request->accepted || request->received_samples == 0 || + worker_protocol::integer(value, "next_sequence") != request->received_chunks || + worker_protocol::integer(value, "chunk_count") != request->received_chunks || + worker_protocol::integer(value, "total_samples") != request->received_samples) + throw std::runtime_error("direct native completion does not match audio"); + request->complete = true; + terminal = true; + } else if (kind == "failed" || kind == "aborted") { + request->error = kind == "failed" ? worker_protocol::string(value, "message") + : "native worker synthesis aborted"; + if (request->error.empty()) + request->error = "native worker synthesis failed"; + request->complete = true; + terminal = true; + } else { + throw std::runtime_error("unsupported direct native worker message: " + kind); + } + request->changed.notify_all(); + } + if (terminal) + retire(id); + } + + void read_loop() { + try { + while (true) { + wait_socket(POLLIN, Clock::time_point::max()); + const auto frame = read_frame(Clock::now() + std::chrono::seconds(timeout_seconds)); + if (direct) + receive_direct(worker_protocol::unpack(frame)); + else if (frame.front() == 'A' || frame.front() == 'F') + receive_audio(frame); + else if (frame.front() == 'J') + receive_control(Value::parse(frame.substr(1))); + else + throw std::runtime_error("unsupported native worker frame type"); + } + } + catch (const std::exception& error) { + fail_connection(error.what()); + } + } + + void abort(uint64_t id, const std::shared_ptr& request) { + { + std::lock_guard lock(request->mutex); + request->discard = true; + request->chunks.clear(); + request->queued_bytes = 0; + if (request->complete) + return; + } + try { + if (direct) + send_direct( + {{"kind", "abort"}, {"request_id", request_prefix + hex_id(id)}}, + Clock::now() + std::chrono::seconds(timeout_seconds)); + else + send(control("abort", id), Clock::now() + std::chrono::seconds(timeout_seconds)); + } + catch (const std::exception& error) { + fail_connection(error.what()); + } + } + + SynthesisResult synthesize( + const PreparedSynthesis& prepared, const Synthesizer::PcmCallback& callback) { + const auto started = Clock::now(); + const auto deadline = started + std::chrono::seconds(timeout_seconds); + Value prepared_text; + if (text_preparation) { + prepared_text = text_preparation->prepare( + prepared.metadata.original_text, prepared.apply_normalize != 0); + if (Clock::now() >= deadline) + throw std::runtime_error("native text preparation timed out"); + } + std::shared_ptr request; + uint64_t id; + { + std::lock_guard lock(pending_mutex); + if (!connection_error.empty()) + throw std::runtime_error(connection_error); + if (pending.size() >= std::max(kMaxPendingRequests, static_cast(width) * 2) || + next_id >= kMaxJsonInteger) + throw std::runtime_error("native worker request capacity exhausted"); + id = next_id++; + request = std::make_shared(); + if (direct) { + request->processed_text = prepared_text.at("text").string(); + request->token_count = prepared_text.at("tokens").array().size(); + } + pending.emplace(id, request); + } + Value submit = control(text_preparation ? "submit_prepared" : "submit", id); + submit["text"] = + text_preparation ? prepared_text.at("text") : Value(prepared.metadata.original_text); + if (text_preparation) + submit["token_ids"] = std::move(prepared_text["tokens"]); + submit["voice"] = "english_sample"; + submit["language"] = prepared.metadata.language_code; + submit["seed"] = prepared.options.seed < 0 ? 0 : prepared.options.seed; + if (prepared.apply_normalize >= 0) + submit["apply_normalize"] = prepared.apply_normalize != 0; + try { + if (direct) { + auto tokens = Message::array(); + for (const auto& token : submit.at("token_ids").array()) + tokens.push_back(static_cast(token.number())); + send_direct( + {{"kind", "submit"}, + {"request_id", request_prefix + hex_id(id)}, + {"session_id", nullptr}, + {"normalized_text", request->processed_text}, + {"token_ids", std::move(tokens)}, + {"speaker_id", "english_sample"}, + {"max_tokens", 512}, + {"temperature", 0.7}, + {"top_k", 80}, + {"sample_rate_hz", 22050}, + {"seed", prepared.options.seed < 0 ? 0 : prepared.options.seed}}, + deadline); + } else { + send(submit, deadline); + } + } + catch (const std::exception& error) { + fail_connection(error.what()); + throw; + } + SynthesisResult result; + result.metadata = prepared.metadata; + auto& stats = result.stats; + stats.sample_rate = 22050; + Clock::time_point last_audio{}; + std::vector gaps; + try { + while (true) { + std::string pcm; + { + std::unique_lock lock(request->mutex); + if (!request->changed.wait_until(lock, deadline, [&] { + return !request->error.empty() || !request->chunks.empty() || + request->complete; + })) + throw std::runtime_error("native worker synthesis timed out"); + if (!request->error.empty()) + throw std::runtime_error(request->error); + result.metadata.processed_text = request->processed_text; + result.metadata.token_count = request->token_count; + if (request->chunks.empty()) + break; + pcm = std::move(request->chunks.front()); + request->chunks.pop_front(); + request->queued_bytes -= pcm.size(); + } + const auto now = Clock::now(); + if (last_audio == Clock::time_point{}) + stats.ttfa_ms = stats.e2e_ttfa_ms = milliseconds(now - started); + else + gaps.push_back(milliseconds(now - last_audio)); + last_audio = now; + result.output_samples += pcm.size() / 2; + ++stats.chunks; + if (callback && !callback(result.metadata, pcm)) { + result.cancelled = true; + abort(id, request); + break; + } + } + } + catch (...) { + abort(id, request); + throw; + } + stats.e2e_chunks = stats.chunks; + stats.samples_written = result.output_samples; + stats.audio_s = static_cast(result.output_samples) / 22050; + stats.elapsed_s = std::chrono::duration(Clock::now() - started).count(); + stats.rtf = stats.audio_s > 0 ? stats.elapsed_s / stats.audio_s : 0; + stats.rtfx = stats.e2e_rtfx = stats.elapsed_s > 0 ? stats.audio_s / stats.elapsed_s : 0; + if (!gaps.empty()) { + double sum = 0; + for (const double gap : gaps) sum += gap; + std::sort(gaps.begin(), gaps.end()); + stats.icl_avg_ms = stats.e2e_icl_avg_ms = sum / gaps.size(); + stats.icl_min_ms = stats.e2e_icl_min_ms = gaps.front(); + stats.icl_max_ms = stats.e2e_icl_max_ms = gaps.back(); + stats.e2e_icl_p95_ms = gaps[static_cast(std::ceil(gaps.size() * 0.95)) - 1]; + stats.e2e_icl_p99_ms = gaps[static_cast(std::ceil(gaps.size() * 0.99)) - 1]; + } + return result; + } + + bool direct = false; + Message direct_config; + std::string request_prefix; + worker_protocol::SlabReader slabs; + std::set retired; + std::deque retired_order; + int fd = -1; + int timeout_seconds; + std::string name; + std::vector voices; + bool normalizes = false; + int width = 1; + std::string execution_mode; + std::unique_ptr text_preparation; + std::thread reader; + std::timed_mutex writes; + std::mutex pending_mutex; + std::map> pending; + std::string connection_error; + uint64_t next_id = 1; +}; + +#else + +struct NativeRuntime::Impl { + Impl(const std::string&, int) { + throw std::runtime_error("Nemotron native worker sockets require a Unix platform"); + } + SynthesisResult synthesize(const PreparedSynthesis&, const Synthesizer::PcmCallback&) { + throw std::runtime_error("Nemotron native worker sockets require a Unix platform"); + } + std::string name; + std::vector voices; + bool normalizes = false; + int width = 1; + std::string execution_mode; +}; + +#endif + +NativeRuntime::NativeRuntime(const std::string& socket_path, int timeout_seconds) + : impl_(std::make_unique(socket_path, timeout_seconds)) {} + +NativeRuntime::~NativeRuntime() = default; + +SynthesisResult +NativeRuntime::synthesize( + const PreparedSynthesis& request, const Synthesizer::PcmCallback& callback) { + return impl_->synthesize(request, callback); +} + +const std::string& +NativeRuntime::model_name() const { + return impl_->name; +} + +const std::vector& +NativeRuntime::speaker_names() const { + return impl_->voices; +} + +bool +NativeRuntime::text_normalization_enabled() const { + return impl_->normalizes; +} + +int +NativeRuntime::batch_size() const { + return impl_->width; +} + +const std::string& +NativeRuntime::mode() const { + return impl_->execution_mode; +} + +} // namespace nemo_speech::tts::nemotron diff --git a/src/tts/nemotron/native_runtime.h b/src/tts/nemotron/native_runtime.h new file mode 100644 index 0000000..0e42f9b --- /dev/null +++ b/src/tts/nemotron/native_runtime.h @@ -0,0 +1,33 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include +#include +#include + +#include "tts/synthesizer.h" + +namespace nemo_speech::tts::nemotron { + +class NativeRuntime { + public: + NativeRuntime(const std::string& socket_path, int timeout_seconds); + ~NativeRuntime(); + NativeRuntime(const NativeRuntime&) = delete; + NativeRuntime& operator=(const NativeRuntime&) = delete; + + SynthesisResult synthesize( + const PreparedSynthesis& request, const Synthesizer::PcmCallback& callback); + const std::string& model_name() const; + const std::vector& speaker_names() const; + bool text_normalization_enabled() const; + int batch_size() const; + const std::string& mode() const; + + private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace nemo_speech::tts::nemotron diff --git a/src/tts/nemotron/native_worker_protocol.cpp b/src/tts/nemotron/native_worker_protocol.cpp new file mode 100644 index 0000000..8ca0cf2 --- /dev/null +++ b/src/tts/nemotron/native_worker_protocol.cpp @@ -0,0 +1,260 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#include "native_worker_protocol.h" + +#ifndef _WIN32 +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace nemo_speech::tts::nemotron::worker_protocol { +namespace { + +bool +valid_utf8(const std::string& text) { + size_t offset = 0; + while (offset < text.size()) { + const auto first = static_cast(text[offset++]); + if (first < 0x80) + continue; + const int count = first >= 0xc2 && first <= 0xdf ? 1 + : first >= 0xe0 && first <= 0xef ? 2 + : first >= 0xf0 && first <= 0xf4 ? 3 + : 0; + if (count == 0 || text.size() - offset < static_cast(count)) + return false; + const auto second = static_cast(text[offset]); + if ((first == 0xe0 && second < 0xa0) || (first == 0xed && second >= 0xa0) || + (first == 0xf0 && second < 0x90) || (first == 0xf4 && second >= 0x90)) + return false; + for (int i = 0; i < count; ++i) + if ((static_cast(text[offset++]) & 0xc0) != 0x80) + return false; + } + return true; +} + +class BoundedSax final : public nlohmann::json_sax { + public: + Message result; + std::string error; + + bool null() override { return append(nullptr); } + bool boolean(bool value) override { return append(value); } + bool number_integer(number_integer_t value) override { return append(value); } + bool number_unsigned(number_unsigned_t value) override { return append(value); } + bool number_float(number_float_t value, const string_t&) override { + return std::isfinite(value) ? append(value) : fail("nonfinite native protocol number"); + } + bool string(string_t& value) override { + return valid_utf8(value) ? append(std::move(value)) : fail("invalid native protocol UTF-8"); + } + bool binary(binary_t& value) override { + return value.has_subtype() ? fail("native protocol extensions are unsupported") + : append(Message::binary(std::move(value))); + } + bool start_object(std::size_t elements) override { return start(true, elements); } + bool start_array(std::size_t elements) override { return start(false, elements); } + bool key(string_t& value) override { + if (stack_.empty() || stack_.size() > 64 || !stack_.back().value.is_object() || + !valid_utf8(value)) + return fail("invalid native protocol object key"); + auto& frame = stack_.back(); + if (frame.value.size() >= 16384 || frame.value.contains(value)) + return fail("duplicate or excessive native protocol object keys"); + frame.key = std::move(value); + frame.pending_key = true; + return true; + } + bool end_object() override { return finish(true); } + bool end_array() override { return finish(false); } + bool parse_error(std::size_t, const std::string&, const nlohmann::detail::exception&) override { + return fail("malformed native protocol payload"); + } + + private: + struct Frame { + Message value; + std::string key; + bool pending_key = false; + }; + std::vector stack_; + bool complete_ = false; + + bool fail(const char* reason) { + if (error.empty()) + error = reason; + return false; + } + bool append(Message value) { + if (stack_.size() > 64) + return fail("native protocol nesting limit exceeded"); + if (stack_.empty()) { + if (complete_) + return fail("multiple native protocol roots"); + result = std::move(value); + complete_ = true; + return true; + } + auto& frame = stack_.back(); + if (frame.value.is_object()) { + if (!frame.pending_key) + return fail("missing native protocol object key"); + frame.value.emplace(std::move(frame.key), std::move(value)); + frame.pending_key = false; + } else { + if (frame.value.size() >= 16384) + return fail("excessive native protocol array entries"); + frame.value.push_back(std::move(value)); + } + return true; + } + bool start(bool object, size_t elements) { + if (stack_.size() > 64 || (elements != size_t(-1) && elements > 16384)) + return fail("native protocol collection limit exceeded"); + stack_.push_back({object ? Message::object() : Message::array(), {}, false}); + return true; + } + bool finish(bool object) { + if (stack_.empty() || stack_.back().value.is_object() != object || + stack_.back().pending_key) + return fail("invalid native protocol collection"); + auto value = std::move(stack_.back().value); + stack_.pop_back(); + return append(std::move(value)); + } +}; + +Message +parse(std::string_view payload, Message::input_format_t format) { + BoundedSax sax; + if (!Message::sax_parse(payload.begin(), payload.end(), &sax, format, true) || + !sax.result.is_object()) + throw std::runtime_error( + sax.error.empty() ? "native protocol root must be a map" : sax.error); + return std::move(sax.result); +} + +} + +Message +unpack(std::string_view payload) { + if (payload.empty() || payload.size() > 16 * 1024 * 1024) + throw std::runtime_error("invalid native MessagePack payload size"); + return parse(payload, Message::input_format_t::msgpack); +} + +Message +load_config(const std::string& path) { + if (path.empty() || path.front() != '/' || path.find('\0') != std::string::npos) + throw std::runtime_error("direct native worker configuration path must be absolute"); + std::ifstream input(path, std::ios::binary); + if (!input) + throw std::runtime_error("cannot read direct native worker configuration"); + std::string content(1048577, '\0'); + input.read(content.data(), static_cast(content.size())); + const auto bytes = input.gcount(); + if (input.bad() || bytes < 2 || bytes > 1048576) + throw std::runtime_error("invalid direct native worker configuration size"); + content.resize(static_cast(bytes)); + return parse(content, Message::input_format_t::json); +} + +std::string +pack(const Message& value) { + const auto bytes = Message::to_msgpack(value); + if (bytes.empty() || bytes.size() > 16 * 1024 * 1024) + throw std::runtime_error("invalid native MessagePack payload size"); + return std::string(reinterpret_cast(bytes.data()), bytes.size()); +} + +uint64_t +integer(const Message& value, const char* key, uint64_t maximum) { + const auto& field = value.at(key); + uint64_t result; + if (field.is_number_unsigned()) + result = field.get(); + else if (field.is_number_integer() && field.get() >= 0) + result = static_cast(field.get()); + else + throw std::runtime_error(std::string("invalid native worker integer: ") + key); + if (result > maximum) + throw std::runtime_error(std::string("native worker integer exceeds limit: ") + key); + return result; +} + +std::string +string(const Message& value, const char* key, size_t maximum) { + const auto& field = value.at(key); + if (!field.is_string()) + throw std::runtime_error(std::string("invalid native worker string: ") + key); + const auto& result = field.get_ref(); + if (result.size() > maximum || !valid_utf8(result)) + throw std::runtime_error(std::string("native worker string exceeds contract: ") + key); + return result; +} + +SlabReader::~SlabReader() { + if (mapping_) + ::munmap(const_cast(mapping_), extent_); +} + +Slab +SlabReader::read(const Message& value) { + constexpr uint64_t slot_bytes = 8 * 1764 * 4; + const auto arena = string(value, "arena", 40); + const auto extent = integer(value, "arena_bytes", 4096 * slot_bytes); + const auto slot = integer(value, "slot", 4095); + const auto generation = integer(value, "generation"); + const auto offset = integer(value, "offset"); + const auto bytes = integer(value, "data_bytes", slot_bytes); + const auto samples = integer(value, "samples", slot_bytes / 4); + const auto sequence = integer(value, "sequence", UINT32_MAX); + bool valid_name = arena.size() == 40 && arena.compare(0, 8, "r20-pcm-") == 0; + for (size_t index = 8; valid_name && index < arena.size(); ++index) + valid_name = (arena[index] >= '0' && arena[index] <= '9') || + (arena[index] >= 'a' && arena[index] <= 'f'); + const auto previous = generations_.find(slot); + if (!valid_name || integer(value, "slab_version") != 1 || + integer(value, "slot_bytes") != slot_bytes || extent == 0 || extent % slot_bytes || + slot >= extent / slot_bytes || offset != slot * slot_bytes || bytes == 0 || + bytes != samples * 4 || offset > extent || bytes > extent - offset || generation == 0 || + (previous != generations_.end() && generation <= previous->second) || + integer(value, "sample_rate_hz") != 22050 || integer(value, "channels") != 1 || + string(value, "encoding") != "pcm_f32_le" || !value.at("is_final").is_boolean() || + value.at("is_final").get()) + throw std::runtime_error("invalid native PCM slab metadata"); + if (!mapping_) { + const std::string path = "/dev/shm/" + arena; + const int fd = ::open(path.c_str(), O_RDONLY | O_CLOEXEC | O_NOFOLLOW); + if (fd < 0) + throw std::runtime_error("cannot open native PCM arena"); + struct stat info {}; + if (::fstat(fd, &info) < 0 || !S_ISREG(info.st_mode) || info.st_size < 0 || + static_cast(info.st_size) != extent) { + ::close(fd); + throw std::runtime_error("native PCM arena extent differs from metadata"); + } + void* mapping = ::mmap(nullptr, static_cast(extent), PROT_READ, MAP_SHARED, fd, 0); + ::close(fd); + if (mapping == MAP_FAILED) + throw std::runtime_error("cannot map native PCM arena"); + mapping_ = static_cast(mapping); + extent_ = static_cast(extent); + arena_ = arena; + } + if (arena != arena_ || extent != extent_) + throw std::runtime_error("native PCM arena identity changed during a connection"); + generations_[slot] = generation; + return {mapping_ + static_cast(offset), static_cast(bytes), samples, sequence}; +} + +} +#endif diff --git a/src/tts/nemotron/native_worker_protocol.h b/src/tts/nemotron/native_worker_protocol.h new file mode 100644 index 0000000..373ea71 --- /dev/null +++ b/src/tts/nemotron/native_worker_protocol.h @@ -0,0 +1,45 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#ifndef _WIN32 +#include +#include +#include +#include +#include + +namespace nemo_speech::tts::nemotron::worker_protocol { + +using Message = nlohmann::json; + +Message unpack(std::string_view payload); +Message load_config(const std::string& path); +std::string pack(const Message& value); +uint64_t integer(const Message& value, const char* key, uint64_t maximum = UINT64_MAX); +std::string string(const Message& value, const char* key, size_t maximum = 1048576); + +struct Slab { + const char* data; + size_t bytes; + uint64_t samples; + uint64_t sequence; +}; + +class SlabReader { + public: + SlabReader() = default; + ~SlabReader(); + SlabReader(const SlabReader&) = delete; + SlabReader& operator=(const SlabReader&) = delete; + Slab read(const Message& value); + + private: + const char* mapping_ = nullptr; + size_t extent_ = 0; + std::string arena_; + std::map generations_; +}; + +} +#endif diff --git a/src/tts/nemotron/pcm_conversion.h b/src/tts/nemotron/pcm_conversion.h new file mode 100644 index 0000000..ecd5fe8 --- /dev/null +++ b/src/tts/nemotron/pcm_conversion.h @@ -0,0 +1,108 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include +#include +#include +#include +#include +#include + +#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__)) +#include +#define NEMO_NATIVE_PCM_AVX2 1 +#endif + +namespace nemo_speech::tts::nemotron { + +inline void +pcm_f32_scalar(const char* input, char* output, size_t count) { + static_assert(sizeof(float) == 4 && std::numeric_limits::is_iec559); + for (size_t index = 0; index < count; ++index) { + const auto* bytes = reinterpret_cast(input + index * 4); + const uint32_t bits = uint32_t(bytes[0]) | (uint32_t(bytes[1]) << 8) | + (uint32_t(bytes[2]) << 16) | (uint32_t(bytes[3]) << 24); + float sample; + std::memcpy(&sample, &bits, sizeof(sample)); + if (!std::isfinite(sample)) + throw std::runtime_error("native worker returned nonfinite PCM-F32LE audio"); + int value; + if (sample <= -1.0f) + value = -32768; + else if (sample >= 1.0f) + value = 32767; + else { + const float scaled = sample * (sample < 0.0f ? 32768.0f : 32767.0f); + value = static_cast(std::round(scaled)); + } + const auto encoded = static_cast(value); + output[index * 2] = static_cast(encoded & 255); + output[index * 2 + 1] = static_cast(encoded >> 8); + } +} + +inline std::string +pcm_f32_to_s16_scalar(const char* input, size_t bytes) { + if (bytes % 4 != 0) + throw std::runtime_error("native worker returned malformed PCM-F32LE audio"); + std::string output(bytes / 2, '\0'); + pcm_f32_scalar(input, output.data(), bytes / 4); + return output; +} + +#ifdef NEMO_NATIVE_PCM_AVX2 +__attribute__((target("avx2"))) inline void +pcm_f32_avx2(const char* input, char* output, size_t count) { + const __m256 zero = _mm256_setzero_ps(); + const __m256 one = _mm256_set1_ps(1.0f); + const __m256 negative_one = _mm256_set1_ps(-1.0f); + const __m256i exponent_mask = _mm256_set1_epi32(0x7f800000); + const __m256 sign_mask = _mm256_castsi256_ps(_mm256_set1_epi32(-2147483647 - 1)); + size_t index = 0; + for (; index + 8 <= count; index += 8) { + __m256 samples; + std::memcpy(&samples, input + index * 4, sizeof(samples)); + const __m256i exponents = _mm256_and_si256(_mm256_castps_si256(samples), exponent_mask); + if (_mm256_movemask_epi8(_mm256_cmpeq_epi32(exponents, exponent_mask)) != 0) + throw std::runtime_error("native worker returned nonfinite PCM-F32LE audio"); + samples = _mm256_max_ps(negative_one, _mm256_min_ps(one, samples)); + const __m256 negative = _mm256_cmp_ps(samples, zero, _CMP_LT_OQ); + const __m256 scale = + _mm256_blendv_ps(_mm256_set1_ps(32767.0f), _mm256_set1_ps(32768.0f), negative); + const __m256 scaled = _mm256_mul_ps(samples, scale); + const __m256 magnitude = _mm256_andnot_ps(sign_mask, scaled); + const __m256 integral = _mm256_round_ps(magnitude, _MM_FROUND_TO_ZERO | _MM_FROUND_NO_EXC); + const __m256 increment = _mm256_and_ps( + _mm256_cmp_ps(_mm256_sub_ps(magnitude, integral), _mm256_set1_ps(0.5f), _CMP_GE_OQ), + one); + const __m256 rounded = + _mm256_or_ps(_mm256_add_ps(integral, increment), _mm256_and_ps(scaled, sign_mask)); + const __m256i integers = _mm256_cvttps_epi32(rounded); + const __m128i packed = _mm_packs_epi32( + _mm256_castsi256_si128(integers), _mm256_extracti128_si256(integers, 1)); + std::memcpy(output + index * 2, &packed, sizeof(packed)); + } + pcm_f32_scalar(input + index * 4, output + index * 2, count - index); +} +#endif + +inline std::string +pcm_f32_to_s16(const char* input, size_t bytes) { + if (bytes % 4 != 0) + throw std::runtime_error("native worker returned malformed PCM-F32LE audio"); + std::string output(bytes / 2, '\0'); +#ifdef NEMO_NATIVE_PCM_AVX2 + if (__builtin_cpu_supports("avx2")) + pcm_f32_avx2(input, output.data(), bytes / 4); + else +#endif + pcm_f32_scalar(input, output.data(), bytes / 4); + return output; +} + +} // namespace nemo_speech::tts::nemotron + +#ifdef NEMO_NATIVE_PCM_AVX2 +#undef NEMO_NATIVE_PCM_AVX2 +#endif diff --git a/src/tts/synthesizer.cpp b/src/tts/synthesizer.cpp index 01accf9..b7c59f0 100644 --- a/src/tts/synthesizer.cpp +++ b/src/tts/synthesizer.cpp @@ -5,12 +5,14 @@ #include #include #include +#include #include #include #include #include #include "audio_resampler.h" +#include "nemotron/native_runtime.h" #include "tts/preproc/text_normalizer.h" namespace nemo_speech::tts { @@ -48,47 +50,113 @@ add_preparation_timing(MagpieSynthesisStats& stats, const PreparedSynthesis& req stats.e2e_rtfx = stats.rtfx; } +MagpieSynthesisOptions +native_options(MagpieSynthesisOptions options, const MagpieRuntimeConfig& defaults) { + if (options.seed < 0) + options.seed = defaults.seed < 0 ? 0 : defaults.seed; + if (options.steps <= 0) + options.steps = defaults.steps; + if (options.top_k <= 0) + options.top_k = defaults.top_k; + if (!options.override_temperature && defaults.override_temperature) { + options.temperature = defaults.temperature; + options.override_temperature = true; + } + if (options.steps > 0 && options.steps != 512) + throw std::invalid_argument("Nemotron requires steps=512"); + if (options.top_k > 0 && options.top_k != 80) + throw std::invalid_argument("Nemotron requires top_k=80"); + if (options.override_temperature && + (!std::isfinite(options.temperature) || std::abs(options.temperature - 0.7f) > 1e-6f)) + throw std::invalid_argument("Nemotron requires temperature=0.7"); + if (options.override_cfg_scale || defaults.override_cfg_scale) + throw std::invalid_argument("Nemotron does not support cfg_scale"); + return options; +} + +bool +valid_utf8(const std::string& text) { + size_t offset = 0; + while (offset < text.size()) { + const auto first = static_cast(text[offset++]); + if (first < 0x80) + continue; + const int continuation = first >= 0xc2 && first <= 0xdf ? 1 + : first >= 0xe0 && first <= 0xef ? 2 + : first >= 0xf0 && first <= 0xf4 ? 3 + : 0; + if (continuation == 0 || text.size() - offset < static_cast(continuation)) + return false; + const auto second = static_cast(text[offset]); + if ((first == 0xe0 && second < 0xa0) || (first == 0xed && second >= 0xa0) || + (first == 0xf0 && second < 0x90) || (first == 0xf4 && second >= 0x90)) + return false; + for (int i = 0; i < continuation; ++i) + if ((static_cast(text[offset++]) & 0xc0) != 0x80) + return false; + } + return true; +} + } // namespace struct Synthesizer::Impl { explicit Impl(SynthesizerConfig config) - : default_language_code( + : backend(std::move(config.backend)), defaults(config.runtime), + default_language_code( config.default_language_code.empty() ? "en-US" : std::move(config.default_language_code)), default_voice_name(std::move(config.default_voice_name)), - default_speaker(config.runtime.speaker), runtime(std::move(config.runtime)), - text_normalizer(std::move(config.text_normalizer_model_dir)) { + default_speaker(config.runtime.speaker) { + if (backend == "nemotron") { + native_options({}, defaults); + const auto language = lower_ascii(default_language_code); + if (language != "en-us" && language != "en" && language != "english") + throw std::invalid_argument("Nemotron supports only English"); + default_language_code = "en-US"; + native = std::make_unique( + config.native_socket, config.native_timeout_seconds); + return; + } + if (backend != "magpie") + throw std::invalid_argument("unknown TTS backend '" + backend + "'"); + runtime = std::make_unique(std::move(config.runtime)); + text_normalizer = + std::make_unique(std::move(config.text_normalizer_model_dir)); if (!config.tokenizer_model_dir.empty()) { tokenizer = std::make_unique( std::move(config.tokenizer_model_dir), config.tokenizer); - if (tokenizer->profile_id() != runtime.tokenizer_profile() || - tokenizer->text_vocab_size() != runtime.text_vocab_size()) { + if (tokenizer->profile_id() != runtime->tokenizer_profile() || + tokenizer->text_vocab_size() != runtime->text_vocab_size()) { throw std::invalid_argument( "Magpie model/tokenizer mismatch: GGUF requires tokenizer profile '" + - runtime.tokenizer_profile() + "' with text vocabulary " + - std::to_string(runtime.text_vocab_size()) + ", but tokenizer directory is '" + + runtime->tokenizer_profile() + "' with text vocabulary " + + std::to_string(runtime->text_vocab_size()) + ", but tokenizer directory is '" + tokenizer->profile_id() + "' with text vocabulary " + std::to_string(tokenizer->text_vocab_size())); } } } + std::string backend; + MagpieRuntimeConfig defaults; std::string default_language_code; std::string default_voice_name; int default_speaker = 0; - MagpieTtsRuntime runtime; + std::unique_ptr runtime; + std::unique_ptr native; std::unique_ptr tokenizer; - preproc::TextNormalizer text_normalizer; + std::unique_ptr text_normalizer; }; Synthesizer::Synthesizer(SynthesizerConfig config) : impl_(std::make_unique(std::move(config))) { if (!impl_->default_voice_name.empty()) impl_->default_speaker = resolve_speaker(impl_->default_voice_name); - if (impl_->default_speaker < 0 || impl_->default_speaker >= impl_->runtime.speaker_count()) { + if (impl_->default_speaker < 0 || impl_->default_speaker >= speaker_count()) { throw std::invalid_argument( "default speaker " + std::to_string(impl_->default_speaker) + " must be in [0, " + - std::to_string(impl_->runtime.speaker_count() - 1) + "]"); + std::to_string(speaker_count() - 1) + "]"); } } @@ -103,7 +171,7 @@ Synthesizer::resolve_speaker(const std::string& voice_name) const { try { const int value = std::stoi(voice_name, &consumed); if (consumed == voice_name.size()) { - if (value >= 0 && value < impl_->runtime.speaker_count()) + if (value >= 0 && value < speaker_count()) return value; throw std::invalid_argument("voice speaker index is outside the model speaker range"); } @@ -117,10 +185,10 @@ Synthesizer::resolve_speaker(const std::string& voice_name) const { } std::string wanted = lower_ascii(voice_name); - const std::string model_prefix = lower_ascii(impl_->runtime.model_name()) + "."; + const std::string model_prefix = lower_ascii(model_name()) + "."; if (wanted.rfind(model_prefix, 0) == 0) wanted.erase(0, model_prefix.size()); - const auto& names = impl_->runtime.speaker_names(); + const auto& names = speaker_names(); for (size_t i = 0; i < names.size(); ++i) { if (lower_ascii(names[i]) == wanted) return static_cast(i); @@ -132,6 +200,36 @@ PreparedSynthesis Synthesizer::prepare(const SynthesisRequest& request) const { if (request.text.empty()) throw std::invalid_argument("text is required"); + if (impl_->native) { + if (request.text.find_first_not_of(" \t\r\n") == std::string::npos) + throw std::invalid_argument("text is required"); + if (request.text.size() > 1000000 || !valid_utf8(request.text)) + throw std::invalid_argument( + "Nemotron text must be valid UTF-8 and at most 1000000 bytes"); + if (request.output_sample_rate != 0 && request.output_sample_rate != 22050) + throw std::invalid_argument("Nemotron requires sample_rate=22050"); + if (request.apply_normalize < -1 || request.apply_normalize > 1) + throw std::invalid_argument("apply_normalize must be -1, 0, or 1"); + const auto language = lower_ascii( + request.language_code.empty() ? impl_->default_language_code : request.language_code); + if (language != "en" && language != "en-us" && language != "english") + throw std::invalid_argument("Nemotron supports only English"); + PreparedSynthesis prepared; + prepared.metadata.original_text = request.text; + prepared.metadata.language_code = "en-US"; + prepared.metadata.sample_rate = 22050; + prepared.metadata.tokenizer_name = "nemotron-native"; + prepared.metadata.chunk_count = 1; + prepared.options = native_options(request.options, impl_->defaults); + if (prepared.options.speaker < 0) + prepared.options.speaker = resolve_speaker( + request.voice_name.empty() ? impl_->default_voice_name : request.voice_name); + if (prepared.options.speaker != 0) + throw std::invalid_argument("Nemotron supports only voice english_sample"); + prepared.metadata.speaker = prepared.options.speaker; + prepared.apply_normalize = request.apply_normalize; + return prepared; + } if (!impl_->tokenizer) { throw std::invalid_argument("text synthesis requires a tokenizer model directory"); } @@ -163,7 +261,7 @@ Synthesizer::prepare(const SynthesisRequest& request) const { [&](const std::string& chunk, bool is_final) { const auto normalizer_start = std::chrono::steady_clock::now(); std::string normalized = - impl_->text_normalizer.normalize(chunk, prepared.metadata.language_code); + impl_->text_normalizer->normalize(chunk, prepared.metadata.language_code); prepared.normalizer_ms += std::chrono::duration( std::chrono::steady_clock::now() - normalizer_start) .count(); @@ -201,6 +299,17 @@ Synthesizer::prepare(const SynthesisRequest& request) const { SynthesisResult Synthesizer::synthesize(const PreparedSynthesis& request, const PcmCallback& callback) { + if (impl_->native) { + if (!request.tokens.empty() || !request.token_chunks.empty()) + throw std::invalid_argument("Nemotron does not support pretokenized synthesis"); + SynthesisRequest validation; + validation.text = request.metadata.original_text; + validation.language_code = request.metadata.language_code; + validation.output_sample_rate = request.metadata.sample_rate; + validation.options = request.options; + validation.apply_normalize = request.apply_normalize; + return impl_->native->synthesize(prepare(validation), callback); + } SynthesisResult result; result.metadata = request.metadata; const int output_rate = request.metadata.sample_rate; @@ -229,7 +338,7 @@ Synthesizer::synthesize(const PreparedSynthesis& request, const PcmCallback& cal return emit(std::string(reinterpret_cast(converted.data()), converted.size())); }; - result.stats = impl_->runtime.synthesize(request.token_chunks, request.options, process); + result.stats = impl_->runtime->synthesize(request.token_chunks, request.options, process); if (!result.cancelled && resampler) { std::vector tail; resampler->finish(&tail); @@ -259,6 +368,8 @@ SynthesisResult Synthesizer::synthesize_tokens( const std::vector& tokens, const MagpieSynthesisOptions& options, int output_sample_rate, const PcmCallback& callback) { + if (impl_->native) + throw std::invalid_argument("Nemotron does not support pretokenized synthesis"); if (tokens.empty()) throw std::invalid_argument("text token list is empty"); const int output_rate = output_sample_rate == 0 ? sample_rate() : output_sample_rate; @@ -288,38 +399,43 @@ SynthesisResult Synthesizer::warmup(const std::string& text, int steps) { if (text.empty()) return {}; + if (impl_->native && impl_->native->batch_size() > 1) + return {}; SynthesisRequest request; request.text = text; request.language_code = impl_->default_language_code; request.options.speaker = impl_->default_speaker; - request.options.steps = steps; + request.options.steps = impl_->native ? -1 : steps; return synthesize(request); } int Synthesizer::sample_rate() const { - return impl_->runtime.sample_rate(); + return impl_->native ? 22050 : impl_->runtime->sample_rate(); } int Synthesizer::speaker_count() const { - return impl_->runtime.speaker_count(); + return impl_->native ? static_cast(impl_->native->speaker_names().size()) + : impl_->runtime->speaker_count(); } const std::vector& Synthesizer::speaker_names() const { - return impl_->runtime.speaker_names(); + return impl_->native ? impl_->native->speaker_names() : impl_->runtime->speaker_names(); } std::vector Synthesizer::supported_language_codes() const { + if (impl_->native) + return {"en-US"}; return impl_->tokenizer ? impl_->tokenizer->supported_language_codes() : std::vector{}; } const std::string& Synthesizer::model_name() const { - return impl_->runtime.model_name(); + return impl_->native ? impl_->native->model_name() : impl_->runtime->model_name(); } const std::string& @@ -334,7 +450,24 @@ Synthesizer::default_speaker() const { bool Synthesizer::text_normalization_enabled() const { - return impl_->text_normalizer.enabled(); + return impl_->native ? impl_->native->text_normalization_enabled() + : impl_->text_normalizer->enabled(); +} + +const std::string& +Synthesizer::backend_name() const { + return impl_->backend; +} + +int +Synthesizer::backend_batch_size() const { + return impl_->native ? impl_->native->batch_size() : 1; +} + +const std::string& +Synthesizer::backend_mode() const { + static const std::string serial = "serial"; + return impl_->native ? impl_->native->mode() : serial; } } // namespace nemo_speech::tts diff --git a/src/tts/synthesizer.h b/src/tts/synthesizer.h index 7939512..9bf70b8 100644 --- a/src/tts/synthesizer.h +++ b/src/tts/synthesizer.h @@ -16,6 +16,9 @@ namespace nemo_speech::tts { struct SynthesizerConfig { + std::string backend = "magpie"; + std::string native_socket; + int native_timeout_seconds = 300; MagpieRuntimeConfig runtime; std::string tokenizer_model_dir; std::string text_normalizer_model_dir; @@ -29,6 +32,7 @@ struct SynthesisRequest { std::string language_code; std::string voice_name; int output_sample_rate = 0; // 0 = model rate + int apply_normalize = -1; MagpieSynthesisOptions options; }; @@ -48,6 +52,7 @@ struct PreparedSynthesis { std::vector tokens; std::vector> token_chunks; MagpieSynthesisOptions options; + int apply_normalize = -1; double normalizer_ms = 0.0; double tokenizer_ms = 0.0; }; @@ -87,6 +92,9 @@ class Synthesizer { const std::string& default_language_code() const; int default_speaker() const; bool text_normalization_enabled() const; + const std::string& backend_name() const; + int backend_batch_size() const; + const std::string& backend_mode() const; private: int resolve_speaker(const std::string& voice_name) const; diff --git a/tests/cpp/CMakeLists.txt b/tests/cpp/CMakeLists.txt index a84542a..14e933b 100644 --- a/tests/cpp/CMakeLists.txt +++ b/tests/cpp/CMakeLists.txt @@ -28,6 +28,18 @@ if(TARGET nemo_speech_http_server) endif() find_package(Python3 COMPONENTS Interpreter QUIET) +if(Python3_Interpreter_FOUND AND TARGET nemo_speech_cli AND NEMO_SPEECH_BUILD_TTS AND NEMO_SPEECH_BUILD_HTTP AND UNIX) + add_test(NAME nemotron_http COMMAND ${Python3_EXECUTABLE} + ${CMAKE_SOURCE_DIR}/tests/nemotron_http_test.py --binary $) + set_tests_properties(nemotron_http PROPERTIES TIMEOUT 60) + add_test(NAME nemotron_normalization COMMAND ${Python3_EXECUTABLE} + ${CMAKE_SOURCE_DIR}/tests/nemotron_normalization_test.py) + if(Python3_VERSION VERSION_GREATER_EQUAL 3.11) + add_test(NAME nemotron_bridge COMMAND ${Python3_EXECUTABLE} -m unittest discover + -s ${CMAKE_SOURCE_DIR}/scripts/nemotron_native/tests -v) + set_tests_properties(nemotron_bridge PROPERTIES TIMEOUT 60) + endif() +endif() if(Python3_Interpreter_FOUND AND (NEMO_SPEECH_BUILD_ASR OR NEMO_SPEECH_BUILD_DIAR OR NEMO_SPEECH_BUILD_NMT OR NEMO_SPEECH_BUILD_TTS)) @@ -123,6 +135,19 @@ if(NEMO_SPEECH_BUILD_ASR OR NEMO_SPEECH_BUILD_DIAR) add_subdirectory(asr) endif() if(NEMO_SPEECH_BUILD_TTS) + add_executable(test_nemotron_pcm common/test_nemotron_pcm.cpp) + target_include_directories(test_nemotron_pcm PRIVATE ${CMAKE_SOURCE_DIR}/src) + add_test(NAME nemotron_pcm COMMAND test_nemotron_pcm) + if(CMAKE_SYSTEM_NAME STREQUAL "Linux") + add_executable(test_nemotron_worker_protocol + common/test_nemotron_worker_protocol.cpp + ${CMAKE_SOURCE_DIR}/src/tts/nemotron/native_worker_protocol.cpp) + target_include_directories(test_nemotron_worker_protocol PRIVATE + ${CMAKE_SOURCE_DIR}/src + ${CMAKE_SOURCE_DIR}/llama.cpp/vendor) + add_test(NAME nemotron_worker_protocol COMMAND test_nemotron_worker_protocol) + set_tests_properties(nemotron_worker_protocol PROPERTIES TIMEOUT 30 SKIP_RETURN_CODE 77) + endif() add_subdirectory(tts) endif() if(NEMO_SPEECH_BUILD_NMT) diff --git a/tests/cpp/common/test_nemotron_pcm.cpp b/tests/cpp/common/test_nemotron_pcm.cpp new file mode 100644 index 0000000..861972c --- /dev/null +++ b/tests/cpp/common/test_nemotron_pcm.cpp @@ -0,0 +1,210 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tts/nemotron/pcm_conversion.h" + +using nemo_speech::tts::nemotron::pcm_f32_to_s16; +using nemo_speech::tts::nemotron::pcm_f32_to_s16_scalar; + +namespace { + +float +from_bits(uint32_t bits) { + float value; + std::memcpy(&value, &bits, sizeof(value)); + return value; +} + +std::string +encode(const std::vector& values) { + std::string encoded; + encoded.reserve(values.size() * 4); + for (float value : values) { + uint32_t bits; + std::memcpy(&bits, &value, sizeof(bits)); + for (unsigned shift = 0; shift < 32; shift += 8) + encoded.push_back(static_cast((bits >> shift) & 255)); + } + return encoded; +} + +std::string +oracle(const std::vector& values) { + std::string output; + output.reserve(values.size() * 2); + for (float sample : values) { + if (!std::isfinite(sample)) + throw std::runtime_error("nonfinite oracle input"); + const double scale = sample < 0 ? 32768.0 : 32767.0; + const double clipped = std::max(-1.0, std::min(1.0, static_cast(sample))); + const volatile float product = static_cast(clipped * scale); + const double rounded = std::floor(std::abs(static_cast(product)) + 0.5); + const int signed_value = static_cast(sample < 0 ? -rounded : rounded); + const auto bits = static_cast(signed_value); + output.push_back(static_cast(bits & 255)); + output.push_back(static_cast(bits >> 8)); + } + return output; +} + +void +require_equal(const std::string& actual, const std::string& expected, const std::string& name) { + if (actual != expected) { + const auto difference = + std::mismatch(actual.begin(), actual.end(), expected.begin(), expected.end()); + throw std::runtime_error( + name + " mismatch at sample " + + std::to_string((difference.first - actual.begin()) / 2)); + } +} + +bool +has_avx2() { +#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__)) + return __builtin_cpu_supports("avx2"); +#else + return false; +#endif +} + +void +compare(const std::vector& values, size_t alignment, const std::string& name) { + const auto expected = oracle(values); + const auto payload = encode(values); + const std::string storage(alignment, 'x'); + const auto padded = storage + payload + std::string(32, 'x'); + const char* input = padded.data() + alignment; + require_equal(pcm_f32_to_s16_scalar(input, payload.size()), expected, name + " scalar"); + require_equal(pcm_f32_to_s16(input, payload.size()), expected, name + " dispatch"); +#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__)) + if (has_avx2()) { + std::string output(expected.size() + 32, 'x'); + nemo_speech::tts::nemotron::pcm_f32_avx2(input, output.data() + 1, values.size()); + require_equal(output.substr(1, expected.size()), expected, name + " AVX2"); + if (output.front() != 'x' || output.substr(expected.size() + 1) != std::string(31, 'x')) + throw std::runtime_error("AVX2 output crossed its buffer extent"); + } +#endif +} + +template +void +require_rejected(Function function, const std::string& name) { + try { + function(); + } + catch (const std::runtime_error&) { + return; + } + throw std::runtime_error(name + " accepted invalid input"); +} + +void +test_nonfinite() { + for (uint32_t bits : + {0x7f800000u, 0xff800000u, 0x7fc00000u, 0xffc00000u, 0x7f800001u, 0xff800001u}) { + for (size_t position = 0; position < 25; ++position) { + std::vector values(25, 0.25f); + values[position] = from_bits(bits); + const auto payload = encode(values); + require_rejected( + [&] { pcm_f32_to_s16_scalar(payload.data(), payload.size()); }, "scalar nonfinite"); + require_rejected( + [&] { pcm_f32_to_s16(payload.data(), payload.size()); }, "dispatch nonfinite"); +#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__)) + if (has_avx2()) { + std::string output(values.size() * 2, '\0'); + require_rejected( + [&] { + nemo_speech::tts::nemotron::pcm_f32_avx2( + payload.data(), output.data(), values.size()); + }, + "AVX2 nonfinite"); + } +#endif + } + } +} + +} // namespace + +int +main() { + try { + const float infinity = std::numeric_limits::infinity(); + std::vector boundaries = { + 0.0f, + -0.0f, + 1.0f, + -1.0f, + 2.0f, + -2.0f, + std::numeric_limits::max(), + -std::numeric_limits::max(), + std::numeric_limits::min(), + -std::numeric_limits::min(), + std::numeric_limits::denorm_min(), + -std::numeric_limits::denorm_min(), + std::nextafter(1.0f, 0.0f), + std::nextafter(-1.0f, 0.0f), + std::nextafter(1.0f, infinity), + std::nextafter(-1.0f, -infinity)}; + for (int sign : {-1, 1}) { + const int scale = sign < 0 ? 32768 : 32767; + for (int integer = 0; integer < scale; ++integer) { + const float threshold = static_cast(sign * (integer + 0.5) / scale); + boundaries.push_back(std::nextafter(threshold, -infinity)); + boundaries.push_back(threshold); + boundaries.push_back(std::nextafter(threshold, infinity)); + } + } + compare(boundaries, 1, "half-threshold neighbors"); + std::vector random; + random.reserve(100000); + uint32_t state = 0x32a174f9; + while (random.size() < 100000) { + state ^= state << 13; + state ^= state >> 17; + state ^= state << 5; + if ((state & 0x7f800000u) != 0x7f800000u) + random.push_back(from_bits(state)); + } + compare(random, 3, "deterministic random float bits"); + for (size_t alignment = 0; alignment < 32; ++alignment) { + for (size_t count = 0; count <= 17; ++count) { + compare( + std::vector(boundaries.begin(), boundaries.begin() + count), alignment, + "alignment " + std::to_string(alignment) + " tail " + std::to_string(count)); + } + } + const std::string malformed(64, '\0'); + for (size_t size = 1; size < malformed.size(); ++size) { + if (size % 4) { + require_rejected( + [&] { pcm_f32_to_s16_scalar(malformed.data(), size); }, + "scalar malformed extent"); + require_rejected( + [&] { pcm_f32_to_s16(malformed.data(), size); }, "dispatch malformed extent"); + } + } + test_nonfinite(); + std::cout + << "PCM conversion passed " << boundaries.size() + << " boundary samples, 100000 random samples, all tails and invalid samples; AVX2=" + << has_avx2() << '\n'; + return 0; + } + catch (const std::exception& error) { + std::cerr << error.what() << '\n'; + return 1; + } +} diff --git a/tests/cpp/common/test_nemotron_worker_protocol.cpp b/tests/cpp/common/test_nemotron_worker_protocol.cpp new file mode 100644 index 0000000..4e225d3 --- /dev/null +++ b/tests/cpp/common/test_nemotron_worker_protocol.cpp @@ -0,0 +1,300 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tts/nemotron/native_worker_protocol.h" + +namespace protocol = nemo_speech::tts::nemotron::worker_protocol; +using protocol::Message; + +namespace { +constexpr uint64_t kSlotBytes = 8 * 1764 * 4; +size_t assertions = 0; + +void +require(bool condition, const std::string& name) { + ++assertions; + if (!condition) + throw std::runtime_error(name); +} + +template +void +rejects(Function function, const std::string& name) { + ++assertions; + try { + function(); + } + catch (const std::exception&) { + return; + } + throw std::runtime_error("accepted invalid " + name); +} + +std::string +bytes(std::initializer_list values) { + std::string result; + for (const auto value : values) result.push_back(static_cast(value)); + return result; +} + +std::string +unique_name() { + static uint64_t counter = 0; + std::ostringstream value; + value << "r20-pcm-" << std::hex << std::setfill('0') << std::setw(16) + << static_cast(::getpid()) << std::setw(16) + << (static_cast(std::chrono::steady_clock::now().time_since_epoch().count()) + + ++counter); + return value.str(); +} + +struct Arena { + std::string name = unique_name(); + std::string path = "/dev/shm/" + name; + int fd = -1; + + Arena() { + fd = ::open(path.c_str(), O_RDWR | O_CREAT | O_EXCL | O_CLOEXEC, 0600); + if (fd < 0) + throw std::runtime_error("cannot create test PCM arena"); + if (::ftruncate(fd, 2 * kSlotBytes) != 0) { + ::close(fd); + ::unlink(path.c_str()); + throw std::runtime_error("cannot allocate test PCM arena"); + } + } + ~Arena() { + if (fd >= 0) + ::close(fd); + ::unlink(path.c_str()); + } + void write(uint64_t slot, const std::vector& samples) { + const auto count = samples.size() * sizeof(float); + require( + ::pwrite(fd, samples.data(), count, slot * kSlotBytes) == static_cast(count), + "test arena write failed"); + } + Message metadata(uint64_t slot = 0, uint64_t generation = 1) const { + return { + {"arena", name}, + {"arena_bytes", 2 * kSlotBytes}, + {"slot_bytes", kSlotBytes}, + {"slot", slot}, + {"generation", generation}, + {"offset", slot * kSlotBytes}, + {"data_bytes", 8}, + {"samples", 2}, + {"sequence", 0}, + {"slab_version", 1}, + {"sample_rate_hz", 22050}, + {"channels", 1}, + {"encoding", "pcm_f32_le"}, + {"is_final", false}}; + } +}; + +void +test_parser_and_types() { + for (const uint64_t value : + {uint64_t{0}, uint64_t{1}, (uint64_t{1} << 53) + 1, uint64_t{1} << 63, UINT64_MAX}) { + const auto decoded = protocol::unpack(protocol::pack(Message{{"value", value}})); + require(protocol::integer(decoded, "value") == value, "uint64 lost precision"); + } + for (const Message value : + {Message(true), Message(-1), Message(1.0), Message("1"), Message(nullptr)}) + rejects([&] { protocol::integer(Message{{"value", value}}, "value"); }, "integer type"); + rejects([] { protocol::integer(Message{{"value", 3}}, "value", 2); }, "integer maximum"); + rejects([] { protocol::string(Message{{"value", false}}, "value"); }, "string type"); + rejects([] { protocol::string(Message{{"value", "abc"}}, "value", 2); }, "string length"); + const std::string valid = "Gr\xc3\xbc\xc3\x9f"; + require( + protocol::string(protocol::unpack(protocol::pack(Message{{"value", valid}})), "value") == + valid, + "UTF-8 round trip differs"); + const std::string embedded("a\0b", 3); + require( + protocol::string(protocol::unpack(protocol::pack(Message{{"value", embedded}})), "value") == + embedded, + "embedded NUL differs"); + for (const auto& invalid : + {bytes({0xc0, 0x80}), bytes({0xed, 0xa0, 0x80}), bytes({0xf4, 0x90, 0x80, 0x80}), + bytes({0xe2, 0x82})}) + rejects( + [&] { protocol::unpack(protocol::pack(Message{{"value", invalid}})); }, + "UTF-8 encoding"); + for (const double value : + {std::numeric_limits::infinity(), -std::numeric_limits::infinity(), + std::numeric_limits::quiet_NaN()}) + rejects( + [&] { protocol::unpack(protocol::pack(Message{{"value", value}})); }, + "nonfinite number"); + rejects([] { protocol::unpack(""); }, "empty payload"); + rejects([] { protocol::unpack(std::string(16 * 1024 * 1024 + 1, '\0')); }, "oversized payload"); + rejects([] { protocol::unpack(bytes({0x82, 0xa1, 'x', 1, 0xa1, 'x', 2})); }, "duplicate key"); + rejects([] { protocol::unpack(bytes({0x80, 0x80})); }, "trailing root"); + rejects([] { protocol::unpack(bytes({0x90})); }, "non-map root"); + rejects([] { protocol::unpack(bytes({0x81, 0xa1, 'x', 0xc7, 1, 1, 0})); }, "extension type"); + rejects( + [] { protocol::unpack(bytes({0x81, 0xa1, 'x', 0xdd, 0, 0, 0x40, 1})); }, + "array size bound"); + Message deep = Message::object(); + for (size_t depth = 0; depth < 66; ++depth) deep = Message{{"nested", std::move(deep)}}; + rejects([&] { protocol::unpack(protocol::pack(deep)); }, "nesting bound"); + const auto packed = + protocol::pack(Message{{"kind", "submit"}, {"token_ids", Message::array({1, 131073})}}); + for (size_t length = 0; length < packed.size(); ++length) + rejects( + [&] { protocol::unpack(std::string_view(packed).substr(0, length)); }, + "truncated message"); +} + +void +test_config() { + Arena arena; + const auto write = [&](const std::string& data) { + std::ofstream output(arena.path, std::ios::binary | std::ios::trunc); + output << data; + require(static_cast(output), "test config write failed"); + }; + write("{\"value\":18446744073709551615}"); + require( + protocol::integer(protocol::load_config(arena.path), "value") == UINT64_MAX, + "JSON config uint64 lost precision"); + for (const std::string value : {"{} trailing", "[]", "{\"x\":1,\"x\":2}", ""}) { + write(value); + rejects([&] { protocol::load_config(arena.path); }, "config shape"); + } + write(std::string(1048577, ' ')); + rejects([&] { protocol::load_config(arena.path); }, "config size"); + rejects([] { protocol::load_config("relative.json"); }, "relative config path"); + rejects( + [&] { protocol::load_config(arena.path + std::string("\0extra", 6)); }, "NUL config path"); +} + +void +test_slab_ownership_and_bounds() { + Arena arena; + arena.write(0, {0.25f, -0.5f}); + arena.write(1, {0.5f, -0.25f}); + protocol::SlabReader reader; + const auto initial = reader.read(arena.metadata()); + require( + initial.bytes == 8 && initial.samples == 2 && initial.sequence == 0, "slab shape differs"); + const std::string owned(initial.data, initial.bytes); + arena.write(0, {1.0f, 1.0f}); + const std::vector expected{0.25f, -0.5f}; + require( + std::memcmp(owned.data(), expected.data(), owned.size()) == 0, + "owned PCM changed after arena reuse"); + rejects([&] { reader.read(arena.metadata()); }, "repeated generation"); + reader.read(arena.metadata(0, 2)); + rejects([&] { reader.read(arena.metadata(0, 1)); }, "stale generation"); + const auto other = reader.read(arena.metadata(1)); + require( + other.data - initial.data == static_cast(kSlotBytes), + "upper slot offset differs"); + reader.read(arena.metadata(1, (uint64_t{1} << 53) + 1)); + reader.read(arena.metadata(1, (uint64_t{1} << 53) + 2)); + reader.read(arena.metadata(1, UINT64_MAX)); + rejects([&] { reader.read(arena.metadata(1, UINT64_MAX)); }, "uint64 generation replay"); + const std::vector> invalid = { + {"arena", "../escape"}, + {"arena", "r20-pcm-" + std::string(32, 'G')}, + {"slab_version", 2}, + {"arena_bytes", 0}, + {"arena_bytes", kSlotBytes + 1}, + {"arena_bytes", 4097 * kSlotBytes}, + {"slot_bytes", kSlotBytes - 4}, + {"slot", true}, + {"slot", -1}, + {"slot", 2}, + {"slot", 4096}, + {"generation", 0}, + {"generation", 1.0}, + {"offset", 4}, + {"offset", UINT64_MAX}, + {"data_bytes", 0}, + {"data_bytes", 7}, + {"data_bytes", kSlotBytes + 4}, + {"samples", false}, + {"samples", -1}, + {"samples", 3}, + {"sequence", uint64_t{UINT32_MAX} + 1}, + {"sample_rate_hz", 16000}, + {"channels", 2}, + {"encoding", "pcm_s16_le"}, + {"is_final", true}, + {"is_final", 0}}; + for (const auto& entry : invalid) { + protocol::SlabReader fresh; + auto value = arena.metadata(); + value[entry.first] = entry.second; + rejects([&] { fresh.read(value); }, "slab " + entry.first); + fresh.read(arena.metadata()); + } + auto changed = arena.metadata(0, 3); + changed["arena"] = unique_name(); + rejects([&] { reader.read(changed); }, "arena identity change"); + changed = arena.metadata(0, 3); + changed["arena_bytes"] = kSlotBytes; + rejects([&] { reader.read(changed); }, "arena extent change"); + reader.read(arena.metadata(0, 3)); +} + +void +test_slab_file_contract() { + Arena target; + const auto name = unique_name(); + const auto path = "/dev/shm/" + name; + require(::symlink(target.path.c_str(), path.c_str()) == 0, "test symlink creation failed"); + auto value = target.metadata(); + value["arena"] = name; + protocol::SlabReader reader; + rejects([&] { reader.read(value); }, "symlink arena"); + require(::unlink(path.c_str()) == 0, "test symlink removal failed"); + require(::mkdir(path.c_str(), 0700) == 0, "test directory creation failed"); + rejects([&] { reader.read(value); }, "directory arena"); + require(::rmdir(path.c_str()) == 0, "test directory removal failed"); + require(::ftruncate(target.fd, kSlotBytes) == 0, "test truncate failed"); + rejects([&] { reader.read(target.metadata()); }, "file extent mismatch"); + require(::ftruncate(target.fd, 2 * kSlotBytes) == 0, "test restore failed"); + reader.read(target.metadata()); +} +} + +int +main() { + if (::access("/dev/shm", W_OK | X_OK) != 0) { + std::cout << "SKIP: /dev/shm unavailable for native slab ownership tests\n"; + return 77; + } + try { + test_parser_and_types(); + test_config(); + test_slab_ownership_and_bounds(); + test_slab_file_contract(); + std::cout << "Native worker protocol passed " << assertions << " checks\n"; + return 0; + } + catch (const std::exception& error) { + std::cerr << error.what() << '\n'; + return 1; + } +} diff --git a/tests/nemotron_http_test.py b/tests/nemotron_http_test.py new file mode 100644 index 0000000..b105ed1 --- /dev/null +++ b/tests/nemotron_http_test.py @@ -0,0 +1,340 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import http.client +import io +import json +import os +import signal +import socket +import struct +import subprocess +import sys +import tempfile +import time +import wave +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from nemotron_mock_worker import MockWorker + + +def require(condition, message): + if not condition: + raise AssertionError(message) + + +def free_port(): + with socket.socket() as listener: + listener.bind(("127.0.0.1", 0)) + return listener.getsockname()[1] + + +def open_speech(port, text, response_format="pcm", seed=None): + connection = http.client.HTTPConnection("127.0.0.1", port, timeout=8) + payload = { + "model": "nemotron-tts", + "input": text, + "voice": "english_sample", + "language": "en-US", + "response_format": response_format, + } + if seed is not None: + payload["seed"] = seed + connection.request( + "POST", + "/v1/audio/speech", + json.dumps(payload), + {"Content-Type": "application/json"}, + ) + return connection, connection.getresponse() + + +def speech(port, text, response_format="pcm", seed=None): + connection, response = open_speech(port, text, response_format, seed=seed) + try: + return response.status, response.read(), dict(response.getheaders()) + finally: + response.close() + connection.close() + + +def failed_stream(port, text, require_no_audio=False): + connection, response = open_speech(port, text) + try: + try: + payload = response.read() + except http.client.IncompleteRead as error: + if require_no_audio: + require(not error.partial, "invalid F32 frame published partial PCM audio") + return + require( + response.status >= 400 and payload, + "worker failure became a successfully completed audio response", + ) + finally: + response.close() + connection.close() + + +def run(binary, float_audio=False, runtime_seed=None, preparation=None): + with tempfile.TemporaryDirectory(prefix="nemo-native-http-") as temporary: + folder = Path(temporary) + worker = MockWorker( + folder / "worker.sock", float_audio=float_audio, preparation=preparation + ) + port = free_port() + with (folder / "server.log").open("w+") as log: + process = subprocess.Popen( + [ + str(binary), + "serve", + "--host", + "127.0.0.1", + "--port", + str(port), + "--tts-backend", + "nemotron", + "--native-socket", + worker.path, + "--http-threads", + "8", + "--no-warmup", + ] + + ([] if runtime_seed is None else ["--tts.seed", str(runtime_seed)]), + stdout=log, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + try: + deadline = time.monotonic() + 15 + while time.monotonic() < deadline: + if process.poll() is not None: + raise RuntimeError("server exited during startup") + connection = http.client.HTTPConnection("127.0.0.1", port, timeout=1) + try: + connection.request("GET", "/ready") + response = connection.getresponse() + response.read() + if response.status == 200: + break + except OSError: + time.sleep(0.05) + finally: + connection.close() + else: + raise RuntimeError("server did not become ready") + require(worker.ready.is_set(), "server skipped bridge handshake") + + connection, response = open_speech(port, "stream") + try: + require(response.status == 200, "streaming request failed") + require( + response.getheader("Content-Type", "").startswith("audio/pcm"), + "PCM content type", + ) + require( + response.getheader("x-audio-sample-rate") == "22050", + "missing PCM sample rate", + ) + first = response.read(512) + require(len(first) == 512, "missing first PCM chunk") + require( + not worker.completed, + "worker completed before first audio arrived", + ) + worker.release_stream.set() + rest = response.read() + require( + len(first + rest) == 1024, + "stream PCM was truncated or duplicated", + ) + if float_audio: + expected = struct.pack("<512h", *([-16384] * 256 + [16384] * 256)) + require( + first + rest == expected, + "F32 conversion changed PCM16 bytes", + ) + finally: + worker.release_stream.set() + response.close() + connection.close() + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [ + executor.submit(speech, port, text) for text in ("parallel-A", "parallel-B") + ] + results = [future.result(timeout=10) for future in futures] + require( + worker.parallel_ready.is_set(), + "requests did not overlap inside the worker", + ) + require( + all(status == 200 and len(pcm) == 1024 for status, pcm, _ in results), + "concurrent requests lost audio", + ) + + status, payload, _ = speech(port, "wave", "wav") + require(status == 200, "WAV request failed") + with wave.open(io.BytesIO(payload), "rb") as audio: + require( + ( + audio.getnchannels(), + audio.getsampwidth(), + audio.getframerate(), + audio.getnframes(), + ) + == (1, 2, 22050, 512), + "WAV metadata or sample extent mismatch", + ) + + status, payload, _ = speech(port, "failure", "wav") + require(status >= 400, "worker preparation failure became HTTP success") + require( + b"injected preparation failure" in payload, + "worker error message was lost", + ) + failed_stream(port, "failure") + require( + speech(port, "recovery")[0] == 200, + "request failure poisoned other requests", + ) + + connection, response = open_speech(port, "cancel") + require( + response.status == 200 and len(response.read(512)) == 512, + "cancellation request lacked audio", + ) + response.close() + connection.close() + require( + worker.aborted.wait(5), + "HTTP disconnect did not abort backend request", + ) + require( + speech(port, "after-cancel")[0] == 200, + "cancellation poisoned subsequent requests", + ) + require( + speech(port, "seed-override", seed=11)[0] == 200, + "explicit request seed failed", + ) + + if preparation: + require( + speech(port, " prepared-trim ")[0] == 200, + "prepared request failed", + ) + sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "python")) + from nemotron_cpp_bridge.native_preparation import NativePreparation + + reference = NativePreparation( + preparation["library"], Path(preparation["tokenizer"]) + ) + try: + for request in worker.submissions: + text, tokens = reference.prepare_tokens(request["text"]) + require( + request["kind"] == "submit_prepared", + "frontend skipped native preparation", + ) + require( + (request["text"], request.get("token_ids")) == (text, tokens), + "frontend changed prepared text/tokens", + ) + require( + worker.submissions[-1]["text"] == "prepared-trim", + "frontend did not normalize whitespace", + ) + finally: + reference.close() + + failed_stream(port, "malformed", require_no_audio=float_audio) + require(not worker.errors, "mock protocol errors: " + repr(worker.errors)) + require( + all( + request["voice"] == "english_sample" + and request["language"] == "en-US" + and request["seed"] + == ( + 11 + if request["text"] == "seed-override" + else (0 if runtime_seed is None else runtime_seed) + ) + for request in worker.submissions + ), + "request metadata changed", + ) + except BaseException: + log.flush() + log.seek(0) + print(log.read()) + raise + finally: + if process.poll() is None: + os.killpg(process.pid, signal.SIGTERM) + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGKILL) + process.wait(timeout=5) + worker.close() + print( + f"PASS ({'F32' if float_audio else 'PCM16'} frames): " + "streaming, concurrency, WAV, errors, cancellation, malformed PCM, runtime seed" + ) + + +def reject_runtime_options(binary): + cases = ( + ("--tts.steps", "64", "Nemotron requires steps=512"), + ("--tts.top-k", "10", "Nemotron requires top_k=80"), + ("--tts.temperature", "0.5", "Nemotron requires temperature=0.7"), + ("--tts.cfg-scale", "1.0", "Nemotron does not support cfg_scale"), + ) + with tempfile.TemporaryDirectory(prefix="nemo-native-runtime-") as temporary: + for option, value, error in cases: + result = subprocess.run( + [ + str(binary), + "serve", + "--tts-backend", + "nemotron", + "--native-socket", + str(Path(temporary) / "absent.sock"), + "--no-warmup", + option, + value, + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + timeout=5, + ) + require( + result.returncode != 0 and error in result.stdout, + f"{option}={value} did not reject the incompatible runtime setting: {result.stdout}", + ) + print("PASS: incompatible CLI steps, top-k, temperature, and CFG rejected") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--binary", required=True, type=Path) + parser.add_argument("--preparation-library", type=Path) + parser.add_argument("--tokenizer", type=Path) + arguments = parser.parse_args() + if bool(arguments.preparation_library) != bool(arguments.tokenizer): + parser.error("--preparation-library and --tokenizer are required together") + binary = arguments.binary.resolve(strict=True) + run(binary, runtime_seed=7) + run(binary, float_audio=True) + if arguments.preparation_library: + preparation = { + "library": str(arguments.preparation_library.resolve(strict=True)), + "tokenizer": str(arguments.tokenizer.resolve(strict=True)), + "normalization_config": "null", + } + run(binary, float_audio=True, runtime_seed=7, preparation=preparation) + reject_runtime_options(binary) diff --git a/tests/nemotron_mock_worker.py b/tests/nemotron_mock_worker.py new file mode 100644 index 0000000..4e12ba8 --- /dev/null +++ b/tests/nemotron_mock_worker.py @@ -0,0 +1,188 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import socket +import struct +import threading + + +class MockWorker: + def __init__(self, path, float_audio=False, preparation=None): + self.path = str(path) + self.float_audio = float_audio + self.preparation = preparation + self.ready = threading.Event() + self.release_stream = threading.Event() + self.parallel_ready = threading.Event() + self.aborted = threading.Event() + self.closed = threading.Event() + self.lock = threading.Lock() + self.send_lock = threading.Lock() + self.submissions = [] + self.completed = [] + self.requests = {} + self.errors = [] + self.threads = [] + self.listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + self.listener.bind(self.path) + self.listener.listen(1) + self.listener.settimeout(0.2) + self.connection = None + self.thread = threading.Thread(target=self._run, daemon=True) + self.thread.start() + + def _send(self, payload): + with self.send_lock: + self.connection.sendall(struct.pack(">I", len(payload)) + payload) + + def _json(self, **value): + self._send(b"J" + json.dumps(value).encode()) + + def _read(self, size): + result = bytearray() + while len(result) < size: + data = self.connection.recv(size - len(result)) + if not data: + raise EOFError() + result.extend(data) + return bytes(result) + + def _run(self): + try: + while not self.closed.is_set(): + try: + self.connection, _ = self.listener.accept() + break + except socket.timeout: + continue + if self.connection is None: + return + while not self.closed.is_set(): + length = struct.unpack(">I", self._read(4))[0] + if not 1 <= length <= 1024 * 1024: + raise ValueError("invalid control frame length") + payload = self._read(length) + if payload[:1] != b"J": + raise ValueError("expected JSON control frame") + message = json.loads(payload[1:]) + kind = message["kind"] + if kind == "hello": + if message["protocol_version"] != 1: + raise ValueError("unexpected protocol version") + self._json( + kind="ready", + protocol_version=1, + model_name="nemotron-tts", + sample_rate=22050, + voices=["english_sample"], + batch_size=32, + mode="inflight", + text_normalization_enabled=False, + **({"text_preparation": self.preparation} if self.preparation else {}), + ) + self.ready.set() + elif kind in ("submit", "submit_prepared"): + if kind != ("submit_prepared" if self.preparation else "submit"): + raise ValueError("incorrect preparation dispatch") + request_id = message["request_id"] + event = threading.Event() + with self.lock: + if request_id in self.requests: + raise ValueError("duplicate request ID") + self.requests[request_id] = event + self.submissions.append(message) + if sum(x["text"].startswith("parallel") for x in self.submissions) >= 2: + self.parallel_ready.set() + thread = threading.Thread( + target=self._synthesize, args=(message, event), daemon=True + ) + self.threads.append(thread) + thread.start() + elif kind == "abort": + with self.lock: + event = self.requests.get(message["request_id"]) + if event is not None: + event.set() + self.aborted.set() + else: + raise ValueError(f"unexpected control kind {kind}") + except (EOFError, OSError): + pass + except Exception as error: + self.errors.append(str(error)) + + def _synthesize(self, message, abort): + request_id, text = message["request_id"], message["text"] + try: + if text == "failure": + self._json( + kind="failed", + request_id=request_id, + message="injected preparation failure", + ) + return + self._json( + kind="accepted", + request_id=request_id, + processed_text=text, + token_count=len(message.get("token_ids", [0, 0, 0])), + ) + audio_kind = b"F" if self.float_audio else b"A" + if text == "malformed": + invalid = ( + struct.pack("<9f", *([0.5] * 8 + [float("nan")])) if self.float_audio else b"x" + ) + self._send(audio_kind + struct.pack(">QI", request_id, 0) + invalid) + return + if text.startswith("parallel") and not self.parallel_ready.wait(5): + raise RuntimeError("server serialized concurrent requests") + count = 0 + for sequence in range(1000 if text == "cancel" else 2): + if abort.is_set() or self.closed.is_set(): + self._json(kind="failed", request_id=request_id, message="cancelled") + return + pcm = ( + struct.pack("<256f", *([0.5 if sequence % 2 else -0.5] * 256)) + if self.float_audio + else struct.pack("<256h", *([1000 if sequence % 2 else -1000] * 256)) + ) + self._send(audio_kind + struct.pack(">QI", request_id, sequence) + pcm) + count += 256 + if text == "stream" and sequence == 0 and not self.release_stream.wait(5): + raise RuntimeError("HTTP withheld first audio until completion") + if text == "cancel": + abort.wait(0.01) + self._json( + kind="complete", + request_id=request_id, + total_samples=count, + chunk_count=2, + ) + with self.lock: + self.completed.append(request_id) + except (OSError, EOFError): + pass + except Exception as error: + self.errors.append(str(error)) + try: + self._json(kind="failed", request_id=request_id, message=str(error)) + except OSError: + pass + + def close(self): + self.closed.set() + self.release_stream.set() + self.parallel_ready.set() + for event in self.requests.values(): + event.set() + if self.connection is not None: + try: + self.connection.shutdown(socket.SHUT_RDWR) + except OSError: + pass + self.connection.close() + self.listener.close() + self.thread.join(timeout=2) + for thread in self.threads: + thread.join(timeout=2) diff --git a/tests/nemotron_normalization_test.py b/tests/nemotron_normalization_test.py new file mode 100644 index 0000000..ef10e47 --- /dev/null +++ b/tests/nemotron_normalization_test.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import importlib.util +import os +import threading +import unittest +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +SPEC = importlib.util.spec_from_file_location( + "normalization", + Path(__file__).resolve().parents[1] / "python/nemotron_cpp_bridge/normalization.py", +) +MODULE = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(MODULE) + + +class PreparationTests(unittest.TestCase): + def test_bypass_trims_without_changing_punctuation(self): + preparer = MODULE.TextPreparer() + self.assertFalse(preparer.enabled) + self.assertEqual(preparer.prepare(" \tKeep 2 apples\n", "en-US"), "Keep 2 apples") + self.assertEqual(preparer.prepare("Numbers .", "EN-us", False), "Numbers .") + preparer.close() + preparer.close() + with self.assertRaisesRegex(RuntimeError, "closed"): + preparer.prepare("Hello", "en-US") + + def test_rejects_invalid_request(self): + preparer = MODULE.TextPreparer() + for text, language, normalize in [ + (" ", "en-US", True), + ("ok", "de-DE", True), + ("ok", "en-US", "true"), + (None, "en-US", True), + ("x" * 65537, "en-US", True), + ]: + with self.subTest(text=str(text)[:10], language=language, normalize=normalize): + with self.assertRaises(ValueError): + preparer.prepare(text, language, normalize) + preparer.close() + + def test_pool_does_not_share_handles_and_close_waits(self): + preparer = MODULE.TextPreparer() + preparer.enabled = True + preparer._handles = [1, 2] + entered = threading.Barrier(3) + release = threading.Event() + destroyed = [] + preparer._api = type("Api", (), {"destroy": lambda self, h: destroyed.append(h)})() + + def process(handle, text): + entered.wait(timeout=5) + if not release.wait(timeout=5): + raise RuntimeError("test did not release normalizer") + return handle, text + + preparer._process = process + with ThreadPoolExecutor(max_workers=3) as executor: + first = executor.submit(preparer.prepare, "one", "en-US") + second = executor.submit(preparer.prepare, "two", "en-US") + entered.wait(timeout=5) + closing = executor.submit(preparer.close) + self.assertFalse(closing.done()) + self.assertFalse(destroyed) + release.set() + self.assertNotEqual(first.result(timeout=5)[0], second.result(timeout=5)[0]) + closing.result(timeout=5) + self.assertEqual(sorted(destroyed), [1, 2]) + + +@unittest.skipUnless( + os.environ.get("NEMOTRON_NORMALIZATION_CONFIG"), "native normalizer config not provided" +) +class NativePreparationTests(unittest.TestCase): + def test_native_abi_and_parallel_replay(self): + preparer = MODULE.TextPreparer(os.environ["NEMOTRON_NORMALIZATION_CONFIG"]) + try: + self.assertTrue(preparer.enabled) + texts = [ + " Hello world ", + "I paid $12.50 for 3 books.", + "Call me at 5 p.m.", + "The year is 2026.", + ] + expected = [preparer.prepare(text, "en-US") for text in texts] + self.assertEqual(expected[0], "Hello world.") + self.assertNotIn("$", expected[1]) + self.assertNotIn("12.50", expected[1]) + self.assertEqual(preparer.prepare(" Keep 2 apples ", "en-US", False), "Keep 2 apples") + with ThreadPoolExecutor(max_workers=32) as executor: + actual = list( + executor.map( + lambda index: preparer.prepare(texts[index % 4], "en-US"), range(128) + ) + ) + self.assertEqual(actual, expected * 32) + finally: + preparer.close() + + +if __name__ == "__main__": + unittest.main()